The code:
"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 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 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, 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)

# 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 (
    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):
# Convert one-hot encoded labels back to label encoded
y_class_labels = np.argmax(y_class_categorical, axis=1)
# Convert label encoded labels back to original class names
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 = 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()
"
How can I implement the focal loss to my code for the classification branch only to my MTL model?
The TensorFlow code:
"
tf.keras.losses.CategoricalFocalCrossentropy(
alpha=0.25,
gamma=2.0,
from_logits=False,
label_smoothing=0.0,
axis=-1,
reduction=losses_utils.ReductionV2.AUTO,
name='categorical_focal_crossentropy'
)
"