"import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import cv2
import random
import tensorflow as tf
import tkinter as tk
from tkinter import filedialog
from PIL import ImageTk, Image
from ipywidgets import interact, interactive, fixed, interact_manual, HBox, VBox
import ipywidgets as widgets
from IPython.display import display, clear_output
from tensorflow.keras.preprocessing import image
from tensorflow.keras.optimizers import Adam, SGD, RMSprop, AdamW, Adadelta, Adagrad, Adamax, Adafactor, Nadam, Ftrl
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tqdm import tqdm
import os
from sklearn.utils import shuffle
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from tensorflow.keras.models import Sequential, Model, load_model
from tensorflow.keras.layers import (
    GlobalAveragePooling2D,
    Dropout,
    Dense,
    Conv2D,
    MaxPooling2D,
    Flatten,
    Dropout,
    BatchNormalization,
    Activation,
    concatenate,
    Conv2DTranspose,
    Input,
    Reshape,
    UpSampling2D,
)
from tensorflow.keras.applications import (
    EfficientNetV2B0,
    EfficientNetV2B1,
    EfficientNetV2B2,
    EfficientNetV2B3,
    EfficientNetV2L,
    EfficientNetV2M,
    EfficientNetV2S,
)
from tensorflow.keras.applications import Xception
from tensorflow.keras.applications import VGG16, VGG19
from tensorflow.keras.applications import ResNet50, ResNet101, ResNet152, ResNetRS50, ResNetRS101
from tensorflow.keras.applications import InceptionResNetV2, ConvNeXtXLarge, ConvNeXtBase, DenseNet121, MobileNetV2, NASNetLarge, NASNetMobile
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau, TensorBoard, ModelCheckpoint
from sklearn.metrics import classification_report, confusion_matrix
import ipywidgets as widgets
import io
from PIL import Image
from IPython.display import display, clear_output
from ipywidgets import Output, Button
from warnings import filterwarnings
from google.colab import drive
drive.mount("/content/gdrive")
def load_data(data_folders):
    X_data = []  # Combined data
    y_class_labels = []  # Combined classification labels
    y_seg_labels = []  # Combined segmentation labels
    for folderPath in data_folders:
        for label in labels:
            label_folder_path = os.path.join(folderPath, label)
            for filename in tqdm(os.listdir(label_folder_path)):
                if filename.endswith(".jpg"):
                    img = cv2.imread(os.path.join(label_folder_path, filename))
                    img = cv2.resize(img, (image_size, image_size))
                    X_data.append(img)
                    y_class_labels.append(label)
                    seg_filename = filename.split(".")[0] + ".png"
                    seg_img = cv2.imread(os.path.join(label_folder_path, seg_filename), 0)
                    seg_img = cv2.resize(seg_img, (image_size, image_size))
                    seg_img = np.where(seg_img > 0, 1, 0)  # Convert segmentation mask to binary
                    y_seg_labels.append(seg_img)
    X_data = np.array(X_data)
    y_class_labels = np.array(y_class_labels)
    y_seg_labels = np.array(y_seg_labels)
    X_data, y_class_labels, y_seg_labels = shuffle(X_data, y_class_labels, y_seg_labels, random_state=101)
    return X_data, y_class_labels, y_seg_labels

