5_2dunet_trainig.py

# -*- coding: utf-8 -*-
"""5 - 2DUnet_Trainig.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1KDBzA8Y8dOhHTQoR7POPwsGZ3APP-8PU

# IMPORT DATA AND LIBRARIES
"""

!pip install tensorflow==2.16.0rc0
!PYTHONHASHSEED=0

# Import other modules
from matplotlib import pyplot as plt
import zipfile
from shutil import copyfile
from time import time
import numpy as np
import pandas as pd
import random as python_random
import os
import shutil
import glob

# Import TensorFlow/Keras
import tensorflow as tf
import keras
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Conv2D, Conv2DTranspose, MaxPooling2D, Dropout, Activation, concatenate
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Conv3D, Conv3DTranspose, MaxPooling3D, concatenate, Dropout, Activation, BatchNormalization, GroupNormalization
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, CSVLogger
import tensorflow.keras.backend as K


from google.colab import drive
drive.mount('/content/gdrive')

# File's path on Google Drive
drive_zip_path = '/content/gdrive/MyDrive/DL_Project/data/128.zip'
local_extract_path = '/content/Training_Data'
os.makedirs(local_extract_path, exist_ok=True
            )
# Unzip the ZIP file in the local directory
with zipfile.ZipFile(drive_zip_path, 'r') as zip_ref:
    zip_ref.extractall(local_extract_path)

print("Decompressione completata.")

train_img_dir = "/content/Training_Data/X_train/"
train_mask_dir = "/content/Training_Data/Y_train/"
train_img_list = sorted(os.listdir(train_img_dir))
train_mask_list = sorted(os.listdir(train_mask_dir))

val_img_dir = "/content/Training_Data/X_val/"
val_mask_dir = "/content/Training_Data/Y_val/"
val_img_list = sorted(os.listdir(val_img_dir))
val_mask_list = sorted(os.listdir(val_mask_dir))

"""# DATA GENERATOR"""

# Function to load images from a given directory and list of image file names
def load_img(img_dir, img_list):
    images = []  # Initialize an empty list to store the images
    for image_name in img_list:  # Loop through each image name in the list
        if image_name.split('.')[-1] == 'npy':  # Check if the file is a .npy file
            volume = np.load(img_dir + image_name)  # Load the .npy file as a numpy array
            # Extract 128 images along the third axis (z-axis) from each volume
            for j in range(volume.shape[2]):
                images.append(volume[:, :, j, :])  # Append each slice to the images list
    images = np.array(images)  # Convert the list of images to a numpy array
    return images  # Return the numpy array of images

# Generator function to load and yield image and mask batches
def image_loader(img_dir, img_list, mask_dir, mask_list, batch_size):

    L = len(img_list)  # Get the length of the image list
    N = int(batch_size / 128)

    while True:  # Infinite loop to keep yielding batches

        for i in range(0, L, N):  # Iterate through the image list in steps of 2
            if i + (N-1) < L:  # Check if there are at least two more volumes to load
                img = load_img(img_dir, img_list[i:(i+N)])  # Load images from two volumes
                mask = load_img(mask_dir, mask_list[i:(i+N)])  # Load corresponding masks

                # Combine images and masks into a list of tuples and shuffle them
                combined = list(zip(img, mask))
                np.random.shuffle(combined)
                img[:], mask[:] = zip(*combined)  # Unzip the shuffled list back into images and masks

                yield (img, mask)  # Yield the shuffled images and masks as a batch
            else:
                # Handle the case where there is an odd number of volumes
                img = load_img(img_dir, img_list[i:L])  # Load images from the last volume
                mask = load_img(mask_dir, mask_list[i:L])  # Load corresponding masks

                # Combine images and masks into a list of tuples and shuffle them
                combined = list(zip(img, mask))
                np.random.shuffle(combined)
                img[:], mask[:] = zip(*combined)  # Unzip the shuffled list back into images and masks

                yield (img, mask)  # Yield the shuffled images and masks as a batch

"""# FIXED PARAMETERS"""

#Compute weight of classes

columns = ['0','1', '2', '3']
df = pd.DataFrame(columns = columns)
train_mask_list = sorted(glob.glob('/content/Training_Data/Y_train/*.npy'))
for img in range(len(train_mask_list)):

  temp_image=np.load(train_mask_list[img])
  temp_image = np.argmax(temp_image, axis=3)
  val, counts = np.unique(temp_image, return_counts=True)

  conts_dict = {str(i): 0 for i in range(4)}
  for v, c in zip(val, counts):
    conts_dict[str(v)] = c

  row_df = pd.DataFrame([conts_dict])

  # Add new count row
  df = pd.concat([df, row_df], ignore_index=True)

