05 - Image Segmentation: Semantic, Instance, Panoptic

import torch
import numpy as np
from typing import Union

def compute_iou(pred: torch.Tensor, target: torch.Tensor, threshold: float = 0.5) -> float:
    """
    Intersection over Union for binary segmentation.
    pred, target: tensors [H, W] or [B, H, W] with values in [0, 1]
    Range: 0 (no overlap) -> 1 (perfect overlap)
    """
    pred_binary = (pred >= threshold).bool()
    target_binary = target.bool()

    intersection = (pred_binary & target_binary).float().sum()
    union = (pred_binary | target_binary).float().sum()

    if union == 0:
        return 1.0  # degenerate case: both masks are empty

    return float(intersection / union)

def dice_score(pred: torch.Tensor, target: torch.Tensor,
               threshold: float = 0.5, smooth: float = 1.0) -> float:
    """
    Dice Score (F1 for segmentation): 2*|X intersect Y| / (|X| + |Y|)
    Preferred in medical imaging because less sensitive to class imbalance.
    A model predicting all background gets Dice=0, whereas Accuracy can be 99%+.
    Range: 0 (worst) -> 1 (perfect)
    """
    pred_binary = (pred >= threshold).float()
    target_binary = target.float()

    intersection = (pred_binary * target_binary).sum()
    dice = (2.0 * intersection + smooth) / (pred_binary.sum() + target_binary.sum() + smooth)

    return float(dice)

def compute_multiclass_miou(pred_logits: torch.Tensor, targets: torch.Tensor,
                             num_classes: int, ignore_index: int = 255) -> float:
    """
    Mean IoU for multi-class semantic segmentation.
    pred_logits: [B, C, H, W] - raw logits
    targets: [B, H, W] - class indices 0..num_classes-1
    """
    preds = pred_logits.argmax(dim=1)  # [B, H, W]
    ious = []

    for cls in range(num_classes):
        pred_cls = preds == cls
        true_cls = targets == cls
        valid = targets != ignore_index

        pred_cls = pred_cls & valid
        true_cls = true_cls & valid

        intersection = (pred_cls & true_cls).sum().float()
        union = (pred_cls | true_cls).sum().float()

        if union > 0:
            ious.append(float(intersection / union))

    return float(np.mean(ious)) if ious else 0.0

def hausdorff_distance(pred: np.ndarray, target: np.ndarray) -> float:
    """
    Hausdorff Distance: measures the maximum distance between mask boundaries.
    Useful in medicine for assessing contour precision - a model can have
    high Dice but large Hausdorff if a small boundary region is misaligned.
    Lower is better. Unit: pixels.
    """
    from scipy.spatial.distance import directed_hausdorff

    pred_points = np.argwhere(pred)
    target_points = np.argwhere(target)

    if len(pred_points) == 0 or len(target_points) == 0:
        return float('inf')

    d1 = directed_hausdorff(pred_points, target_points)[0]
    d2 = directed_hausdorff(target_points, pred_points)[0]

    return max(d1, d2)  # symmetric Hausdorff

# Example
pred_tensor = torch.sigmoid(torch.randn(256, 256))
target_tensor = (torch.randn(256, 256) > 0).float()

iou = compute_iou(pred_tensor, target_tensor)
dice = dice_score(pred_tensor, target_tensor)
print(f"IoU: {iou:.3f} | Dice: {dice:.3f}")

# Hausdorff (uses numpy arrays)
pred_np = (pred_tensor >= 0.5).numpy()
target_np = target_tensor.numpy().astype(bool)
hd = hausdorff_distance(pred_np, target_np)
print(f"Hausdorff Distance: {hd:.1f} px")

import torch
import torch.nn as nn
import torch.nn.functional as F

class DoubleConv(nn.Module):
    """U-Net basic block: Conv-BN-ReLU-Conv-BN-ReLU."""

    def __init__(self, in_channels: int, out_channels: int, mid_channels: int | None = None):
        super().__init__()
        if mid_channels is None:
            mid_channels = out_channels
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.double_conv(x)

class DownBlock(nn.Module):
    """Encoder block: MaxPool2d (stride 2) followed by DoubleConv."""

    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.maxpool_conv(x)