def split_data(X_data, y_class_labels, y_seg_labels, train_ratio, val_ratio, test_ratio):
    num_samples = len(X_data)
    num_train_samples = int(train_ratio * num_samples)
    num_val_samples = int(val_ratio * num_samples)
    num_test_samples = int(test_ratio * num_samples)
    X_train = X_data[:num_train_samples]
    y_train_class = y_class_labels[:num_train_samples]
    y_train_seg = y_seg_labels[:num_train_samples]
    X_val = X_data[num_train_samples: num_train_samples + num_val_samples]
    y_val_class = y_class_labels[num_train_samples: num_train_samples + num_val_samples]
    y_val_seg = y_seg_labels[num_train_samples: num_train_samples + num_val_samples]
    X_test = X_data[num_train_samples + num_val_samples: num_train_samples + num_val_samples + num_test_samples]
    y_test_class = y_class_labels[num_train_samples + num_val_samples: num_train_samples + num_val_samples + num_test_samples]
    y_test_seg = y_seg_labels[num_train_samples + num_val_samples: num_train_samples + num_val_samples + num_test_samples]
    # Convert class labels to categorical
    label_encoder = LabelEncoder()
    y_train_class_encoded = label_encoder.fit_transform(y_train_class)
    y_train_class_categorical = to_categorical(y_train_class_encoded)
    y_val_class_encoded = label_encoder.transform(y_val_class)
    y_val_class_categorical = to_categorical(y_val_class_encoded)
    y_test_class_encoded = label_encoder.transform(y_test_class)
    y_test_class_categorical = to_categorical(y_test_class_encoded)
    return (
        X_train,
        y_train_class_categorical,
        y_train_seg,
        X_val,
        y_val_class_categorical,
        y_val_seg,
        X_test,
        y_test_class_categorical,
        y_test_seg,
    )