label_0 = df['0'].sum()
label_1 = df['1'].sum()
label_2 = df['2'].sum()
label_3 = df['3'].sum()
total_labels = label_0 + label_1 + label_2 + label_3
n_classes = 4

wt0 = round((total_labels/(n_classes*label_0)), 2) #round to 2 decimals
wt1 = round((total_labels/(n_classes*label_1)), 2)
wt2 = round((total_labels/(n_classes*label_2)), 2)
wt3 = round((total_labels/(n_classes*label_3)), 2)

print(wt0, wt1, wt2, wt3)

#Define loss function and metrics to use during training
CFC_loss = keras.losses.CategoricalFocalCrossentropy(alpha = [wt0, wt1, wt2, wt3])

IoU_0 = keras.metrics.OneHotIoU(num_classes = 4, target_class_ids = [0])
IoU_1 = keras.metrics.OneHotIoU(num_classes = 4, target_class_ids = [1])
IoU_2 = keras.metrics.OneHotIoU(num_classes = 4, target_class_ids = [2])
IoU_3 = keras.metrics.OneHotIoU(num_classes = 4, target_class_ids = [3])

Mean_IoU = keras.metrics.OneHotMeanIoU(num_classes = 4)

metrics = ["accuracy",  Mean_IoU, IoU_0, IoU_1, IoU_2, IoU_3]

#Define the optimizer
batch_size = 256

lr_schedule = keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate = 5e-4,
    decay_steps  = (len(train_img_list) // 2),
    decay_rate = 0.985)

optim = keras.optimizers.Adam(learning_rate = lr_schedule)

"""# MODEL ARCHITECTURE"""

def conv_block_2d(input, n_filters, dropout_rate):
    x = Conv2D(n_filters, 3, padding='same')(input)
    x = Activation('relu')(x)
    x = Dropout(dropout_rate)(x)
    x = Conv2D(n_filters, 3, padding='same')(x)
    x = Activation('relu')(x)
    return x

def encoder_block_2d(input, num_filters, dropout_rate):
    x = conv_block_2d(input, num_filters, dropout_rate)
    p = MaxPooling2D((2, 2))(x)
    return x, p

def decoder_block_2d(input, skip_features, num_filters, dropout_rate):
    x = Conv2DTranspose(num_filters, (2, 2), strides=(2, 2), padding='same')(input)
    x = concatenate([x, skip_features])
    x = conv_block_2d(x, num_filters, dropout_rate)
    return x

def build_unet_2d(input_shape, n_classes):
    inputs = Input(input_shape)

    s1, p1 = encoder_block_2d(inputs, 64, 0)
    s2, p2 = encoder_block_2d(p1, 128, 0)
    s3, p3 = encoder_block_2d(p2, 256, 0)
    s4, p4 = encoder_block_2d(p3, 512, 0)

    b1 = conv_block_2d(p4, 1024, 0)

    d1 = decoder_block_2d(b1, s4, 512, 0)
    d2 = decoder_block_2d(d1, s3, 256, 0)
    d3 = decoder_block_2d(d2, s2, 128, 0)
    d4 = decoder_block_2d(d3, s1, 64, 0)

    outputs = Conv2D(n_classes, 1, padding='same', activation='softmax')(d4)

    model = Model(inputs, outputs, name='2D_U-Net')
    return model

input_shape = (128, 128, 4)  # esempio di input_shape 2D
n_classes = 4

model3 = build_unet_2d(input_shape, n_classes)
model3.compile(optimizer = optim, loss = CFC_loss, metrics = metrics)

"""# TRAIN BASE MODEL (no batch normalization, dropout or regularization)"""

train_img_datagen = image_loader(train_img_dir, train_img_list, train_mask_dir, train_mask_list, batch_size)
val_img_datagen = image_loader(val_img_dir, val_img_list, val_mask_dir, val_mask_list, batch_size)

steps_per_epoch = len(train_img_list) // 2
val_steps_per_epoch = len(val_img_list) // 2

# path to save best model weights
checkpoint_path = "/content/gdrive/MyDrive/DL_Project/2D_64_base.h5.keras"

checkpoint_callback = ModelCheckpoint(
    filepath = checkpoint_path,
    save_best_only = True,
    monitor = 'val_loss',
    mode = 'min',
    verbose = 1
)

early_stopping_callback = EarlyStopping(
    monitor = 'val_loss',    # Track validation loss
    patience = 10,           # Max number of epochs without improvements
    restore_best_weights = True
)

