Code Overview

This script performs segmentation inference using the MedSAM model. It processes 2D medical images in .npz format, applies a trained segmentation model, and saves the resulting masks along with optional overlay visualizations.

Imports and Configuration

The script begins with necessary imports and configuration settings:

import torch
import numpy as np
import cv2
from segment_anything import sam_model_registry
from segment_anything.modeling import MaskDecoder, PromptEncoder, TwoWayTransformer
from tiny_vit_sam import TinyViT
import argparse
from tqdm import tqdm
from time import time
import os
import matplotlib.pyplot as plt

Argument Parsing

Command-line arguments are defined using argparse:

parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input_dir', type=str, default='./data/Viral Pneumonia/dd_dd/npz/')
parser.add_argument('-o', '--output_dir', type=str, default='./data/Viral Pneumonia/dd_dd/new_sam_segs/')
parser.add_argument("-chk", "--checkpoint", type=str, default="./train/work_dir/medsam_model_best.pth")
parser.add_argument('-lite_medsam_checkpoint_path', type=str, default="work_dir/LiteMedSAM/lite_medsam.pth")
parser.add_argument('-device', type=str, default="cpu")
parser.add_argument('--save_overlay', action='store_true', default=True)
parser.add_argument('-png_save_dir', type=str, default='./overlay/new_sam_overlay')
args = parser.parse_args()

Image Porcessing Functions

Two utility functions handle image resizing and padding:

def resize_longest_side(image, target_length=256):
    oldh, oldw = image.shape[:2]
    scale = target_length / max(oldh, oldw)
    new_size = (int(oldw * scale + 0.5), int(oldh * scale + 0.5))
    return cv2.resize(image, new_size, interpolation=cv2.INTER_AREA)

def pad_image(image, target_size=256):
    h, w = image.shape[:2]
    padh = target_size - h
    padw = target_size - w
    if len(image.shape) == 3:
        return np.pad(image, ((0, padh), (0, padw), (0, 0)))
    return np.pad(image, ((0, padh), (0, padw)))

Model Definition

The MedSAM Lite model combines an image encoder, mask decoder, and prompt encoder:

class MedSAM_Lite(nn.Module):
    def __init__(self, image_encoder, mask_decoder, prompt_encoder):
        super().__init__()
        self.image_encoder = image_encoder
        self.mask_decoder = mask_decoder
        self.prompt_encoder = prompt_encoder

    def forward(self, image, box_np):
        image_embedding = self.image_encoder(image)
        with torch.no_grad():
            box_torch = torch.as_tensor(box_np, dtype=torch.float32, device=image.device)
            if len(box_torch.shape) == 2:
                box_torch = box_torch[:, None, :]
        sparse_embeddings, dense_embeddings = self.prompt_encoder(
            points=None, boxes=box_np, masks=None
        )
        low_res_masks, _ = self.mask_decoder(
            image_embeddings=image_embedding,
            image_pe=self.prompt_encoder.get_dense_pe(),
            sparse_prompt_embeddings=sparse_embeddings,
            dense_prompt_embeddings=dense_embeddings,
            multimask_output=False
        )
        return low_res_masks

Postprocessing

The postprocessing function resizes masks to match original image dimensions:

@torch.no_grad()
def postprocess_masks(self, masks, new_size, original_size):
    masks = masks[..., :new_size[0], :new_size[1]]
    return F.interpolate(masks, size=original_size, mode="bilinear", align_corners=False)

Visualization Functions

Utility functions for mask and bounding box visualization:

def show_mask(mask, ax, mask_color=None, alpha=0.5):
    if mask_color is not None:
        color = np.concatenate([mask_color, [alpha]])
    else:
        color = [251/255, 252/255, 30/255, alpha]
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)

def show_box(box, ax, edgecolor='blue'):
    x0, y0 = box[:2]
    w, h = box[2] - x0, box[3] - y0
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor=edgecolor, facecolor=(0,0,0,0), lw=2))

Bounding Box Calculation

Generates bounding boxes from ground truth masks:

def get_bbox256(mask_256, bbox_shift=3):
    y_indices, x_indices = np.where(mask_256 > 0)
    x_min, x_max = min(x_indices), max(x_indices)
    y_min, y_max = min(y_indices), max(y_indices)
    H, W = mask_256.shape
    x_min = max(0, x_min - bbox_shift)
    x_max = min(W, x_max + bbox_shift)
    y_min = max(0, y_min - bbox_shift)
    y_max = min(H, y_max + bbox_shift)
    return np.array([x_min, y_min, x_max, y_max])

Resizing Bounding Boxes

Adjusts bounding box coordinates to match target image size:

def resize_box_to_256(box, original_size):
    new_box = np.zeros_like(box)
    ratio = 256 / max(original_size)
    for i in range(len(box)):
        new_box[i] = int(box[i] * ratio)
    return new_box

Inference Function

Performs segmentation inference using MedSAM model:

@torch.no_grad()
def medsam_inference(medsam_model, img_embed, box_1024, H, W):
    box_torch = torch.as_tensor(box_1024, dtype=torch.float, device=img_embed.device)
    if len(box_torch.shape) == 2:
        box_torch = box_torch[:, None, :]
    sparse_embeddings, dense_embeddings = medsam_model.prompt_encoder(
        points=None, boxes=box_torch, masks=None
    )
    low_res_logits, _ = medsam_model.mask_decoder(
        image_embeddings=img_embed,
        image_pe=medsam_model.prompt_encoder.get_dense_pe(),
        sparse_prompt_embeddings=sparse_embeddings,
        dense_prompt_embeddings=dense_embeddings,
        multimask_output=False
    )
    low_res_pred = torch.sigmoid(low_res_logits)
    low_res_pred = F.interpolate(low_res_pred, size=(H, W), mode="bilinear", align_corners=False)
    return (low_res_pred.squeeze().cpu().numpy() > 0.5).astype(np.uint8)

Main Inference Pipeline

The core function processes a single .npz file:

def MedSAM_infer_npz_2D(img_npz_file):
    npz_name = os.path.basename(img_npz_file)
    npz_data = np.load(img_npz_file, allow_pickle=True)
    img_3c = npz_data['imgs']
    H, W = 256, 256
    img_256 = cv2.resize(img_3c, (H, W), interpolation=cv2.INTER_NEAREST)
    img_1024 = cv2.resize(img_3c, (1024, 1024), interpolation=cv2.INTER_NEAREST)
    
    # Process ground truth
    gt = npz_data['gts']
    if gt.shape != (H, W):
        gt = cv2.resize(gt.astype(np.uint8), (W, H), interpolation=cv2.INTER_NEAREST)
    gt = pad_image(gt)
    label_ids = np.unique(gt)[1:]
    
    # Generate random label and bounding box
    import random
    gt2D = np.uint8(gt == random.choice(label_ids.tolist()))
    gt2D = np.uint8(gt2D > 0)
    y_indices, x_indices = np.where(gt2D > 0)
    x_min, x_max = min(x_indices), max(x_indices)
    y_min, y_max = min(y_indices), max(y_indices)
    
    # Add random shift to bounding box
    bbox_shift = 5
    x_min = max(0, x_min - random.randint(0, bbox_shift))
    x_max = min(W, x_max + random.randint(0, bbox_shift))
    y_min = max(0, y_min - random.randint(0, bbox_shift))
    y_max = min(H, y_max + random.randint(0, bbox_shift))
    boxes = np.array([[x_min, y_min, x_max, y_max]])
    
    # Process image for inference
    img_1024 = resize_longest_side(img_1024, 1024)
    newh, neww = img_1024.shape[:2]
    img_1024_norm = (img_1024 - img_1024.min()) / max(1e-8, img_1024.max() - img_1024.min())
    img_1024_padded = pad_image(img_1024_norm, 1024)
    img_1024_tensor = torch.tensor(img_1024_padded).float().permute(2, 0, 1).unsqueeze(0).to(device)
    
    # Perform inference
    with torch.no_grad():
        image_embedding = medsam_model.image_encoder(img_1024_tensor)
    
    # Process each bounding box
    segs = np.zeros((H, W), dtype=np.uint8)
    for idx, box in enumerate(boxes, start=1):
        box_1024 = box / np.array([W, H, W, H]) * 1024
        box_1024 = box_1024[None, ...]
        medsam_mask, _ = medsam_inference(medsam_model, image_embedding, box_1024, H, W)
        segs[medsam_mask > 0] = idx
    
    # Save results
    np.savez_compressed(os.path.join(pred_save_dir, npz_name), segs=segs)
    
    # Visualization
    if save_overlay:
        fig, ax = plt.subplots(1, 2, figsize=(10, 5))
        ax[0].imshow(img_256)
        ax[1].imshow(img_256)
        ax[0].set_title("Image")
        ax[1].set_title("SAM Segmentation")
        ax[0].axis('off')
        ax[1].axis('off')
        for i, box in enumerate(boxes):
            color = np.random.rand(3)
            show_box(box, ax[1], edgecolor=color)
            show_mask((segs == i + 1).astype(np.uint8), ax[1], mask_color=color)
        plt.tight_layout()
        plt.savefig(os.path.join(png_save_dir, os.path.splitext(npz_name)[0] + '.png'), dpi=300)
        plt.close()

Execution

The main block processes all .npz files in the input directory:

if __name__ == '__main__':
    img_npz_files = sorted(glob(os.path.join(data_root, '*.npz'), recursive=True))
    efficiency = {'case': [], 'time': []}
    
    for img_npz_file in tqdm(img_npz_files[:-1]):
        start_time = time()
        if os.path.basename(img_npz_file).startswith('3D'):
            MedSAM_infer_npz_3D(img_npz_file)
        else:
            MedSAM_infer_npz_2D(img_npz_file)
        end_time = time()
        
        efficiency['case'].append(os.path.basename(img_npz_file))
        efficiency['time'].append(end_time - start_time)
        print(f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')} - {os.path.basename(img_npz_file)}: {end_time - start_time:.4f}s")
    
    pd.DataFrame(efficiency).to_csv(os.path.join(pred_save_dir, 'efficiency.csv'), index=False)