def build_model(input_shape, num_classes):
    num_filter = 32  # 16/32 best, 8: best classification but no segment
    # Encoder (Done)
    inputs = Input(input_shape)
    conv1 = Conv2D(num_filter * 1, 3, activation="linear", padding="same", strides=1)(inputs)
    bn1 = BatchNormalization()(conv1)
    relu1 = Activation("relu")(bn1)
    conv2 = Conv2D(num_filter * 1, 3, activation="linear", padding="same", strides=1)(relu1)
    bn2 = BatchNormalization()(conv2)
    relu2 = Activation("relu")(bn2)
    down1 = MaxPooling2D(pool_size=(2, 2), strides=2)(relu2)
    conv3 = Conv2D(num_filter * 2, 3, activation="linear", padding="same", strides=1)(down1)
    bn3 = BatchNormalization()(conv3)
    relu3 = Activation("relu")(bn3)
    conv4 = Conv2D(num_filter * 2, 3, activation="linear", padding="same", strides=1)(relu3)
    bn4 = BatchNormalization()(conv4)
    relu4 = Activation("relu")(bn4)
    down2 = MaxPooling2D(pool_size=(2, 2), strides=2)(relu4)
    conv5 = Conv2D(num_filter * 4, 3, activation="linear", padding="same", strides=1)(down2)
    bn5 = BatchNormalization()(conv5)
    relu5 = Activation("relu")(bn5)
    conv6 = Conv2D(num_filter * 4, 3, activation="linear", padding="same", strides=1)(relu5)
    bn6 = BatchNormalization()(conv6)
    relu6 = Activation("relu")(bn6)
    down3 = MaxPooling2D(pool_size=(2, 2), strides=2)(relu6)
    conv7 = Conv2D(num_filter * 8, 3, activation="linear", padding="same", strides=1)(down3)
    bn7 = BatchNormalization()(conv7)
    relu7 = Activation("relu")(bn7)
    conv8 = Conv2D(num_filter * 8, 3, activation="linear", padding="same", strides=1)(relu7)
    bn8 = BatchNormalization()(conv8)
    relu8 = Activation("relu")(bn8)
    # Middle
    down4 = MaxPooling2D(pool_size=(2, 2), strides=2)(relu8)
    conv9 = Conv2D(num_filter * 16, 3, activation="linear", padding="same", strides=1)(down4)
    bn9 = BatchNormalization()(conv9)
    relu9 = Activation("relu")(bn9)
    conv10 = Conv2D(num_filter * 16, 3, activation="linear", padding="same", strides=1)(relu9)
    bn10 = BatchNormalization()(conv10)
    relu10 = Activation("relu")(bn10)
    up1 = UpSampling2D(size=(2, 2), interpolation="bilinear")(relu10)
    # Decoder (Done)
    concat1 = concatenate([up1, relu8], axis=-1)  # , axis=3
    conv11 = Conv2D(num_filter * 8, 3, activation="linear", padding="same", strides=1)(concat1)
    bn11 = BatchNormalization()(conv11)
    relu11 = Activation("relu")(bn11)
    conv12 = Conv2D(num_filter * 8, 3, activation="linear", padding="same", strides=1)(relu11)
    bn12 = BatchNormalization()(conv12)
    relu12 = Activation("relu")(bn12)
    up2 = UpSampling2D(size=(2, 2), interpolation="bilinear")(relu12)
    concat2 = concatenate([up2, relu6], axis=-1)  # , axis=3
    conv13 = Conv2D(num_filter * 4, 3, activation="linear", padding="same", strides=1)(concat2)
    bn13 = BatchNormalization()(conv13)
    relu13 = Activation("relu")(bn13)
    conv14 = Conv2D(num_filter * 4, 3, activation="linear", padding="same", strides=1)(relu13)
    bn14 = BatchNormalization()(conv14)
    relu14 = Activation("relu")(bn14)
    up3 = UpSampling2D(size=(2, 2), interpolation="bilinear")(relu14)
    concat3 = concatenate([up3, relu4], axis=-1)  # , axis=3
    conv15 = Conv2D(num_filter * 2, 3, activation="linear", padding="same", strides=1)(concat3)
    bn15 = BatchNormalization()(conv15)
    relu15 = Activation("relu")(bn15)
    conv16 = Conv2D(num_filter * 2, 3, activation="linear", padding="same", strides=1)(relu15)
    bn16 = BatchNormalization()(conv16)
    relu16 = Activation("relu")(bn16)
    up4 = UpSampling2D(size=(2, 2), interpolation="bilinear")(relu16)
    concat4 = concatenate([up4, relu2], axis=-1)  # , axis=3
    conv17 = Conv2D(num_filter * 1, 3, activation="linear", padding="same", strides=1)(concat4)
    bn17 = BatchNormalization()(conv17)
    relu17 = Activation("relu")(bn17)
    conv18 = Conv2D(num_filter * 1, 3, activation="linear", padding="same", strides=1)(relu17)
    bn18 = BatchNormalization()(conv18)
    relu18 = Activation("relu")(bn18)
    # Segmentation branch
    segmentation_output = Conv2D(1, 1, activation="sigmoid", name="segmentation_output")(relu18)  # original
    # Classification branch (Not done)
    gap1 = GlobalAveragePooling2D()(relu8)
    gap2 = GlobalAveragePooling2D()(relu10)
    gap3 = GlobalAveragePooling2D()(relu12)
    conv20 = Conv2D(16, 3, activation="linear", padding="same", strides=1)(segmentation_output)
    bn20 = BatchNormalization()(conv20)
    relu20 = Activation("relu")(bn20)
    down5 = MaxPooling2D(pool_size=(4, 4), strides=4)(relu20)
    conv21 = Conv2D(32, 3, activation="linear", padding="same", strides=1)(down5)
    bn21 = BatchNormalization()(conv21)
    relu21 = Activation("relu")(bn21)
    down6 = MaxPooling2D(pool_size=(4, 4), strides=4)(relu21)
    conv22 = Conv2D(64, 3, activation="linear", padding="same", strides=1)(down6)
    bn22 = BatchNormalization()(conv22)
    relu22 = Activation("relu")(bn22)
    down7 = MaxPooling2D(pool_size=(4, 4), strides=4)(relu22)
    flatten1 = Flatten()(down7)
    concat5 = concatenate([gap1, gap2, gap3, flatten1], axis=-1)
    # FC layers
    fc1 = Dense(1024, activation="relu")(concat5)
    dropout1 = Dropout(0.5)(fc1)
    fc2 = Dense(1024, activation="relu")(dropout1)
    dropout2 = Dropout(0.5)(fc2)
    classification_output = Dense(num_classes, activation="softmax", name="classification_output")(dropout2)
    # Define the model
    model = Model(inputs=inputs, outputs=[classification_output, segmentation_output])
    return model

