"
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
import ipywidgets as widgets
from ipywidgets import Output, Button
from warnings import filterwarnings
from IPython.display import display, clear_output
from tensorflow.keras.preprocessing import image
from tensorflow.keras.optimizers import Adam
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.utils import to_categorical
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau, TensorBoard, ModelCheckpoint
from sklearn.metrics import classification_report, confusion_matrix
from google.colab import drive
drive.mount("/content/gdrive")
def load_data(data_folders):
    X_data = []  
    y_class_labels = []  
    y_seg_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) 
                    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, class_data_counts):
    X_train = []
    y_train_class = []
    y_train_seg = []
    X_val = []
    y_val_class = []
    y_val_seg = []
    X_test = []
    y_test_class = []
    y_test_seg = []
    for label, count in class_data_counts.items():
        label_indices = np.where(y_class_labels == label)[0]
        class_X_data = X_data[label_indices]
        class_y_class_labels = y_class_labels[label_indices]
        class_y_seg_labels = y_seg_labels[label_indices]
        train_count = count[0]
        val_count = count[1]
        test_count = count[2]
        class_X_train = class_X_data[:train_count]
        class_y_train_class = class_y_class_labels[:train_count]
        class_y_train_seg = class_y_seg_labels[:train_count]
        class_X_val = class_X_data[train_count: train_count + val_count]
        class_y_val_class = class_y_class_labels[train_count: train_count + val_count]
        class_y_val_seg = class_y_seg_labels[train_count: train_count + val_count]
        class_X_test = class_X_data[train_count + val_count: train_count + val_count + test_count]
        class_y_test_class = class_y_class_labels[train_count + val_count: train_count + val_count + test_count]
        class_y_test_seg = class_y_seg_labels[train_count + val_count: train_count + val_count + test_count]
        X_train.extend(class_X_train)
        y_train_class.extend(class_y_train_class)
        y_train_seg.extend(class_y_train_seg)
        X_val.extend(class_X_val)
        y_val_class.extend(class_y_val_class)
        y_val_seg.extend(class_y_val_seg)
        X_test.extend(class_X_test)
        y_test_class.extend(class_y_test_class)
        y_test_seg.extend(class_y_test_seg)
    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 (
        np.array(X_train),
        np.array(y_train_class_categorical),
        np.array(y_train_seg),
        np.array(X_val),
        np.array(y_val_class_categorical),
        np.array(y_val_seg),
        np.array(X_test),
        np.array(y_test_class_categorical),
        np.array(y_test_seg),
    )

def count_labels(y_class_categorical, label_encoder):
    y_class_labels = np.argmax(y_class_categorical, axis=1)
    y_class_names = label_encoder.inverse_transform(y_class_labels)
    unique, counts = np.unique(y_class_names, return_counts=True)
    return dict(zip(unique, counts))

def build_model(input_shape, num_classes):
    num_filter = 16  
    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)
    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)
    concat1 = concatenate([up1, relu8], axis=-1)  
    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)  
    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)  
    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)  
    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_output = Conv2D(1, 1, activation="sigmoid", name="segmentation_output")(relu18) 
    gap1 = GlobalAveragePooling2D()(relu8)
    gap2 = GlobalAveragePooling2D()(relu10)
    gap3 = GlobalAveragePooling2D()(relu12)
    conv20 = Conv2D(num_filter * 4, 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(num_filter * 4, 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(num_filter * 1, 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)
    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)
    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

Categorical_Focal_loss = tf.keras.losses.CategoricalFocalCrossentropy(
        alpha=0.25,
        gamma=2.0,
        from_logits=False,
        label_smoothing=0.0,
        axis=-1,)

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(learning_rate=0.001),
        loss={"classification_output": Categorical_Focal_loss, "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}):
        best_model = load_model("multitask_best_weights.h5")
        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)
        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)
        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_class_acc

