Explanation of the Breast Tumor Segmentation Code

Introduction

This code builds a segmentation model to detect tumors in breast MRI images. The tumors can be:

• Benign (non-cancerous)
• Malignant (cancerous)
• Or Normal (no tumor)

We are training a U-Net model (with a ResNet34 backbone) to find tumors in MRI images and classify them into these categories by segmenting the exact tumor region in the image.

Step-by-Step Breakdown

Install Segmentation Model

pip install segmentation_models

1. Import Libraries

We load essential Python libraries:

• numpy, os: For handling files and arrays.
• cv2, scipy, matplotlib: For image processing and visualization.
• segmentation_models: For building the U-Net model.

import os

import cv2
import matplotlib.pyplot as plt
import numpy as np  # linear algebra
import scipy
import scipy.ndimage
import segmentation_models as sm
import tensorflow as tf

2. Load Dataset Information

We have a dataset containing three folders:

• Normal
• Benign
• Malignant

We read each folder name and assign an index:

normal = 0
benign = 1
malignant = 2

BASE_PATH = "/kaggle/input/breast-ultrasound-images-dataset/Dataset_BUSI_with_GT"
unique_classes = []
for path in os.listdir(BASE_PATH):
    unique_classes.append(path)
print(unique_classes)

# assign index to classes
class_index = [unique_classes[1], unique_classes[0], unique_classes[2]]
for c in class_index:
    print(c, "-", class_index.index(c))

3. Prepare Paths for Images and Masks

For each image:

• Find its corresponding mask (where the tumor area is marked).
• Save the image path, mask path, and its label (normal, benign, or malignant).

Why?
✅ So we can later load the actual image and its tumor mask for training.

images = []
masks = []
labels = []
for folder in os.listdir(BASE_PATH):
    class_path = os.path.join(BASE_PATH, folder)
    for img in os.listdir(class_path):
        if "_mask" not in img:
            img_path = os.path.join(class_path, img)
            msk_path = img_path.replace(".png", "_mask.png")
            # check if mask exist
            if os.path.exists(msk_path):
                images.append(img_path)
                masks.append(msk_path)
                labels.append(folder)

4. Image Preprocessing Functions

• load_image() – Loads an image in grayscale.
• padding() – Pads images to make them square. Why? Models prefer fixed input sizes.
• resize() and resize_mask() – Resizes images and masks to 256x256 pixels.
• preprocess() – Normalizes pixel values to the range 0-1.

Why?
✅ So that all images are the same size and scale, which the neural network requires.

def load_image(img_path):
    """Load single image as Grayscale"""
    # load image as grayscale
    img = cv2.imread(img_path, 0)
    return img


def padding(img, msk):
    """Pad images to make them square"""
    size = np.max(img.shape)

    offset_x = (size - img.shape[0]) // 2
    offset_y = (size - img.shape[1]) // 2

    blank_image = np.zeros((size, size))
    blank_mask = np.zeros((size, size))

    blank_image[
        offset_x : offset_x + img.shape[0], offset_y : offset_y + img.shape[1]
    ] = img
    blank_mask[
        offset_x : offset_x + img.shape[0], offset_y : offset_y + img.shape[1]
    ] = msk
    return blank_image, blank_mask


def resize_mask(mask):
    """Resize mask, its different because mask pixel value can change because of resize"""
    new_size = np.array([input_images_size, input_images_size]) / mask.shape
    mask = scipy.ndimage.interpolation.zoom(mask, new_size)
    return mask


def resize(img):
    """Resize image"""
    img = cv2.resize(img, (input_images_size, input_images_size))
    return img


def preprocess(img):
    """Image preprocessing
    Normalize image
    """
    img = img / 255.0
    return img


def inverse_preprocess(img):
    """Inverse of preprocessing"""
    img = img * 255
    return img

5. Merging Mask with Tumor Type

Each mask pixel is either:

• 0 (background)
• 1 (tumor area)

But we also care if the tumor is benign or malignant. So we create a 2-channel mask:

• Channel 1: Benign tumor mask
• Channel 2: Malignant tumor mask

If the tumor is normal, the mask is left empty.

def load_data(img_path, msk_path, label):
    """Load image, mask and repalce mask value with class index
    0 = normal
    1 = benign
    2 = malignant
    """
    img = load_image(img_path)
    msk = load_image(msk_path)
    img, msk = padding(img, msk)
    label_indx = class_index.index(label)
    msk[msk == 255] = 1
    msk = msk.astype("uint8")
    img = resize(img)
    msk = resize_mask(msk)
    new_mask = np.zeros((input_images_size, input_images_size, 2))
    if label_indx != 0:
        new_mask[:, :, label_indx - 1] = msk
    #     print(np.unique(msk), label, label_indx)
    return img, new_mask

🔹 6. Batch Loading

load_batch() loads a small group of images and their masks, preparing them for training.

Why use batches?
✅ Because training on the entire dataset at once would be too big for memory.