def segmentation_loss(y_true, y_pred):
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred, tf.float32)
    bce_loss = tf.keras.losses.binary_crossentropy(y_true, y_pred)
    smooth = 1e-5
    intersection = tf.reduce_sum(y_true * y_pred)
    union = tf.reduce_sum(y_true) + tf.reduce_sum(y_pred)
    dice_loss = 1.0 - 2.0 * (intersection + smooth) / (union + smooth)
    segmentation_loss = bce_loss + 1 * dice_loss
    return segmentation_loss

def train_model(model, X_train, y_train_class, y_train_seg, X_val, y_val_class, y_val_seg, batch_size, epochs):
    checkpoint = ModelCheckpoint(
        "multitask_best_weights.h5",
        monitor="val_classification_output_accuracy",
        save_best_only=True,
        mode="max",
        verbose=1,)
    reduce_lr = ReduceLROnPlateau(
        monitor="val_classification_output_accuracy",
        factor=0.3,
        patience=2,
        min_delta=0.001,
        mode="auto",
        verbose=1,)
    tensorboard = TensorBoard(log_dir="logs")
    model.compile(
        optimizer=Adam(lr=0.001),
        loss={"classification_output": "categorical_crossentropy", "segmentation_output": segmentation_loss},
        metrics={"classification_output": "accuracy", "segmentation_output": "accuracy"},
        loss_weights={"classification_output": 1, "segmentation_output": 1},)
    history = model.fit(
        X_train,
        {"classification_output": y_train_class, "segmentation_output": y_train_seg},
        validation_data=(X_val, {"classification_output": y_val_class, "segmentation_output": y_val_seg}),
        epochs=epochs,
        verbose=1,
        batch_size=batch_size,
        callbacks=[checkpoint, reduce_lr, tensorboard],)
    return history

def evaluate_model(model, X_test, y_test_class, y_test_seg):
    with tf.keras.utils.custom_object_scope({"segmentation_loss": segmentation_loss}):
        # Load the best model weights
        best_model = load_model("multitask_best_weights.h5")
        # Evaluate the model on test data
        test_loss, test_class_loss, test_seg_loss, test_class_acc, test_seg_acc = best_model.evaluate(
            X_test, {"classification_output": y_test_class, "segmentation_output": y_test_seg})
        print("Test Classification Loss:", test_class_loss)
        print("Test Segmentation Loss:", test_seg_loss)
        print("Test Classification Accuracy:", test_class_acc)
        print("Test Segmentation Accuracy:", test_seg_acc)
        # Evaluate the model on validation data
        val_loss, val_class_loss, val_seg_loss, val_class_acc, val_seg_acc = best_model.evaluate(
            X_val, {'classification_output': y_val_class, 'segmentation_output': y_val_seg})
        print("Validation Classification Loss:", val_class_loss)
        print("Validation Segmentation Loss:", val_seg_loss)
        print("Validation Classification Accuracy:", val_class_acc)
        print("Validation Segmentation Accuracy:", val_seg_acc)
        # Evaluate the model on training data
        train_loss, train_class_loss, train_seg_loss, train_class_acc, train_seg_acc = best_model.evaluate(X_train, {'classification_output': y_train_class, 'segmentation_output': y_train_seg})
        print("Train Classification Loss:", train_class_loss)
        print("Train Segmentation Loss:", train_seg_loss)
        print("Train Classification Accuracy:", train_class_acc)
        print("Train Segmentation Accuracy:", train_seg_acc)
        # Return test classification accuracy
        return test_class_acc

