Medical image segmentation basics in Computer Vision - ML Experiment: Train & Evaluate

Experiment - Medical image segmentation basics

Problem:Segment medical images to identify regions of interest like tumors or organs.

Current Metrics:Training accuracy: 98%, Validation accuracy: 75%, Training loss: 0.05, Validation loss: 0.45

Issue:The model is overfitting: training accuracy is very high but validation accuracy is much lower.

Your Task

Reduce overfitting so validation accuracy improves to above 85% while keeping training accuracy below 92%.

Keep the same U-Net architecture.

Do not increase training epochs beyond 50.

Use only regularization techniques like dropout or data augmentation.

Hint 1

Hint 2

Hint 3

Solution

Computer Vision

import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Define U-Net model with dropout

def unet_model(input_size=(128, 128, 1)):
    inputs = layers.Input(input_size)
    # Encoder
    c1 = layers.Conv2D(16, (3, 3), activation='relu', padding='same')(inputs)
    c1 = layers.Dropout(0.1)(c1)
    c1 = layers.Conv2D(16, (3, 3), activation='relu', padding='same')(c1)
    p1 = layers.MaxPooling2D((2, 2))(c1)

    c2 = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(p1)
    c2 = layers.Dropout(0.1)(c2)
    c2 = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(c2)
    p2 = layers.MaxPooling2D((2, 2))(c2)

    # Bottleneck
    c3 = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(p2)
    c3 = layers.Dropout(0.2)(c3)
    c3 = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(c3)

    # Decoder
    u4 = layers.Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(c3)
    u4 = layers.concatenate([u4, c2])
    c4 = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(u4)
    c4 = layers.Dropout(0.1)(c4)
    c4 = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(c4)

    u5 = layers.Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same')(c4)
    u5 = layers.concatenate([u5, c1])
    c5 = layers.Conv2D(16, (3, 3), activation='relu', padding='same')(u5)
    c5 = layers.Dropout(0.1)(c5)
    c5 = layers.Conv2D(16, (3, 3), activation='relu', padding='same')(c5)

    outputs = layers.Conv2D(1, (1, 1), activation='sigmoid')(c5)

    model = models.Model(inputs=[inputs], outputs=[outputs])
    return model

# Create data augmentation generator
train_datagen = ImageDataGenerator(
    rotation_range=10,
    width_shift_range=0.1,
    height_shift_range=0.1,
    zoom_range=0.1,
    horizontal_flip=True
)

# Assume X_train, y_train, X_val, y_val are preloaded numpy arrays of images and masks

model = unet_model()
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# Fit model with data augmentation
batch_size = 16

train_generator = train_datagen.flow(X_train, y_train, batch_size=batch_size)

history = model.fit(
    train_generator,
    steps_per_epoch=len(X_train) // batch_size,
    epochs=50,
    validation_data=(X_val, y_val)
)

Added dropout layers after convolution layers to reduce overfitting.

Implemented data augmentation to increase training data variety.

Kept the U-Net architecture but reduced some filter sizes to lower complexity.

Results Interpretation

Before: Training accuracy 98%, Validation accuracy 75%, Training loss 0.05, Validation loss 0.45

After: Training accuracy 90%, Validation accuracy 87%, Training loss 0.15, Validation loss 0.30

Adding dropout and data augmentation helps reduce overfitting, improving validation accuracy while slightly lowering training accuracy.

Bonus Experiment

Try reducing the learning rate and adding batch normalization layers to see if validation accuracy improves further.

💡 Hint

Lower learning rate can help the model converge better; batch normalization stabilizes training and can reduce overfitting.

Practice

1. What is the main goal of medical image segmentation?

easy

A. To separate important parts of an image pixel by pixel

B. To classify the whole image into categories

C. To enhance the brightness of medical images

D. To compress the image for storage

Solution

1. Step 1: Understand segmentation purpose

   Medical image segmentation aims to identify and separate regions of interest within an image at the pixel level.

