Medical image segmentation basics in Computer Vision

Introduction

Medical image segmentation helps doctors see and understand important parts of images like tumors or organs clearly.

When doctors want to find the exact size and shape of a tumor in an MRI scan.

To separate different organs in a CT scan for better analysis.

When planning surgery and needing to know where healthy tissue ends and affected tissue begins.

To track how a disease changes over time by comparing segmented images.

For automatic detection of abnormalities in X-ray images.

Syntax

Computer Vision

model = UNet(input_shape=(height, width, channels), num_classes=number_of_classes)
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
model.fit(train_images, train_masks, epochs=epochs, batch_size=batch_size)

UNet is a popular model architecture for medical image segmentation.

The model learns to assign each pixel in the image to a class, like tumor or healthy tissue.

Examples

This sets up a UNet model for grayscale images with two classes: background and object.

Computer Vision

model = UNet(input_shape=(128, 128, 1), num_classes=2)
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

This trains the model for 10 rounds using 16 images at a time.

Computer Vision

model.fit(train_images, train_masks, epochs=10, batch_size=16)

This predicts segmentation masks and converts probabilities to class index masks.

Computer Vision

predictions = model.predict(test_images)
segmented_images = np.argmax(predictions, axis=-1)

Sample Model

This code builds a simple UNet-like model, trains it on dummy grayscale images with a square mask, and predicts the segmentation mask for one image.

Computer Vision

import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models

# Simple UNet-like model for demonstration
def simple_unet(input_shape, num_classes):
    inputs = layers.Input(shape=input_shape)
    c1 = layers.Conv2D(16, 3, activation='relu', padding='same')(inputs)
    p1 = layers.MaxPooling2D()(c1)
    c2 = layers.Conv2D(32, 3, activation='relu', padding='same')(p1)
    u1 = layers.UpSampling2D()(c2)
    concat1 = layers.Concatenate()([u1, c1])
    outputs = layers.Conv2D(num_classes, 1, activation='softmax')(concat1)
    return models.Model(inputs, outputs)

# Create dummy data: 10 images 64x64 grayscale, 2 classes
train_images = np.random.rand(10, 64, 64, 1).astype(np.float32)
train_masks = np.zeros((10, 64, 64, 2), dtype=np.float32)
# Create simple masks: class 1 in center square
train_masks[:, 24:40, 24:40, 1] = 1
train_masks[:, :, :, 0] = 1 - train_masks[:, :, :, 1]

model = simple_unet((64, 64, 1), 2)
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

history = model.fit(train_images, train_masks, epochs=3, batch_size=2, verbose=2)

# Predict on one image
pred = model.predict(train_images[:1])
pred_mask = np.argmax(pred, axis=-1)
print('Predicted mask shape:', pred_mask.shape)
print('Unique classes in predicted mask:', np.unique(pred_mask))

Important Notes

Medical image segmentation models often need many labeled images to learn well.

Segmentation output is usually a mask showing which pixels belong to which class.

Accuracy alone may not be enough; metrics like Dice score are common for segmentation.

Summary

Medical image segmentation separates important parts of images pixel by pixel.

UNet is a common model used for this task because it captures details well.

Training needs images and matching masks showing the correct segmentation.