def plot_performance(history):
    bg_color = '#303030'  
    text_color = 'white' 
    plt.rcParams['figure.facecolor'] = bg_color
    plt.rcParams['axes.facecolor'] = bg_color
    plt.rcParams['axes.edgecolor'] = 'grey'
    plt.rcParams['axes.labelcolor'] = text_color
    plt.rcParams['text.color'] = text_color
    plt.rcParams['xtick.color'] = text_color
    plt.rcParams['ytick.color'] = text_color

    # 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), dpi=200)
    plt.plot(classification_train_accuracy, label="Training Accuracy")
    plt.plot(classification_val_accuracy, label="Validation Accuracy")
    plt.title("Classification Accuracy", fontweight='bold')
    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), dpi=200)
    plt.plot(classification_train_loss, "b", label="Training Loss")
    plt.plot(classification_val_loss, "r", label="Validation Loss")
    plt.title("Classification Loss", fontweight='bold')
    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), dpi=200)
    plt.plot(segmentation_train_accuracy, label="Training Accuracy")
    plt.plot(segmentation_val_accuracy, label="Validation Accuracy")
    plt.title("Segmentation Accuracy", fontweight='bold')
    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), dpi=200)
    plt.plot(segmentation_train_loss, "b", label="Training Loss")
    plt.plot(segmentation_val_loss, "r", label="Validation Loss")
    plt.title("Segmentation Loss", fontweight='bold')
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend()
    plt.show()

    # Reset to defaults for other plots that may not want these settings
    plt.rcParams['figure.facecolor'] = 'white'
    plt.rcParams['axes.facecolor'] = 'white'
    plt.rcParams['axes.edgecolor'] = 'black'
    plt.rcParams['axes.labelcolor'] = 'black'
    plt.rcParams['text.color'] = 'black'
    plt.rcParams['xtick.color'] = 'black'
    plt.rcParams['ytick.color'] = 'black'

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'  
    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()

out = Output()
def predict_image(image_path):
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  
    img_resized = cv2.resize(img, (image_size, image_size))
    img_expanded = np.expand_dims(img_resized, axis=0)
    class_output, seg_output = model.predict(img_expanded)
    predicted_class_index = np.argmax(class_output)
    predicted_class = labels[predicted_class_index]
    with out:
        clear_output(wait=True)
        fig, axs = plt.subplots(1, 3, figsize=(15, 5))
        bg_color = '#303030'  
        fig.patch.set_facecolor(bg_color)
        for ax in axs:
            ax.set_facecolor(bg_color)
            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")
        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()

def on_upload(change):
    uploaded_filename = next(iter(uploader.value))
    uploaded_file = uploader.value[uploaded_filename]
    with open(uploaded_filename, 'wb') as f:
        f.write(uploaded_file['content'])
    predict_image(uploaded_filename)

clear_button = Button(description='Clear')

def clear_output_button_clicked(button):
    out.clear_output()
clear_button.on_click(clear_output_button_clicked)
image_size = 224
labels = ["bridge", "excess", "good", "insuff", "no"]
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",]
X_data, y_class_labels, y_seg_labels = load_data(data_folders)
class_data_counts = {
    "bridge": [80, 80, 80],
    "excess": [80, 80, 80],
    "good": [80, 80, 80],
    "insuff": [80, 80, 80],
    "no": [80, 80, 80]}
    
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, class_data_counts)
label_encoder = LabelEncoder()
label_encoder.fit(y_class_labels)
train_counts = count_labels(y_train_class, label_encoder)
val_counts = count_labels(y_val_class, label_encoder)
test_counts = count_labels(y_test_class, label_encoder)
print("Train counts:     ", train_counts,"  Total in train set:", sum(train_counts.values()))
print("Validation counts:", val_counts, "  Total in validation set:", sum(val_counts.values()))
print("Test counts:      ", test_counts,"  Total in test set:", sum(test_counts.values()))
input_shape = (image_size, image_size, 3)
num_classes = len(labels)
model = build_model(input_shape, num_classes)
model.summary()
model = build_model(input_shape, num_classes)
batch_size = 16
epochs = 2
history = train_model(model, X_train, y_train_class, y_train_seg, X_val, y_val_class, y_val_seg, batch_size, epochs)
test_class_acc = evaluate_model(model, X_test, y_test_class, y_test_seg)
plot_performance(history)
"
The above is the Python code with Keras to train the multi-task learning model, this model will take the image as input and provide one classification output and one segmentation output. I want to plot some images with their segmentation masks and labels randomly before training (at least one for each class). Use '#303030' as background color.