class UpBlock(nn.Module):
    """Decoder block: Upsample + skip connection concatenation + DoubleConv."""

    def __init__(self, in_channels: int, out_channels: int, bilinear: bool = True):
        super().__init__()
        if bilinear:
            # Bilinear: lighter, no learned parameters for upsampling
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
        else:
            # Transposed conv: learned upsampling, slightly more parameters
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2,
                                          kernel_size=2, stride=2)
            self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        x1 = self.up(x1)

        # Handle size mismatch (input may not be power of 2)
        diff_h = x2.size(2) - x1.size(2)
        diff_w = x2.size(3) - x1.size(3)
        x1 = F.pad(x1, [diff_w // 2, diff_w - diff_w // 2,
                         diff_h // 2, diff_h - diff_h // 2])

        # Skip connection: concatenate encoder + decoder features
        return self.conv(torch.cat([x2, x1], dim=1))

class UNet(nn.Module):
    """
    U-Net for binary or multi-class segmentation.

    Architecture (channels at each resolution level):
    Input -> [64] -> [128] -> [256] -> [512] -> [1024] (bottleneck)
                                               -> [512] -> [256] -> [128] -> [64] -> Output

    Skip connections: each encoder level feeds directly to the symmetric decoder level.
    """

    def __init__(self, in_channels: int = 1, num_classes: int = 1, bilinear: bool = True):
        super().__init__()
        self.inc   = DoubleConv(in_channels, 64)
        self.down1 = DownBlock(64, 128)
        self.down2 = DownBlock(128, 256)
        self.down3 = DownBlock(256, 512)
        factor = 2 if bilinear else 1
        self.down4 = DownBlock(512, 1024 // factor)

        self.up1 = UpBlock(1024, 512 // factor, bilinear)
        self.up2 = UpBlock(512,  256 // factor, bilinear)
        self.up3 = UpBlock(256,  128 // factor, bilinear)
        self.up4 = UpBlock(128,  64,            bilinear)
        self.out = nn.Conv2d(64, num_classes, kernel_size=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Encoder path (save skip connections)
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)

        # Decoder path (use skip connections)
        x = self.up1(x5, x4)
        x = self.up2(x,  x3)
        x = self.up3(x,  x2)
        x = self.up4(x,  x1)
        return self.out(x)

# Architecture test
model = UNet(in_channels=3, num_classes=1)
x = torch.randn(2, 3, 256, 256)
y = model(x)
print(f"Input: {x.shape} -> Output: {y.shape}")
# Input: torch.Size([2, 3, 256, 256]) -> Output: torch.Size([2, 1, 256, 256])

total_params = sum(p.numel() for p in model.parameters())
print(f"Parameters: {total_params:,}")

import torch
import torch.nn as nn

class DiceLoss(nn.Module):
    """
    Dice Loss for binary segmentation.
    Handles class imbalance naturally - ideal for medical images
    (e.g., 95% background, 5% lesion/tumor).
    Loss = 1 - Dice coefficient.
    """

    def __init__(self, smooth: float = 1.0):
        super().__init__()
        self.smooth = smooth

    def forward(self, pred_logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        preds = torch.sigmoid(pred_logits)
        preds_flat = preds.view(-1)
        targets_flat = targets.view(-1)

        intersection = (preds_flat * targets_flat).sum()
        dice = (2.0 * intersection + self.smooth) / (
            preds_flat.sum() + targets_flat.sum() + self.smooth
        )
        return 1.0 - dice

class CombinedLoss(nn.Module):
    """
    BCE + Dice combined: the best tradeoff for medical segmentation.
    BCE: optimizes each pixel individually - provides dense gradients
    Dice: optimizes global overlap between prediction and ground truth
    The combination avoids the instability of Dice alone with small batches.
    """

    def __init__(self, bce_weight: float = 0.5, dice_weight: float = 0.5):
        super().__init__()
        self.bce = nn.BCEWithLogitsLoss()
        self.dice = DiceLoss()
        self.bce_weight = bce_weight
        self.dice_weight = dice_weight

    def forward(self, pred_logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        bce_loss  = self.bce(pred_logits, targets.float())
        dice_loss = self.dice(pred_logits, targets.float())
        return self.bce_weight * bce_loss + self.dice_weight * dice_loss

def train_unet(
    model: UNet,
    train_loader,
    val_loader,
    num_epochs: int = 50,
    learning_rate: float = 1e-4
) -> dict:
    """
    Complete U-Net training with:
    - Combined BCE+Dice loss
    - AdamW + CosineAnnealingLR scheduler
    - Early stopping on validation Dice score
    - Best model checkpoint
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)

    criterion = CombinedLoss(bce_weight=0.5, dice_weight=0.5)
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=learning_rate, weight_decay=1e-5
    )
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=num_epochs, eta_min=1e-6
    )

    history = {'train_loss': [], 'val_loss': [], 'val_dice': []}
    best_dice = 0.0
    patience  = 15
    no_improve = 0

    for epoch in range(num_epochs):
        # Training phase
        model.train()
        train_loss = 0.0
        for images, masks in train_loader:
            images, masks = images.to(device), masks.to(device)

            pred_logits = model(images)
            loss = criterion(pred_logits, masks)

            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            train_loss += loss.item()

        scheduler.step()

        # Validation phase
        model.eval()
        val_loss = 0.0
        val_dice_scores = []

        with torch.no_grad():
            for images, masks in val_loader:
                images, masks = images.to(device), masks.to(device)
                pred_logits = model(images)
                val_loss += criterion(pred_logits, masks).item()

                preds = torch.sigmoid(pred_logits)
                for p, m in zip(preds, masks):
                    val_dice_scores.append(dice_score(p, m))

        avg_train_loss = train_loss / len(train_loader)
        avg_val_loss   = val_loss   / len(val_loader)
        avg_val_dice   = sum(val_dice_scores) / len(val_dice_scores)

        history['train_loss'].append(avg_train_loss)
        history['val_loss'].append(avg_val_loss)
        history['val_dice'].append(avg_val_dice)

        if avg_val_dice > best_dice:
            best_dice = avg_val_dice
            torch.save(model.state_dict(), 'best_unet.pth')
            no_improve = 0
        else:
            no_improve += 1

        print(f"Epoch {epoch+1:2d}/{num_epochs} | "
              f"Loss: {avg_train_loss:.4f}/{avg_val_loss:.4f} | "
              f"Dice: {avg_val_dice:.4f} | Best: {best_dice:.4f}")

        if no_improve >= patience:
            print(f"Early stopping at epoch {epoch+1}")
            break

    print(f"Training complete. Best Dice Score: {best_dice:.4f}")
    return history

# pip install segment-anything
# Download checkpoint: https://github.com/facebookresearch/segment-anything

import numpy as np
import cv2
from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator

def load_sam_model(
    model_type: str = 'vit_h',
    checkpoint_path: str = 'sam_vit_h_4b8939.pth',
    device: str = 'cuda'
):
    """
    Load SAM model.
    Types: 'vit_h' (632M params, best quality), 'vit_l' (308M), 'vit_b' (91M, fastest)
    """
    sam = sam_model_registry[model_type](checkpoint=checkpoint_path)
    sam.to(device=device)
    return sam

def segment_with_point_prompt(
    sam_model,
    image: np.ndarray,
    point_coords: list[tuple[int, int]],
    point_labels: list[int]   # 1=foreground, 0=background
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Segment using point prompts.
    Returns 3 candidate masks sorted by confidence score (best first).
    """
    predictor = SamPredictor(sam_model)
    predictor.set_image(image)  # runs ViT encoder once

    masks, scores, logits = predictor.predict(
        point_coords=np.array(point_coords),
        point_labels=np.array(point_labels),
        multimask_output=True   # generate 3 candidate masks
    )

    sorted_idx = np.argsort(scores)[::-1]
    return masks[sorted_idx], scores[sorted_idx], logits[sorted_idx]

def segment_with_box_prompt(
    sam_model,
    image: np.ndarray,
    box: tuple[int, int, int, int]  # [x1, y1, x2, y2]
) -> tuple[np.ndarray, float]:
    """
    Segment using bounding box prompt. Box-prompted SAM is typically
    more precise than point-prompted for well-defined objects.
    """
    predictor = SamPredictor(sam_model)
    predictor.set_image(image)

    masks, scores, _ = predictor.predict(
        box=np.array([box]),
        multimask_output=False  # single mask with box prompt
    )
    return masks[0], float(scores[0])

def automatic_segmentation(sam_model, image: np.ndarray) -> list[dict]:
    """
    Automatic segmentation: SAM segments EVERYTHING in the image
    without any prompt, using a grid of seed points.
    Returns list of dicts with 'segmentation', 'area', 'predicted_iou'.
    """
    mask_generator = SamAutomaticMaskGenerator(
        model=sam_model,
        points_per_side=32,           # 32x32 grid = 1024 seed points
        pred_iou_thresh=0.88,          # filter low-quality masks
        stability_score_thresh=0.95,   # filter unstable masks
        crop_n_layers=1,               # multi-crop for small objects
        crop_n_points_downscale_factor=2,
        min_mask_region_area=100       # remove very small regions
    )

    masks = mask_generator.generate(image)
    masks = sorted(masks, key=lambda x: x['area'], reverse=True)

    print(f"SAM found {len(masks)} segments")
    for i, mask in enumerate(masks[:5]):
        print(f"  Segment {i+1}: area={mask['area']} "
              f"iou_score={mask['predicted_iou']:.3f}")

    return masks

def visualize_sam_results(image: np.ndarray, masks: list[dict],
                           alpha: float = 0.4) -> np.ndarray:
    """Visualize all SAM masks with random colors and contours."""
    result = image.copy()
    np.random.seed(42)

    for mask_info in masks:
        mask = mask_info['segmentation']  # bool array [H, W]
        color = np.random.randint(50, 255, 3)

        overlay = result.copy()
        overlay[mask] = color
        result = cv2.addWeighted(result, 1 - alpha, overlay, alpha, 0)

        # Draw contour for crisp boundary visualization
        contours, _ = cv2.findContours(
            mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        cv2.drawContours(result, contours, -1, color.tolist(), 2)

    return result

# Usage
sam = load_sam_model('vit_b', 'sam_vit_b_01ec64.pth')  # lightweight variant

image = cv2.imread('image.jpg')
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# Single click on object center
masks, scores, _ = segment_with_point_prompt(
    sam, image_rgb,
    point_coords=[(320, 240)],   # click at object center
    point_labels=[1]              # 1 = foreground
)
print(f"Best mask confidence: {scores[0]:.3f}")

# Box prompt - usually more precise
mask, score = segment_with_box_prompt(sam, image_rgb, box=(100, 50, 400, 350))
print(f"Box segmentation score: {score:.3f}")

# pip install sam2
# SAM2 released by Meta AI, August 2024

import torch
import numpy as np
from sam2.build_sam import build_sam2_video_predictor

def segment_video_with_sam2(
    video_path: str,
    initial_frame: int,
    initial_points: list[tuple[int, int]],
    checkpoint: str = 'sam2_hiera_large.pt',
    config: str = 'sam2_hiera_l.yaml'
) -> dict[int, np.ndarray]:
    """
    Segment and track an object across all video frames.
    Initialize with point prompts on the first frame, then SAM2
    automatically propagates the mask through the entire video.

    Args:
        video_path: path to video file
        initial_frame: frame index where the object is annotated
        initial_points: list of (x, y) foreground point coordinates
        checkpoint: SAM2 model weights
        config: SAM2 model configuration

    Returns:
        Dict mapping frame_index -> binary mask [H, W]
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    predictor = build_sam2_video_predictor(config, checkpoint, device=device)

    with torch.inference_mode(), torch.autocast('cuda', dtype=torch.bfloat16):
        # Initialize state on the video
        state = predictor.init_state(video_path=video_path)
        predictor.reset_state(state)

        # Add foreground point prompts on the initial frame
        frame_idx, obj_ids, masks = predictor.add_new_points_or_box(
            inference_state=state,
            frame_idx=initial_frame,
            obj_id=1,
            points=np.array(initial_points),
            labels=np.ones(len(initial_points), dtype=np.int32)
        )

        # Propagate segmentation across the entire video
        video_masks = {}
        for frame_idx, obj_ids, masks in predictor.propagate_in_video(state):
            # Convert soft mask to binary
            mask = (masks[0][0] > 0.0).cpu().numpy()
            video_masks[frame_idx] = mask

    print(f"Segmentation complete: {len(video_masks)} frames processed")
    return video_masks

def apply_masks_to_video(video_path: str, masks: dict[int, np.ndarray],
                          output_path: str, color: tuple = (0, 255, 0),
                          alpha: float = 0.4) -> None:
    """Overlay tracked masks on the original video and save."""
    import cv2

    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    w   = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h   = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))

    frame_idx = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        if frame_idx in masks:
            mask = masks[frame_idx].astype(np.uint8)
            overlay = frame.copy()
            overlay[mask.astype(bool)] = color  # green overlay
            frame = cv2.addWeighted(frame, 1 - alpha, overlay, alpha, 0)

            # Draw contour
            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            cv2.drawContours(frame, contours, -1, color, 2)

        writer.write(frame)
        frame_idx += 1

    cap.release()
    writer.release()
    print(f"Video with masks saved to {output_path}")

# Usage
masks = segment_video_with_sam2(
    video_path='surgery.mp4',
    initial_frame=0,
    initial_points=[(320, 240), (350, 260)]  # click on surgical instrument
)
apply_masks_to_video('surgery.mp4', masks, 'surgery_tracked.mp4')

class SAMLabelingAssistant:
    """
    Interactive labeling tool using SAM for automatic mask generation.
    Workflow: Human clicks -> SAM proposes mask -> Human accepts/refines -> Saved.
    Reduces annotation time from ~3min to ~20sec per object (85% reduction).
    """

    def __init__(self, sam_model):
        self.predictor = SamPredictor(sam_model)
        self.current_mask = None
        self.annotations = []

    def load_image(self, image: np.ndarray) -> None:
        """Load image and compute encoder embeddings (runs ViT once)."""
        self.image = image
        self.predictor.set_image(image)
        self.annotations = []

    def click_segment(self, x: int, y: int, is_foreground: bool = True) -> np.ndarray:
        """Single click returns best mask candidate from SAM."""
        label = 1 if is_foreground else 0
        masks, scores, _ = self.predictor.predict(
            point_coords=np.array([[x, y]]),
            point_labels=np.array([label]),
            multimask_output=True
        )
        self.current_mask = masks[np.argmax(scores)]
        return self.current_mask

    def refine_with_points(self, foreground_points: list,
                            background_points: list) -> np.ndarray:
        """
        Iteratively refine the mask by adding more positive/negative points.
        Passes the previous mask as a prior to guide the new prediction.
        """
        all_points = foreground_points + background_points
        all_labels = [1] * len(foreground_points) + [0] * len(background_points)

        masks, scores, _ = self.predictor.predict(
            point_coords=np.array(all_points),
            point_labels=np.array(all_labels),
            mask_input=self.current_mask[None],  # use previous mask as prior
            multimask_output=False
        )
        self.current_mask = masks[0]
        return self.current_mask

    def accept_mask(self, class_id: int) -> None:
        """Accept current mask and add to annotation list."""
        if self.current_mask is not None:
            self.annotations.append({
                'mask': self.current_mask.copy(),
                'class_id': class_id,
                'area': int(self.current_mask.sum())
            })
            self.current_mask = None

    def export_annotations(self, output_path: str) -> None:
        """Export annotations as binary PNG masks."""
        import json
        metadata = []
        for i, ann in enumerate(self.annotations):
            mask_img = (ann['mask'] * 255).astype(np.uint8)
            cv2.imwrite(f"{output_path}/mask_{i:04d}.png", mask_img)
            metadata.append({'mask_file': f'mask_{i:04d}.png', 'class_id': ann['class_id']})

        with open(f"{output_path}/annotations.json", 'w') as f:
            json.dump(metadata, f, indent=2)

        print(f"Exported {len(self.annotations)} annotations to {output_path}")

import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image
from pathlib import Path
import albumentations as A
from albumentations.pytorch import ToTensorV2
import numpy as np

class LungXrayDataset(Dataset):
    """
    PyTorch Dataset for lung segmentation from chest X-rays (Montgomery CXR).
    Handles grayscale images with domain-specific augmentation.
    """

    def __init__(self, image_dir: str, mask_dir: str,
                 img_size: int = 512, augment: bool = True):
        self.image_paths = sorted(Path(image_dir).glob('*.png'))
        self.mask_dir = Path(mask_dir)
        self.img_size = img_size

        if augment:
            self.transform = A.Compose([
                A.RandomResizedCrop(img_size, img_size, scale=(0.8, 1.0)),
                A.HorizontalFlip(p=0.5),
                # Conservative rotation: chest X-rays are never upside down
                A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1,
                                   rotate_limit=15, p=0.7),
                A.OneOf([
                    A.GaussNoise(var_limit=(10, 50)),
                    A.GaussianBlur(blur_limit=3),
                    A.MedianBlur(blur_limit=3),
                ], p=0.3),
                A.RandomBrightnessContrast(brightness_limit=0.2,
                                           contrast_limit=0.2, p=0.5),
                # CLAHE: crucial for X-ray contrast enhancement
                A.CLAHE(clip_limit=2, tile_grid_size=(8, 8), p=0.3),
                A.Normalize(mean=[0.485], std=[0.229]),  # grayscale stats
                ToTensorV2()
            ])
        else:
            self.transform = A.Compose([
                A.Resize(img_size, img_size),
                A.Normalize(mean=[0.485], std=[0.229]),
                ToTensorV2()
            ])

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        img_path  = self.image_paths[idx]
        mask_path = self.mask_dir / img_path.name

        # Load as grayscale (single channel)
        image = np.array(Image.open(img_path).convert('L'))
        mask  = np.array(Image.open(mask_path).convert('L'))

        # Binarize mask: pixels > 127 are lung tissue
        mask = (mask > 127).astype(np.float32)

        transformed = self.transform(image=image, mask=mask)
        image_tensor = transformed['image']
        mask_tensor  = transformed['mask'].unsqueeze(0)  # add channel dim

        return image_tensor, mask_tensor

def run_lung_segmentation_pipeline():
    """
    End-to-end pipeline: data loading -> model training -> evaluation -> metrics.
    Expected results on Montgomery dataset: Dice ~0.97, IoU ~0.94
    """
    # --- Data loading ---
    train_dataset = LungXrayDataset(
        'data/train/images', 'data/train/masks',
        img_size=512, augment=True
    )
    val_dataset = LungXrayDataset(
        'data/val/images', 'data/val/masks',
        img_size=512, augment=False
    )

    train_loader = DataLoader(train_dataset, batch_size=8,
                              shuffle=True, num_workers=4, pin_memory=True)
    val_loader   = DataLoader(val_dataset, batch_size=4,
                              shuffle=False, num_workers=4, pin_memory=True)

    # Lighter U-Net for small dataset: avoid overfitting with fewer parameters
    model = UNet(in_channels=1, num_classes=1,
                 features=[32, 64, 128, 256])

    # --- Training ---
    history = train_unet(model, train_loader, val_loader, num_epochs=100)

    # --- Final evaluation ---
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.load_state_dict(torch.load('best_unet.pth', map_location=device))
    model.eval()

    all_dice = []
    all_iou  = []
    all_hd   = []

    with torch.no_grad():
        for images, masks in val_loader:
            images, masks = images.to(device), masks.to(device)
            preds = torch.sigmoid(model(images))

            for p, m in zip(preds.cpu(), masks.cpu()):
                all_dice.append(dice_score(p, m))
                all_iou.append(compute_iou(p, m))

                # Hausdorff needs numpy
                p_np = (p.squeeze(0).numpy() >= 0.5)
                m_np = m.squeeze(0).numpy().astype(bool)
                all_hd.append(hausdorff_distance(p_np, m_np))

    print(f"\n=== Final Results on Montgomery CXR ===")
    print(f"Dice Score:           {np.mean(all_dice):.4f} +/- {np.std(all_dice):.4f}")
    print(f"IoU:                  {np.mean(all_iou):.4f} +/- {np.std(all_iou):.4f}")
    print(f"Hausdorff Distance:   {np.mean(all_hd):.1f} px (lower is better)")

run_lung_segmentation_pipeline()