2_3d_unet_basicimplementation.py

# -*- coding: utf-8 -*-
"""2-3D_Unet_BasicImplementation.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1aCFnq2ajiHTLDrJSUMeqEuO07IHXSMqt
"""

!pip install tensorflow==2.16.0rc0
!pip install segmentation_models_3D
!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 random as python_random
import os

# Import TensorFlow/Keras
import keras
import tensorflow as tf
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
import tensorflow.keras.backend as K

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

# Percorso del file ZIP su Google Drive
# drive_zip_path = '/content/gdrive/MyDrive/DL_PROJECT/'
# Local path
local_zip_path = '/content/gdrive/MyDrive/DL_PROJECT/Data/128.zip'

local_extract_path = '/content/Training_Data'
os.makedirs(local_extract_path, exist_ok=True)

# Decomprimere il file ZIP nella directory locale
with zipfile.ZipFile(local_zip_path, 'r') as zip_ref:
    zip_ref.extractall(local_extract_path)

print("Decompressione completata.")

"""### Defining Data Generator"""

def load_img(img_dir, img_list):
    images=[]
    for i, image_name in enumerate(img_list):
        if (image_name.split('.')[1] == 'npy'):

            image = np.load(img_dir + image_name)

            images.append(image)

    images = np.array(images)

    return images

def image_loader(img_dir, img_list, mask_dir, mask_list, batch_size):

    L = len(img_list)

    while True:

        batch_start = 0
        batch_end = batch_size

        while batch_start < L:

            limit = min(L, batch_end)

            img = load_img(img_dir, img_list[batch_start:limit])
            mask = load_img(mask_dir, mask_list[batch_start:limit])

            yield(img, mask)

        batch_start += batch_size
        batch_end += batch_size

"""#### Testing data generator"""

from matplotlib import pyplot as plt
import random

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

batch_size = 4

train_img_datagen = image_loader(train_img_dir, train_img_list, train_mask_dir, train_mask_list, batch_size)

img, mask = train_img_datagen.__next__()

img_num = random.randint(0,img.shape[0]-1)
test_img=img[img_num]
test_mask=mask[img_num]
test_mask=np.argmax(test_mask, axis=3)

n_slice=random.randint(0, test_mask.shape[2])
plt.figure(figsize=(12, 8))

plt.subplot(221)
plt.imshow(test_img[:,:,n_slice, 0], cmap='gray')
plt.title('Image flair')
plt.subplot(222)
plt.imshow(test_img[:,:,n_slice, 1], cmap='gray')
plt.title('Image t1ce')
plt.subplot(223)
plt.imshow(test_img[:,:,n_slice, 2], cmap='gray')
plt.title('Image t2')
plt.subplot(224)
plt.imshow(test_mask[:,:,n_slice])
plt.title('Mask')
plt.show()

"""### Loss, Metrics, optimizer etc.."""

import pandas as pd
import glob

### COUNT OF WEIGTHS
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)):

  #print(img)
  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])

  # Uso di pd.concat per aggiungere la nuova riga
  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)

import tensorflow.keras.backend as K

def DiceCoefficient(y_true, y_pred, smooth = 1e-6):

    y_pred = tf.keras.activations.softmax(y_pred, axis = -1)

    # Cast to float32 datatype
    y_true = K.cast(y_true, 'float32')
    y_pred = K.cast(y_pred, 'float32')

    # Flatten label and prediction tensors
    inputs = K.flatten(y_pred)
    targets = K.flatten(y_true)

    intersection = K.sum(inputs * targets)
    dice = (2 * intersection + smooth) / (K.sum(targets) + K.sum(inputs) + smooth)
    return dice

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", DiceCoefficient, Mean_IoU, IoU_0, IoU_1, IoU_2, IoU_3]

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

batch_size = 4

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

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

print(steps_per_epoch, val_steps_per_epoch)

lr_schedule = keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate = 1e-3,
    decay_steps  = (len(train_img_list) // batch_size),
    decay_rate = 0.95)

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

"""### DEFINE MODEL ARCHITECTURE - 3D U-Net"""

