amal90888/unet-segmentation-model

UNet Model for COVID-19 CT Scan Segmentation

Model Overview

This UNet-based segmentation model is designed for automated segmentation of COVID-19 infected lung regions in CT scans. It enhances the classic U-Net with attention mechanisms to improve focus on infected regions.

• —Architecture: UNet + Attention Gates
• —Dataset: COVID-19 CT scans from Coronacases.org, Radiopaedia.org, and Zenodo Repository
• —Task: Image Segmentation (Lung Infection)
• —Metrics: Dice Coefficient, IoU, Hausdorff Distance, ASSD

Training Details

• —Dataset Size: 20 CT scans (512 × 512 × 301 slices)
• —Preprocessing:
• —Normalization of pixel intensities [0,1]
• —HU Thresholding: [-1000, 1500]
• —Image resizing to 128 × 128 pixels
• —Binarization of masks (0 = background, 1 = infected regions)
• —Augmentation:
• —Rotations: ±5 degrees
• —Elastic transformations, Gaussian blur
• —Brightness/contrast variations
• —Final dataset: 2,252 CT slices
• —Training:
• —Optimizer: Adam (learning rate = 1e-4)
• —Loss Function: Weighted BCE-Dice Loss + Surface Loss
• —Batch Size: 16
• —Epochs: 25
• —Training Platform: NVIDIA Tesla T4 (Google Colab Pro)

Model Performance

📌 Key Findings: ✔ Augmentation improved segmentation accuracy significantly ✔ Attention U-Net outperformed other segmentation models

How to Use the Model

Load the Model

TensorFlow/Keras

import os
from huggingface_hub import hf_hub_download
from tensorflow.keras.models import load_model
from keras.saving import register_keras_serializable
import tensorflow.keras.backend as K

#  Set Keras backend (optional)
os.environ["KERAS_BACKEND"] = "jax"

#  Register and define missing functions
@register_keras_serializable()
def dice_coef(y_true, y_pred, smooth=1e-6):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)

@register_keras_serializable()
def gl_sl(*args, **kwargs):
    pass  # Placeholder function (update if needed)

#  Download the model from Hugging Face
model_path = hf_hub_download(repo_id="amal90888/unet-segmentation-model", filename="unet_model.keras")

#  Load the model with registered custom objects
unet = load_model(model_path, custom_objects={"dice_coef": dice_coef, "gl_sl": gl_sl}, compile=False)

#  Recompile with fresh optimizer and correct loss function
from tensorflow.keras.optimizers import Adam
unet.compile(optimizer=Adam(learning_rate=1e-4), loss="binary_crossentropy", metrics=["accuracy", dice_coef])

print(" Model loaded and recompiled successfully!")