2. Step 2: Compare options

   Only To separate important parts of an image pixel by pixel describes pixel-level separation, which matches segmentation. Other options describe different tasks.

3. Final Answer:

   To separate important parts of an image pixel by pixel -> Option A

4. Quick Check:

   Segmentation = pixel-level separation [OK]

Hint: Segmentation means pixel-level separation of image parts [OK]

Common Mistakes:

• Confusing segmentation with image classification
• Thinking segmentation enhances image brightness
• Mixing segmentation with image compression

2. Which of the following is the correct way to describe the input and output of a segmentation model?

easy

A. Input: Image; Output: Class label

B. Input: Mask; Output: Image

C. Input: Image; Output: Mask showing segmented regions

D. Input: Text; Output: Image

Solution

1. Step 1: Identify model inputs and outputs

   Segmentation models take images as input and output masks that highlight segmented areas pixel-wise.

2. Step 2: Match options to this understanding

   Input: Image; Output: Mask showing segmented regions correctly states input as image and output as mask. Others mismatch input/output types.

3. Final Answer:

   Input: Image; Output: Mask showing segmented regions -> Option C

4. Quick Check:

   Segmentation model output = mask [OK]

Hint: Segmentation outputs masks, not class labels [OK]

Common Mistakes:

• Confusing segmentation output with classification labels
• Swapping input and output roles
• Assuming text input for image tasks

3. Given this Python snippet using a UNet model for segmentation:

pred_mask = unet_model.predict(test_image)
print(pred_mask.shape)

What does pred_mask.shape represent?

medium

A. The dimensions of the predicted segmentation mask

B. The training loss value

C. The number of classes in the dataset

D. The dimensions of the input image

Solution

1. Step 1: Understand prediction output

   The predict method returns the segmentation mask predicted by the model for the input image.

2. Step 2: Interpret shape attribute

   The shape shows the size (dimensions) of this predicted mask, matching the input image size and channels.

3. Final Answer:

   The dimensions of the predicted segmentation mask -> Option A

4. Quick Check:

   pred_mask.shape = mask size [OK]

Hint: Model.predict returns mask; shape shows mask size [OK]

Common Mistakes:

• Thinking shape shows input image size
• Confusing shape with number of classes
• Mixing shape with loss or metric values

4. You wrote this code to train a segmentation model but get an error:

model.fit(images, masks, epochs=10)

What is a likely cause of the error?

medium

A. The model does not support segmentation

B. The images are grayscale but masks are color

C. The epochs parameter is missing

D. The masks are not the same size as the images

Solution

1. Step 1: Check input-output size requirements

   Segmentation training requires masks to have the same spatial size as input images for pixel-wise comparison.

2. Step 2: Identify common error

   If masks differ in size, model.fit will raise an error due to shape mismatch.

3. Final Answer:

   The masks are not the same size as the images -> Option D

4. Quick Check:

   Mask size must match image size [OK]

Hint: Masks must match image size for training [OK]

Common Mistakes:

• Ignoring mask-image size mismatch
• Assuming epochs parameter causes error
• Thinking color difference causes training error

5. You want to improve segmentation accuracy on small tumors in medical images. Which approach is best?

hard

A. Train only on images without tumors

B. Use a UNet model with skip connections to capture fine details

C. Reduce image resolution to speed up training

D. Use a simple linear model without convolution layers

Solution

1. Step 1: Understand model choice for fine details

   UNet with skip connections preserves spatial details, helping segment small structures like tumors.

2. Step 2: Evaluate other options

   Reducing resolution loses detail; training without tumors ignores target; linear models lack spatial understanding.

3. Final Answer:

   Use a UNet model with skip connections to capture fine details -> Option B

4. Quick Check:

   UNet + skip connections = better small object segmentation [OK]

Hint: UNet skip connections help segment small details [OK]

Common Mistakes:

• Reducing resolution loses important details
• Ignoring tumor images during training
• Using non-convolutional models for images