def load_batch(images, masks, labels):
    """Load Batch of data"""
    batch_x = []
    batch_y = []
    for i, m, l in zip(images, masks, labels):
        img, msk = load_data(i, m, l)
        img = preprocess(img)
        batch_x.append(img)
        batch_y.append(msk)
    return np.array(batch_x), np.array(batch_y)

🔹 7. Visualizing Sample Data

We visualize some examples to check if:

• The images and masks are loaded correctly.
• The tumors are properly highlighted.

for i in [0, 500, 600]:
    indx = i
    img, msk = load_data(images[indx], masks[indx], labels[indx])
    print(np.min(img), np.max(img), img.shape)
    plt.figure(figsize=(20, 20))
    plt.subplot(1, 3, 1)
    plt.axis("off")
    plt.imshow(img)
    plt.subplot(1, 3, 2)
    plt.axis("off")
    plt.imshow(msk[:, :, 0])
    plt.subplot(1, 3, 3)
    plt.axis("off")
    plt.imshow(msk[:, :, 1])
    plt.show()

8. Build the Segmentation Model

We use the popular U-Net architecture with a ResNet34 encoder.
Why U-Net?
✅ U-Net is excellent for medical image segmentation due to its design of downsampling and upsampling paths.

We define:

• Loss function: Combination of Dice Loss (good for overlapping areas) and Focal Loss (helps focus on difficult examples).
• Optimizer: Adam.
• Metrics: IOU (Intersection over Union) and F-Score.

sm.set_framework("tf.keras")

sm.framework()

BACKBONE = "resnet34"
LR = 0.00001
model = sm.Unet(
    BACKBONE,
    classes=2,
    activation="sigmoid",
    input_shape=(input_images_size, input_images_size, channel),
    encoder_weights=None,
)

optim = tf.keras.optimizers.Adam(LR)

dice_loss = sm.losses.DiceLoss()
focal_loss = sm.losses.BinaryFocalLoss()
total_loss = dice_loss + (1 * focal_loss)

metrics = [sm.metrics.IOUScore(threshold=0.5), sm.metrics.FScore(threshold=0.5)]

model.compile(optim, total_loss, metrics)

9. Training the Model

We train the model for 100 epochs, and for each batch:

• Load batch of images and masks.
• Train on that batch.
• Record the loss, IOU, and F-Score.

Why so many epochs?
✅ Because medical images need careful, long training to achieve high accuracy.

batch_size = 4
history = {"epoch": []}
for e in range(100):
    print("epoch:", e, end=" > ")
    indexes = list(range(len(images)))
    temp_history = {"loss": [], "IOU": [], "F-Score": []}
    for b in range(0, len(images), batch_size):
        bs = b
        be = bs + batch_size
        batch_index = indexes[bs:be]
        batch_x, batch_y = load_batch(
            images[batch_index], masks[batch_index], labels[batch_index]
        )
        batch_x = np.expand_dims(batch_x, axis=-1)
        batch_y = np.expand_dims(batch_y, axis=-1)
        batch_y = batch_y.astype("float32")
        loss = model.train_on_batch(batch_x, batch_y)
        temp_history["loss"].append(loss[0])
        temp_history["IOU"].append(loss[1])
        temp_history["F-Score"].append(loss[2])
    print(
        "loss",
        np.round(np.mean(temp_history["loss"]), 4),
        "IOU",
        np.round(np.mean(temp_history["IOU"]), 4),
        "F-Score",
        np.round(np.mean(temp_history["F-Score"]), 4),
    )
    history["epoch"].append(temp_history)

🔹 10. Saving the Model

After training, we save the weights of the model to use later without retraining.

model.save_weights("breast_tumor_segmentation")

🔹 11. Testing the Model

We load a few test images, run them through the model, and visualize:

• The original image.
• Predicted benign tumor area.
• Predicted malignant tumor area.

This helps us visually confirm if the model is correctly segmenting the tumor areas.

for i in [0, 500, 600]:
    indx = i
    img, msk = load_data(images[indx], masks[indx], labels[indx])
    print(np.min(img), np.max(img), img.shape)
    print(img.shape)

    img2 = preprocess(img)
    pred = model.predict(np.array([img2]))
    pred = pred[0]

    plt.figure(figsize=(20, 20))
    plt.subplot(1, 3, 1)
    plt.axis("off")
    plt.imshow(img)
    plt.subplot(1, 3, 2)
    plt.axis("off")
    plt.imshow(pred[:, :, 0])
    plt.subplot(1, 3, 3)
    plt.axis("off")
    plt.imshow(pred[:, :, 1])
    plt.show()

What’s the logic of this code?

Step Reason Merge masks To handle multiple tumor types in one mask. Preprocess images To make them uniform for the model. Use U-Net It’s ideal for precise medical image segmentation. Combine losses To handle class imbalance and improve boundary segmentation. Train in batches To fit large datasets in memory. Visualize results To ensure the model works as expected.

Final Thoughts

This code is a full pipeline to:

• Prepare data.
• Train a deep learning model.
• Predict and visualize tumor segmentation.

Reference

Full Code: https://www.kaggle.com/code/faheemkhaskheli9/breast-tumor-segmentation