csv = CSVLogger("/content/gdrive/MyDrive/DL_Project/2Dc_64_base.csv")

history = model3.fit(train_img_datagen,
                    steps_per_epoch = steps_per_epoch,
                    epochs = 50,
                    verbose = 1,
                    validation_data = val_img_datagen,
                    validation_steps = val_steps_per_epoch,
                    callbacks = [checkpoint_callback, early_stopping_callback, csv]
)

loss = history.history['loss']
val_loss = history.history['val_loss']
epochs = range(1, len(loss) + 1)
plt.plot(epochs, loss, 'y', label='Training loss')
plt.plot(epochs, val_loss, 'r', label='Validation loss')
plt.title('Training and validation loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()

acc = history.history['accuracy']
val_acc = history.history['val_accuracy']

plt.plot(epochs, acc, 'y', label='Training Accuracy')
plt.plot(epochs, val_acc, 'r', label='Validation Accuracy')
plt.title('Training and validation Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
plt.show()



a1 = history.history['one_hot_io_u']
a2 = history.history['val_one_hot_io_u']
b1 = history.history['one_hot_io_u_1']
b2 = history.history['val_one_hot_io_u_1']
c1 = history.history['one_hot_io_u_2']
c2 = history.history['val_one_hot_io_u_2']
d1 = history.history['one_hot_io_u_3']
d2 = history.history['val_one_hot_io_u_3']
e1 = history.history['one_hot_mean_io_u']
e2 = history.history['val_one_hot_mean_io_u']

colors = ['b', 'g', 'c', 'm', 'y', 'k']
line_styles = ['-', '--', '-.', ':']

plt.figure(figsize=(12, 8))

plt.plot(epochs, a1, color=colors[0], linestyle=line_styles[0], label='Training IoU (0)')
plt.plot(epochs, a2, color=colors[0], linestyle=line_styles[1], label='Validation IoU (0)')
plt.plot(epochs, b1, color=colors[1], linestyle=line_styles[0], label='Training IoU (1)')
plt.plot(epochs, b2, color=colors[1], linestyle=line_styles[1], label='Validation IoU (1)')
plt.plot(epochs, c1, color=colors[2], linestyle=line_styles[0], label='Training IoU (2)')
plt.plot(epochs, c2, color=colors[2], linestyle=line_styles[1], label='Validation IoU (2)')
plt.plot(epochs, d1, color=colors[3], linestyle=line_styles[0], label='Training IoU (3)')
plt.plot(epochs, d2, color=colors[3], linestyle=line_styles[1], label='Validation IoU (3)')
plt.plot(epochs, e1, color=colors[4], linestyle=line_styles[0], label='Training Mean IoU')
plt.plot(epochs, e2, color=colors[4], linestyle=line_styles[1], label='Validation Mean IoU')

plt.title('Training and validation IoU')
plt.xlabel('Epochs')
plt.ylabel('IoU')
plt.legend()
plt.show()

#We try to continue training until epoch 100
history = model3.fit(train_img_datagen,
                    steps_per_epoch = steps_per_epoch,
                    initial_epoch = 50,
                    epochs = 100,
                    verbose = 1,
                    validation_data = val_img_datagen,
                    validation_steps = val_steps_per_epoch,
                    callbacks = [checkpoint_callback, early_stopping_callback, csv]
)

loss = history.history['loss']
val_loss = history.history['val_loss']
epochs = range(1, len(loss) + 1)
plt.plot(epochs, loss, 'y', label='Training loss')
plt.plot(epochs, val_loss, 'r', label='Validation loss')
plt.title('Training and validation loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()

acc = history.history['accuracy']
val_acc = history.history['val_accuracy']

plt.plot(epochs, acc, 'y', label='Training Accuracy')
plt.plot(epochs, val_acc, 'r', label='Validation Accuracy')
plt.title('Training and validation Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
plt.show()



a1 = history.history['one_hot_io_u']
a2 = history.history['val_one_hot_io_u']
b1 = history.history['one_hot_io_u_1']
b2 = history.history['val_one_hot_io_u_1']
c1 = history.history['one_hot_io_u_2']
c2 = history.history['val_one_hot_io_u_2']
d1 = history.history['one_hot_io_u_3']
d2 = history.history['val_one_hot_io_u_3']
e1 = history.history['one_hot_mean_io_u']
e2 = history.history['val_one_hot_mean_io_u']

colors = ['b', 'g', 'c', 'm', 'y', 'k']
line_styles = ['-', '--', '-.', ':']

plt.figure(figsize=(12, 8))

plt.plot(epochs, a1, color=colors[0], linestyle=line_styles[0], label='Training IoU (0)')
plt.plot(epochs, a2, color=colors[0], linestyle=line_styles[1], label='Validation IoU (0)')
plt.plot(epochs, b1, color=colors[1], linestyle=line_styles[0], label='Training IoU (1)')
plt.plot(epochs, b2, color=colors[1], linestyle=line_styles[1], label='Validation IoU (1)')
plt.plot(epochs, c1, color=colors[2], linestyle=line_styles[0], label='Training IoU (2)')
plt.plot(epochs, c2, color=colors[2], linestyle=line_styles[1], label='Validation IoU (2)')
plt.plot(epochs, d1, color=colors[3], linestyle=line_styles[0], label='Training IoU (3)')
plt.plot(epochs, d2, color=colors[3], linestyle=line_styles[1], label='Validation IoU (3)')
plt.plot(epochs, e1, color=colors[4], linestyle=line_styles[0], label='Training Mean IoU')
plt.plot(epochs, e2, color=colors[4], linestyle=line_styles[1], label='Validation Mean IoU')

plt.title('Training and validation IoU')
plt.xlabel('Epochs')
plt.ylabel('IoU')
plt.legend()
plt.show()

model3.save('/content/gdrive/MyDrive/DL_Project/2D_Unet/64_Base.keras')

"""# ADD BATCH NORMALIZATION"""

#Define the optimizer
batch_size = 256

lr_schedule = keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate = 5e-4,
    decay_steps  = (len(train_img_list) // 4),
    decay_rate = 0.985)

optim = keras.optimizers.Adam(learning_rate = lr_schedule)

#Modify convolutional block adding batch normalization
def conv_block_2d(input, n_filters, dropout_rate):
    x = Conv2D(n_filters, 3, padding='same')(input)
    X = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = Dropout(dropout_rate)(x)
    x = Conv2D(n_filters, 3, padding='same')(x)
    X = BatchNormalization()(x)
    x = Activation('relu')(x)
    return x

#Compile new model
model2 = build_unet_2d(input_shape, n_classes)
model2.compile(optimizer = optim, loss = CFC_loss, metrics = metrics)
model2.summary()

train_img_datagen = image_loader(train_img_dir, train_img_list, train_mask_dir, train_mask_list, batch_size)
val_img_datagen = image_loader(val_img_dir, val_img_list, val_mask_dir, val_mask_list, batch_size)

steps_per_epoch = len(train_img_list) // 2
val_steps_per_epoch = len(val_img_list) // 2

# path to save best model weights
checkpoint_path = "/content/gdrive/MyDrive/DL_Project/2D_UNet_64_BatchNorm.h5.keras"

checkpoint_callback = ModelCheckpoint(
    filepath = checkpoint_path,
    save_best_only = True,
    monitor = 'val_loss',
    mode = 'min',
    verbose = 1
)

early_stopping_callback = EarlyStopping(
    monitor = 'val_loss',    # Track validation loss
    patience = 10,           # Max number of epochs without improvements
    restore_best_weights = True
)

csv = CSVLogger("/content/gdrive/MyDrive/DL_Project/history64_BatchNorm.csv")

history = model2.fit(train_img_datagen,
                    steps_per_epoch = steps_per_epoch,
                    epochs = 30,
                    verbose = 1,
                    validation_data = val_img_datagen,
                    validation_steps = val_steps_per_epoch,
                    callbacks = [checkpoint_callback, early_stopping_callback, csv]
)

loss = history.history['loss']
val_loss = history.history['val_loss']
epochs = range(1, len(loss) + 1)
plt.plot(epochs, loss, 'y', label='Training loss')
plt.plot(epochs, val_loss, 'r', label='Validation loss')
plt.title('Training and validation loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()

acc = history.history['accuracy']
val_acc = history.history['val_accuracy']

plt.plot(epochs, acc, 'y', label='Training Accuracy')
plt.plot(epochs, val_acc, 'r', label='Validation Accuracy')
plt.title('Training and validation Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
plt.show()



a1 = history.history['one_hot_io_u']
a2 = history.history['val_one_hot_io_u']
b1 = history.history['one_hot_io_u_1']
b2 = history.history['val_one_hot_io_u_1']
c1 = history.history['one_hot_io_u_2']
c2 = history.history['val_one_hot_io_u_2']
d1 = history.history['one_hot_io_u_3']
d2 = history.history['val_one_hot_io_u_3']
e1 = history.history['one_hot_mean_io_u']
e2 = history.history['val_one_hot_mean_io_u']

colors = ['b', 'g', 'c', 'm', 'y', 'k']
line_styles = ['-', '--', '-.', ':']

plt.figure(figsize=(12, 8))

plt.plot(epochs, a1, color=colors[0], linestyle=line_styles[0], label='Training IoU (0)')
plt.plot(epochs, a2, color=colors[0], linestyle=line_styles[1], label='Validation IoU (0)')
plt.plot(epochs, b1, color=colors[1], linestyle=line_styles[0], label='Training IoU (1)')
plt.plot(epochs, b2, color=colors[1], linestyle=line_styles[1], label='Validation IoU (1)')
plt.plot(epochs, c1, color=colors[2], linestyle=line_styles[0], label='Training IoU (2)')
plt.plot(epochs, c2, color=colors[2], linestyle=line_styles[1], label='Validation IoU (2)')
plt.plot(epochs, d1, color=colors[3], linestyle=line_styles[0], label='Training IoU (3)')
plt.plot(epochs, d2, color=colors[3], linestyle=line_styles[1], label='Validation IoU (3)')
plt.plot(epochs, e1, color=colors[4], linestyle=line_styles[0], label='Training Mean IoU')
plt.plot(epochs, e2, color=colors[4], linestyle=line_styles[1], label='Validation Mean IoU')

plt.title('Training and validation IoU')
plt.xlabel('Epochs')
plt.ylabel('IoU')
plt.legend()
plt.show()

"""# ADD BATCH NORM AND 2D SPATIAL DROPOUT"""

#Define the optimizer
batch_size = 256

lr_schedule = keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate = 5e-4,
    decay_steps  = (len(train_img_list) // 2),
    decay_rate = 0.985)

optim = keras.optimizers.Adam(learning_rate = lr_schedule)

from tensorflow.keras.layers import SpatialDropout2D

def conv_block_2d(input, n_filters):
    x = Conv2D(n_filters, 3, padding='same')(input)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = SpatialDropout2D(0.3)(x)
    x = Conv2D(n_filters, 3, padding='same')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    return x

def encoder_block_2d(input, num_filters):
    x = conv_block_2d(input, num_filters)
    p = MaxPooling2D((2, 2))(x)
    return x, p

def decoder_block_2d(input, skip_features, num_filters):
    x = Conv2DTranspose(num_filters, (2, 2), strides=(2, 2), padding='same')(input)
    x = concatenate([x, skip_features])
    x = conv_block_2d(x, num_filters)
    return x

def build_unet_2d(input_shape, n_classes):
    inputs = Input(input_shape)

    s1, p1 = encoder_block_2d(inputs, 64)
    s2, p2 = encoder_block_2d(p1, 128)
    s3, p3 = encoder_block_2d(p2, 256)
    s4, p4 = encoder_block_2d(p3, 512)

    b1 = conv_block_2d(p4, 1024)

    d1 = decoder_block_2d(b1, s4, 512)
    d2 = decoder_block_2d(d1, s3, 256)
    d3 = decoder_block_2d(d2, s2, 128)
    d4 = decoder_block_2d(d3, s1, 64)

    outputs = Conv2D(n_classes, 1, padding='same', activation='softmax')(d4)

    model = Model(inputs, outputs, name='2D_U-Net')
    return model

input_shape = (128, 128, 4)
n_classes = 4

model4 = build_unet_2d(input_shape, n_classes)
model4.compile(optimizer = optim, loss = CFC_loss, metrics = metrics)

train_img_datagen = image_loader(train_img_dir, train_img_list, train_mask_dir, train_mask_list, batch_size)
val_img_datagen = image_loader(val_img_dir, val_img_list, val_mask_dir, val_mask_list, batch_size)

steps_per_epoch = len(train_img_list) // 2
val_steps_per_epoch = len(val_img_list) // 2

# path to save best model weights
checkpoint_path = "/content/gdrive/MyDrive/DL_Project/2D_64_Complex.h5.keras"

checkpoint_callback = ModelCheckpoint(
    filepath = checkpoint_path,
    save_best_only = True,
    monitor = 'val_loss',
    mode = 'min',
    verbose = 1
)

early_stopping_callback = EarlyStopping(
    monitor = 'val_loss',    # Track validation loss
    patience = 10,           # Max number of epochs without improvements
    restore_best_weights = True
)

csv = CSVLogger("/content/gdrive/MyDrive/DL_Project/2D_64_Complex.csv")

history = model4.fit(train_img_datagen,
                    steps_per_epoch = steps_per_epoch,
                    epochs = 50,
                    verbose = 1,
                    validation_data = val_img_datagen,
                    validation_steps = val_steps_per_epoch,
                    callbacks = [checkpoint_callback, early_stopping_callback, csv]
)