def conv_block(input, n_filters, dropout_rate):

    x = Conv3D(n_filters, 3, padding='same')(input)
    x = Activation('relu')(x)
    x = Dropout(dropout_rate)(x)

    x = Conv3D(n_filters, 3, padding='same')(x)
    x = Activation('relu')(x)

    return x

def encoder_block(input, num_filters, dropout_rate):

    x = conv_block(input, num_filters, dropout_rate)
    p = MaxPooling3D((2, 2, 2))(x)

    return x, p

def decoder_block(input, skip_features, num_filters, dropout_rate):

    x = Conv3DTranspose(num_filters, (2, 2, 2), strides=(2, 2, 2), padding='same')(input)
    x = concatenate([x, skip_features])
    x = conv_block(x, num_filters, dropout_rate)

    return x

def build_unet(input_shape, n_classes, num_filters):
    inputs = Input(input_shape)

    s1, p1 = encoder_block(inputs, num_filters, 0.1)
    s2, p2 = encoder_block(p1, 2*num_filters, 0.1)
    s3, p3 = encoder_block(p2, 4*num_filters, 0.2)
    s4, p4 = encoder_block(p3, 8*num_filters, 0.2)

    b1 = conv_block(p4, 16*num_filters, 0.3)

    d1 = decoder_block(b1, s4, 8*num_filters, 0.2)
    d2 = decoder_block(d1, s3, 4*num_filters, 0.2)
    d3 = decoder_block(d2, s2, 2*num_filters, 0.1)
    d4 = decoder_block(d3, s1, num_filters, 0.1)

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

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

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

model = build_unet(input_shape, n_classes, 32)
model.compile(optimizer = optim, loss = CFC_loss, metrics = metrics)
model.summary()

"""### **Compile and fit**"""

# Callback per salvare il miglior modello basato sulla loss di validazione
checkpoint_path = "/content/gdrive/MyDrive/DL_PROJECT/3D_Unet_first.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',    # Monitorare la loss di validazione
    patience = 10,           # Numero di epoche con nessun miglioramento dopo le quali l'addestramento verrà fermato
    restore_best_weights = True # Ripristina i pesi del miglior modello visto durante l'addestramento
  )

csv = keras.callbacks.CSVLogger("/content/gdrive/MyDrive/DL_PROJECT/hist_3D_first", separator=",", append = False)

batch_size = 4

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

print(steps_per_epoch, val_steps_per_epoch)

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

train_mask_list

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)

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

Unet_3D.save("/content/gdrive/MyDrive/DL_PROJECT/3_D_first.keras")

Unet_3D.save_weights("/content/gdrive/MyDrive/DL_PROJECT/3_D_first.weights.h5", overwrite=True)

history.history["loss"]

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

history.history.keys()

history.history['one_hot_mean_io_u_1']
val_acc = history.history['val_one_hot_mean_io_u_1']

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

lr_schedule = keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate = 1e-3,
    decay_steps  = (len(train_img_list) // batch_size),
    decay_rate = 0.95)

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

def build_unet(input_shape, n_classes, num_filters):
    inputs = Input(input_shape)

    s1, p1 = encoder_block(inputs, num_filters, 0.3)
    s2, p2 = encoder_block(p1, 2*num_filters, 0.3)
    s3, p3 = encoder_block(p2, 4*num_filters, 0.3)

    b1 = conv_block(p3, 8*num_filters, 0.3)

    d1 = decoder_block(b1, s3, 4*num_filters, 0.3)
    d2 = decoder_block(d1, s2, 2*num_filters, 0.3)
    d3 = decoder_block(d2, s1, num_filters, 0.3)

    outputs = Conv3D(n_classes, 1, padding = 'same', activation='softmax')(d3)

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

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

model_2 = build_unet(input_shape, n_classes, 16)
model_2.compile(optimizer = optim, loss = CFC_loss, metrics = metrics)
model_2.summary()

history_2=model_2.fit(train_img_datagen,
                  steps_per_epoch=steps_per_epoch,
                  epochs=20,
                  verbose=1,
                  validation_data=val_img_datagen,
                  validation_steps=val_steps_per_epoch,
                  callbacks = [checkpoint_callback, early_stopping_callback, csv]
)