def plot_performance(history):
    # Plot classification accuracy
    classification_train_accuracy = history.history["classification_output_accuracy"]
    classification_val_accuracy = history.history["val_classification_output_accuracy"]
    plt.figure(figsize=(7, 3))
    plt.plot(classification_train_accuracy, label="Training Accuracy")
    plt.plot(classification_val_accuracy, label="Validation Accuracy")
    plt.title("Classification Accuracy")
    plt.xlabel("Epochs")
    plt.ylabel("Accuracy")
    plt.legend()
    plt.show()
    # Plot classification loss
    classification_train_loss = history.history["classification_output_loss"]
    classification_val_loss = history.history["val_classification_output_loss"]
    plt.figure(figsize=(7, 3))
    plt.plot(classification_train_loss, "b", label="Training Loss")
    plt.plot(classification_val_loss, "r", label="Validation Loss")
    plt.title("Classification Loss")
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend()
    plt.show()
    # Plot segmentation accuracy
    segmentation_train_accuracy = history.history["segmentation_output_accuracy"]
    segmentation_val_accuracy = history.history["val_segmentation_output_accuracy"]
    plt.figure(figsize=(7, 3))
    plt.plot(segmentation_train_accuracy, label="Training Accuracy")
    plt.plot(segmentation_val_accuracy, label="Validation Accuracy")
    plt.title("Segmentation Accuracy")
    plt.xlabel("Epochs")
    plt.ylabel("Accuracy")
    plt.legend()
    plt.show()
    # Plot segmentation loss
    segmentation_train_loss = history.history["segmentation_output_loss"]
    segmentation_val_loss = history.history["val_segmentation_output_loss"]
    plt.figure(figsize=(7, 3))
    plt.plot(segmentation_train_loss, "b", label="Training Loss")
    plt.plot(segmentation_val_loss, "r", label="Validation Loss")
    plt.title("Segmentation Loss")
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend()
    plt.show()

def test_model(model, X_test, y_test_class, y_test_seg, labels, num_images=5):
    random_indices = np.random.choice(len(X_test), num_images, replace=False)
    selected_images = X_test[random_indices]
    selected_labels = y_test_class[random_indices]
    selected_masks = y_test_seg[random_indices]
    predictions = model.predict(selected_images)
    predicted_classes = np.argmax(predictions[0], axis=1)
    predicted_masks = predictions[1]

    bg_color = '#303030'  # Lighter than black, as per the previous specification
    text_color = 'white'

    for i in range(num_images):
        plt.figure(figsize=(10, 5))

        plt.subplot(2, 3, 1)
        plt.imshow(selected_images[i])
        plt.title(f"Input Image\nGround Truth Label: {labels[np.argmax(selected_labels[i])]}\nPredicted Label: {labels[predicted_classes[i]]}", color=text_color)
        plt.axis("off")
        plt.gcf().set_facecolor(bg_color)

        plt.subplot(2, 3, 2)
        plt.imshow(selected_masks[i].squeeze(), cmap="gray")
        plt.title("Ground Truth Segmentation Mask", color=text_color)
        plt.axis("off")

        plt.subplot(2, 3, 3)
        plt.imshow(predicted_masks[i].squeeze(), cmap="gray")
        plt.title("Predicted Segmentation Mask", color=text_color)
        plt.axis("off")
        plt.tight_layout()
        plt.show()

# preprocess and predict on the selected image
out = Output()
'''
def predict_image(image_path):
    # Load and preprocess the image
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
    img_resized = cv2.resize(img, (image_size, image_size))
    img_expanded = np.expand_dims(img_resized, axis=0)
    # Predict using the model
    class_output, seg_output = model.predict(img_expanded)
    # Get the predicted class
    predicted_class_index = np.argmax(class_output)
    predicted_class = labels[predicted_class_index]
    # Clear the previous output
    with out:
        clear_output(wait=True)
        # Visualize the input and segmentation output
        fig, axs = plt.subplots(1, 3, figsize=(15, 5))
        axs[0].imshow(img)
        axs[0].set_title("Input Image")
        axs[0].axis("off")

        axs[1].imshow(seg_output[0, :, :, 0], cmap='gray')
        axs[1].set_title("Segmentation Output")
        axs[1].axis("off")

        class_probabilities = class_output[0]
        x_pos = np.arange(len(labels))
        axs[2].bar(x_pos, class_probabilities, align='center', alpha=0.7, color='blue')
        axs[2].set_xticks(x_pos)
        axs[2].set_xticklabels(labels, rotation=45, ha="right")
        axs[2].set_ylabel('Probability')
        axs[2].set_ylim([0, 1])
        axs[2].set_title(f'Classification Output\nPredicted Class: {predicted_class}')  # Increased title font size

        for i, v in enumerate(class_probabilities):
            axs[2].text(i, v + 0.02, f"{v:.2f}", color='black', ha='center')

        plt.tight_layout()
        plt.show()
'''

