Spaces:
Runtime error
Runtime error
| """ | |
| model.py | |
| -------- | |
| Model definitions for Real-World Lung CT Analysis: | |
| 1. load_pretrained_model() -- existing 2D inference model | |
| 2. build_unet_3d() -- [NEW] 3D volumetric segmentation model | |
| 3. build_nodule_detector() -- [NEW] Nodule classification CNN | |
| 4. Custom losses and metrics (Dice, IoU, Focal Loss for class imbalance) | |
| """ | |
| import tensorflow as tf | |
| from tensorflow import keras | |
| from tensorflow.keras import layers | |
| # =========================================================================== | |
| # LOSSES & METRICS | |
| # =========================================================================== | |
| def dice_coefficient(y_true, y_pred, smooth=1.0): | |
| """Dice Similarity Coefficient (works for 2D and 3D).""" | |
| y_true_f = keras.backend.flatten(tf.cast(y_true, tf.float32)) | |
| y_pred_f = keras.backend.flatten(tf.cast(y_pred, tf.float32)) | |
| intersection = tf.reduce_sum(y_true_f * y_pred_f) | |
| return (2.0 * intersection + smooth) / ( | |
| tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) + smooth | |
| ) | |
| def dice_loss(y_true, y_pred): | |
| return 1.0 - dice_coefficient(y_true, y_pred) | |
| def iou_score(y_true, y_pred, smooth=1.0): | |
| """Intersection over Union (Jaccard Index).""" | |
| y_true_f = keras.backend.flatten(tf.cast(y_true, tf.float32)) | |
| y_pred_f = keras.backend.flatten(tf.cast(y_pred, tf.float32)) | |
| intersection = tf.reduce_sum(y_true_f * y_pred_f) | |
| union = (tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) - intersection) | |
| return (intersection + smooth) / (union + smooth) | |
| def combined_loss(y_true, y_pred): | |
| """BCE + Dice Loss for highly accurate segmentation boundaries.""" | |
| bce = keras.losses.binary_crossentropy( | |
| tf.cast(y_true, tf.float32), | |
| tf.cast(y_pred, tf.float32) | |
| ) | |
| dl = dice_loss(y_true, y_pred) | |
| return 0.5 * bce + 0.5 * dl | |
| def focal_loss(gamma=2.0, alpha=0.25): | |
| """ | |
| Focal Loss handles severe class imbalance (e.g., millions of healthy lung | |
| pixels vs a few nodule pixels). It down-weights well-classified examples. | |
| """ | |
| def focal_loss_fixed(y_true, y_pred): | |
| pt_1 = tf.where(tf.equal(y_true, 1), y_pred, tf.ones_like(y_pred)) | |
| pt_0 = tf.where(tf.equal(y_true, 0), y_pred, tf.zeros_like(y_pred)) | |
| return -keras.backend.sum(alpha * keras.backend.pow(1. - pt_1, gamma) * keras.backend.log(pt_1 + keras.backend.epsilon())) \ | |
| -keras.backend.sum((1 - alpha) * keras.backend.pow(pt_0, gamma) * keras.backend.log(1. - pt_0 + keras.backend.epsilon())) | |
| return focal_loss_fixed | |
| # =========================================================================== | |
| # 2D SEGMENTATION | |
| # =========================================================================== | |
| def load_pretrained_model(model_path="LungCT_Seg.keras"): | |
| """Load the original 2D Keras model.""" | |
| return keras.models.load_model(model_path, compile=False) | |
| # =========================================================================== | |
| # 3D VOLUMETRIC SEGMENTATION [NEW] | |
| # =========================================================================== | |
| def conv_block_3d(x, filters, kernel_size=3, dropout_rate=0.1): | |
| """3D Convolutional Block: Conv3D -> BN -> ReLU -> Dropout -> Conv3D""" | |
| x = layers.Conv3D(filters, kernel_size, padding="same", use_bias=False)(x) | |
| x = layers.BatchNormalization()(x) | |
| x = layers.Activation("relu")(x) | |
| x = layers.Dropout(dropout_rate)(x) | |
| x = layers.Conv3D(filters, kernel_size, padding="same", use_bias=False)(x) | |
| x = layers.BatchNormalization()(x) | |
| x = layers.Activation("relu")(x) | |
| return x | |
| def build_unet_3d(input_shape=(64, 64, 64, 1), base_filters=16, depth=3): | |
| """ | |
| 3D U-Net for Volumetric CT Segmentation. | |
| Unlike 2D models that process slice-by-slice, a 3D U-Net processes a cubic | |
| volume (patch) of the scan. This preserves spatial context across the Z-axis | |
| (depth), which is crucial for distinguishing spherical nodules from cylindrical | |
| blood vessels. | |
| Args: | |
| input_shape : tuple, (D, H, W, C). Default 64x64x64 patches. | |
| base_filters : int, filters in first block. Kept small (16) to fit in GPU RAM. | |
| depth : int, number of downsampling levels. | |
| Returns: | |
| model : keras.Model (Output shape = input_shape) | |
| """ | |
| inputs = keras.Input(shape=input_shape, name="ct_volume_input") | |
| # ── Encoder ────────────────────────────────────────────────────────────── | |
| encoder_outputs = [] | |
| x = inputs | |
| for i in range(depth): | |
| filters = base_filters * (2 ** i) | |
| x = conv_block_3d(x, filters) | |
| encoder_outputs.append(x) | |
| x = layers.MaxPooling3D(pool_size=(2, 2, 2))(x) | |
| # ── Bottleneck ──────────────────────────────────────────────────────────── | |
| bottleneck_filters = base_filters * (2 ** depth) | |
| x = conv_block_3d(x, bottleneck_filters) | |
| # ── Decoder ─────────────────────────────────────────────────────────────── | |
| for i in reversed(range(depth)): | |
| filters = base_filters * (2 ** i) | |
| x = layers.UpSampling3D(size=(2, 2, 2))(x) | |
| x = layers.Concatenate()([x, encoder_outputs[i]]) | |
| x = conv_block_3d(x, filters) | |
| # ── Output ──────────────────────────────────────────────────────────────── | |
| outputs = layers.Conv3D(1, kernel_size=1, activation="sigmoid", name="seg_mask_3d")(x) | |
| return keras.Model(inputs, outputs, name="UNet_3D") | |
| # =========================================================================== | |
| # NODULE / ANOMALY DETECTOR [NEW] | |
| # =========================================================================== | |
| def build_nodule_detector(input_shape=(32, 32, 1)): | |
| """ | |
| CNN Classifier for identifying lung nodules in localized 2D patches. | |
| This model is used by the sliding-window detector in predict.py. | |
| It takes a small 32x32 image patch (cropped from inside the lung mask) | |
| and predicts whether it contains a nodule. | |
| Architecture: Standard VGG-style Conv2D + Dense classification head. | |
| Args: | |
| input_shape : tuple, expected (32, 32, 1) or (64, 64, 1) | |
| Returns: | |
| model : keras.Model (Output = scalar probability [0, 1]) | |
| """ | |
| inputs = keras.Input(shape=input_shape, name="patch_input") | |
| x = layers.Conv2D(32, 3, padding="same", activation="relu")(inputs) | |
| x = layers.BatchNormalization()(x) | |
| x = layers.MaxPooling2D(2)(x) # 16x16 | |
| x = layers.Conv2D(64, 3, padding="same", activation="relu")(x) | |
| x = layers.BatchNormalization()(x) | |
| x = layers.MaxPooling2D(2)(x) # 8x8 | |
| x = layers.Conv2D(128, 3, padding="same", activation="relu")(x) | |
| x = layers.BatchNormalization()(x) | |
| x = layers.MaxPooling2D(2)(x) # 4x4 | |
| # Global Average Pooling prevents overfitting better than Flatten() | |
| x = layers.GlobalAveragePooling2D()(x) | |
| x = layers.Dense(128, activation="relu")(x) | |
| x = layers.Dropout(0.5)(x) | |
| # Single output neuron: 1 = Nodule, 0 = Healthy tissue/vessel | |
| outputs = layers.Dense(1, activation="sigmoid", name="nodule_prob")(x) | |
| return keras.Model(inputs, outputs, name="Nodule_Detector") | |