def predict_image(image_path):
    # Load and preprocess the image
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
    img_resized = cv2.resize(img, (image_size, image_size))
    img_expanded = np.expand_dims(img_resized, axis=0)
    # Predict using the model
    class_output, seg_output = model.predict(img_expanded)
    # Get the predicted class
    predicted_class_index = np.argmax(class_output)
    predicted_class = labels[predicted_class_index]
    # Clear the previous output
    with out:
        clear_output(wait=True)
        # Configure specific style options for lighter background
        fig, axs = plt.subplots(1, 3, figsize=(15, 5))

        # Lighten background
        bg_color = '#303030'  # Lighter than black for contrast
        fig.patch.set_facecolor(bg_color)
        for ax in axs:
            ax.set_facecolor(bg_color)
            # Set spines color to lighter tone
            ax.spines['bottom'].set_color('grey')
            ax.spines['top'].set_color('grey')
            ax.spines['right'].set_color('grey')
            ax.spines['left'].set_color('grey')
            ax.tick_params(axis='x', colors='white')
            ax.tick_params(axis='y', colors='white')

        axs[0].imshow(img)
        axs[0].set_title("Input Image", color='white', fontweight='bold')
        axs[0].axis("off")

        axs[1].imshow(seg_output[0, :, :, 0], cmap='gray')
        axs[1].set_title("Segmentation Output", color='white', fontweight='bold')
        axs[1].axis("off")

        # Adjusted plot for classification
        class_probabilities = class_output[0]
        x_pos = np.arange(len(labels))
        bars = axs[2].bar(x_pos, class_probabilities, align='center', alpha=0.7, color='cyan')
        axs[2].set_xticks(x_pos)
        axs[2].set_xticklabels(labels, rotation=45, ha="right", color='white')
        axs[2].set_ylabel('Probability', color='white')
        axs[2].set_ylim([0, 1])
        axs[2].set_title(f'Classification Output\nPredicted Class: {predicted_class}', color='white', fontweight='bold')

        for i, v in enumerate(class_probabilities):
            axs[2].text(i, v + 0.02, f"{v:.2f}", color='white', ha='center')

        plt.tight_layout()
        plt.show()

# Define a callback function to handle the uploaded file
def on_upload(change):
    # Get the uploaded file
    uploaded_filename = next(iter(uploader.value))
    uploaded_file = uploader.value[uploaded_filename]
    # Save the uploaded file
    with open(uploaded_filename, 'wb') as f:
        f.write(uploaded_file['content'])
    # Predict and visualize the image
    predict_image(uploaded_filename)

# Define a function to handle the button click event and clear the output
clear_button = Button(description='Clear')
def clear_output_button_clicked(button):
    out.clear_output()
clear_button.on_click(clear_output_button_clicked)

# Set image size
image_size = 224
# Define labels
labels = ["bridge", "excess", "good", "insuff", "no"]
# Set data folders
data_folders = [
    "/content/gdrive/MyDrive/Deep learning/FYP_2/4 Dataset Ratio 60 20 20/jit012/jit0/b_dip/train",
    "/content/gdrive/MyDrive/Deep learning/FYP_2/4 Dataset Ratio 60 20 20/jit012/jit0/b_dip/val",
    "/content/gdrive/MyDrive/Deep learning/FYP_2/4 Dataset Ratio 60 20 20/jit012/jit0/b_dip/test",]
# Load data
X_data, y_class_labels, y_seg_labels = load_data(data_folders)
# Set image size
image_size = 224
# Split data
train_ratio = 0.2
val_ratio = 0.2
test_ratio = 0.6
X_train, y_train_class, y_train_seg, X_val, y_val_class, y_val_seg, X_test, y_test_class, y_test_seg = split_data(
    X_data, y_class_labels, y_seg_labels, train_ratio, val_ratio, test_ratio)
# Print the counts
print("Number of train images:", len(X_train))
print("Number of train binary masks:", len(y_train_seg))
print("Number of validation images:", len(X_val))
print("Number of validation binary masks:", len(y_val_seg))
print("Number of test images:", len(X_test))
print("Number of test binary masks:", len(y_test_seg))
input_shape = (image_size, image_size, 3)
num_classes = len(labels)
saved_model_path = "/content/gdrive/MyDrive/Deep learning/FYP_2/4 Dataset Ratio 60 20 20/jit012/jit0/b_dip/model/pcbmtl_bdip.h5"
custom_objects = {"segmentation_loss": segmentation_loss}
# Load model
model = load_model(saved_model_path, custom_objects=custom_objects)
# Test model
test_model(model, X_test, y_test_class, y_test_seg, labels, num_images=5)
uploader = widgets.FileUpload(accept='image/*', multiple=False)
uploader.observe(on_upload, names='value')
buttons_layout = HBox([uploader, clear_button])
display(buttons_layout)
display(out)
"
The above is the Python code with Keras to do inference using the multi-task learning model, this model will take the image as input and provide one classification output and one segmentation output, this code let user to select one image and performance inference, however the predicted output always produce the wrong classification output. Check and correct these functions:"def predict_image(image_path):
    # Load and preprocess the image
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
    img_resized = cv2.resize(img, (image_size, image_size))
    img_expanded = np.expand_dims(img_resized, axis=0)
    # Predict using the model
    class_output, seg_output = model.predict(img_expanded)
    # Get the predicted class
    predicted_class_index = np.argmax(class_output)
    predicted_class = labels[predicted_class_index]
    # Clear the previous output
    with out:
        clear_output(wait=True)
        # Configure specific style options for lighter background
        fig, axs = plt.subplots(1, 3, figsize=(15, 5))

        # Lighten background
        bg_color = '#303030'  # Lighter than black for contrast
        fig.patch.set_facecolor(bg_color)
        for ax in axs:
            ax.set_facecolor(bg_color)
            # Set spines color to lighter tone
            ax.spines['bottom'].set_color('grey')
            ax.spines['top'].set_color('grey')
            ax.spines['right'].set_color('grey')
            ax.spines['left'].set_color('grey')
            ax.tick_params(axis='x', colors='white')
            ax.tick_params(axis='y', colors='white')

        axs[0].imshow(img)
        axs[0].set_title("Input Image", color='white', fontweight='bold')
        axs[0].axis("off")

        axs[1].imshow(seg_output[0, :, :, 0], cmap='gray')
        axs[1].set_title("Segmentation Output", color='white', fontweight='bold')
        axs[1].axis("off")

        # Adjusted plot for classification
        class_probabilities = class_output[0]
        x_pos = np.arange(len(labels))
        bars = axs[2].bar(x_pos, class_probabilities, align='center', alpha=0.7, color='cyan')
        axs[2].set_xticks(x_pos)
        axs[2].set_xticklabels(labels, rotation=45, ha="right", color='white')
        axs[2].set_ylabel('Probability', color='white')
        axs[2].set_ylim([0, 1])
        axs[2].set_title(f'Classification Output\nPredicted Class: {predicted_class}', color='white', fontweight='bold')

        for i, v in enumerate(class_probabilities):
            axs[2].text(i, v + 0.02, f"{v:.2f}", color='white', ha='center')

        plt.tight_layout()
        plt.show()

# Define a callback function to handle the uploaded file
def on_upload(change):
    # Get the uploaded file
    uploaded_filename = next(iter(uploader.value))
    uploaded_file = uploader.value[uploaded_filename]
    # Save the uploaded file
    with open(uploaded_filename, 'wb') as f:
        f.write(uploaded_file['content'])
    # Predict and visualize the image
    predict_image(uploaded_filename)

# Define a function to handle the button click event and clear the output
clear_button = Button(description='Clear')
def clear_output_button_clicked(button):
    out.clear_output()
clear_button.on_click(clear_output_button_clicked)"