files / stage2 /segmentation /train_segmentation.py

Backup missing sandbox files: stat_hardening_g, transfer_learning, segmentation, tools (1411 files, 1.13GB) (part 3)

e59a249 verified about 2 months ago

15.4 kB

"""
MT (Malignant Tumor) Segmentation — ConvNeXt-Tiny + FPN
Binary segmentation: foreground = malignant tumor region, background = everything else.
Dataset: G source-C images with CystoDS polygon annotations rasterized to binary masks.
Splits: inherit G's patient-disjoint train/val split. No test set (source D has no masks).
Input: 384×384 letterbox.
"""
import os, sys, time, json, io, random
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW
from PIL import Image, ImageDraw
import timm
from datasets import load_from_disk

# ── Paths ──
BASE = '~/kccitadmin_home/jupyterlab/sandbox/stage2/segmentation'
G_PATH = '~/kccitadmin_home/jupyterlab/sandbox/datasets/G_hf_dataset'
C_PARQUET = '~/kccitadmin_home/.cache/huggingface/hub/datasets--milkyroad--C/snapshots/b6315dd34959f8b262a988eaff70238c495f8fc7/data'
MASK_DIR = os.path.join(BASE, 'data/masks')
MANIFEST = os.path.join(BASE, 'data/manifest.csv')
RESULTS = os.path.join(BASE, 'results')
os.makedirs(RESULTS, exist_ok=True)

# ── Config ──
IMG_SIZE = 384
BATCH_SIZE = 32
EPOCHS = 80
LR = 1e-4
WEIGHT_DECAY = 1e-2
SEED = 42
BACKBONE = 'convnext_tiny.fb_in22k_ft_in1k'
DEVICE = 'cuda'
NUM_WORKERS = 4

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)

# ── Dataset ──
class MTSegDataset(Dataset):
    def __init__(self, manifest_df, g_dd, c_df, mask_dir, split, img_size=384, augment=False):
        self.split = split
        self.img_size = img_size
        self.augment = augment

        # Filter: only MT images with non-empty masks
        df = manifest_df[(manifest_df['split'] == split) & (manifest_df['has_mt_mask'] == 1)].copy()
        self.filenames = df['filename'].tolist()
        self.patient_ids = df['patient_id'].tolist()
        self.stems = df['stem'].tolist()

        # Load G images for these filenames — HF Dataset returns PIL Image directly
        fn_set = set(self.filenames)
        self.g_images = {}
        for split_name in ['train', 'val', 'test']:
            for item in g_dd[split_name]:
                fn = item['original_filename']
                if fn in fn_set:
                    img = item['image']
                    # Could be PIL Image or dict with bytes
                    if isinstance(img, dict) and 'bytes' in img:
                        self.g_images[fn] = img['bytes']
                    elif isinstance(img, bytes):
                        self.g_images[fn] = img
                    else:
                        # PIL Image — convert to bytes
                        import io as _io
                        buf = _io.BytesIO()
                        img.save(buf, format='PNG')
                        self.g_images[fn] = buf.getvalue()

        # Load C images as fallback
        self.c_lookup = {}
        for _, row in c_df.iterrows():
            fn = row['filename']
            if fn in fn_set and fn not in self.g_images:
                img = row['image']
                if isinstance(img, dict) and 'bytes' in img:
                    self.c_lookup[fn] = img['bytes']
                elif isinstance(img, bytes):
                    self.c_lookup[fn] = img
                else:
                    buf = _io.BytesIO()
                    img.save(buf, format='PNG')
                    self.c_lookup[fn] = buf.getvalue()

        print(f"  [{split}] {len(self.filenames)} MT seg images, {df['patient_id'].nunique()} patients")

    def __len__(self):
        return len(self.filenames)

    def _get_image(self, fname):
        if fname in self.g_images:
            return Image.open(io.BytesIO(self.g_images[fname])).convert('RGB')
        elif fname in self.c_lookup:
            return Image.open(io.BytesIO(self.c_lookup[fname])).convert('RGB')
        raise FileNotFoundError(f"Image {fname} not found")

    def __getitem__(self, idx):
        fname = self.filenames[idx]
        stem = self.stems[idx]

        img = self._get_image(fname)
        W, H = img.size  # original 352x240

        mask = np.load(os.path.join(MASK_DIR, f"{stem}.npy")).astype(np.float32)  # (H, W) binary

        # Letterbox resize
        img, mask = self._letterbox(img, mask, self.img_size)

        if self.augment:
            img, mask = self._augment(img, mask)

        # To tensor
        img_t = torch.from_numpy(np.array(img, dtype=np.float32)).permute(2, 0, 1) / 255.0
        # Normalize with ImageNet stats
        mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
        std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
        img_t = (img_t - mean) / std
        mask_t = torch.from_numpy(mask).long()  # (H, W) {0, 1}

        return img_t, mask_t

    def _letterbox(self, img, mask, target_size):
        """Letterbox: resize keeping aspect ratio, pad with zeros."""
        W, H = img.size
        scale = min(target_size / W, target_size / H)
        new_w = int(W * scale)
        new_h = int(H * scale)

        img_resized = img.resize((new_w, new_h), Image.BILINEAR)
        # Mask: use nearest to keep binary
        mask_pil = Image.fromarray(mask.astype(np.uint8))
        mask_resized = mask_pil.resize((new_w, new_h), Image.NEAREST)

        # Pad to target_size
        pad_w = target_size - new_w
        pad_h = target_size - new_h
        top = pad_h // 2
        bottom = pad_h - top
        left = pad_w // 2
        right = pad_w - left

        img_pad = Image.new('RGB', (target_size, target_size), (0, 0, 0))
        img_pad.paste(img_resized, (left, top))
        mask_pad = Image.new('L', (target_size, target_size), 0)
        mask_pad.paste(mask_resized, (left, top))

        return img_pad, np.array(mask_pad, dtype=np.float32)

    def _augment(self, img, mask):
        """Joint augmentation: hflip, vflip, rotation, scale, color jitter."""
        # Random horizontal flip
        if random.random() < 0.5:
            img = img.transpose(Image.FLIP_LEFT_RIGHT)
            mask = np.fliplr(mask).copy()
        # Random vertical flip
        if random.random() < 0.5:
            img = img.transpose(Image.FLIP_TOP_BOTTOM)
            mask = np.flipud(mask).copy()
        # Random rotation (-15 to 15 degrees)
        if random.random() < 0.5:
            angle = random.uniform(-15, 15)
            img = img.rotate(angle, fillcolor=(0, 0, 0))
            mask_pil = Image.fromarray(mask.astype(np.uint8))
            mask_pil = mask_pil.rotate(angle, fillcolor=0)
            mask = np.array(mask_pil, dtype=np.float32)
        # Color jitter (brightness/contrast only, applied to image)
        if random.random() < 0.5:
            factor = random.uniform(0.8, 1.2)
            img = Image.fromarray(np.clip(np.array(img, dtype=np.float32) * factor, 0, 255).astype(np.uint8))
        return img, mask


# ── Model: ConvNeXt-Tiny + FPN ──
class FPNDecoder(nn.Module):
    """Simple FPN decoder for binary segmentation."""
    def __init__(self, in_channels=[96, 192, 384, 768], out_channels=256):
        super().__init__()
        self.lateral_convs = nn.ModuleList([
            nn.Conv2d(c, out_channels, 1) for c in in_channels
        ])
        self.merge_convs = nn.ModuleList([
            nn.Conv2d(out_channels, out_channels, 3, padding=1) for _ in in_channels
        ])
        # Final segmentation head
        self.head = nn.Sequential(
            nn.Conv2d(out_channels * len(in_channels), out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, 2, 1),  # binary: background + foreground
        )

    def forward(self, features):
        # features: [C2, C3, C4, C5] from coarse to fine
        # Lateral connections
        laterals = [conv(f) for conv, f in zip(self.lateral_convs, features)]
        # Top-down pathway
        for i in range(len(laterals) - 2, -1, -1):
            laterals[i] = laterals[i] + F.interpolate(laterals[i + 1], size=laterals[i].shape[2:], mode='bilinear', align_corners=False)
        # Smooth
        outs = [conv(lat) for conv, lat in zip(self.merge_convs, laterals)]
        # Upsample all to finest resolution
        target_size = outs[0].shape[2:]
        upscaled = [F.interpolate(o, size=target_size, mode='bilinear', align_corners=False) for o in outs[1:]]
        concat = torch.cat([outs[0]] + upscaled, dim=1)
        return self.head(concat)


class SegModel(nn.Module):
    def __init__(self, backbone_name=BACKBONE, pretrained=True):
        super().__init__()
        self.backbone = timm.create_model(backbone_name, pretrained=pretrained, features_only=True, out_indices=(0, 1, 2, 3))
        self.decoder = FPNDecoder(in_channels=[f['num_chs'] for f in self.backbone.feature_info])

    def forward(self, x):
        features = self.backbone(x)
        out = self.decoder(features)
        # Upsample to input size
        out = F.interpolate(out, size=x.shape[2:], mode='bilinear', align_corners=False)
        return out


# ── Metrics ──
def dice_score(pred, target, smooth=1e-6):
    """Dice for foreground class (class 1)."""
    pred_binary = (pred == 1).float()
    target_binary = (target == 1).float()
    intersection = (pred_binary * target_binary).sum()
    return (2.0 * intersection + smooth) / (pred_binary.sum() + target_binary.sum() + smooth)


def iou_score(pred, target, smooth=1e-6):
    pred_binary = (pred == 1).float()
    target_binary = (target == 1).float()
    intersection = (pred_binary * target_binary).sum()
    union = pred_binary.sum() + target_binary.sum() - intersection
    return (intersection + smooth) / (union + smooth)


# ── Loss: Dice + CE ──
class DiceCELoss(nn.Module):
    def __init__(self, dice_weight=1.0, ce_weight=1.0):
        super().__init__()
        self.dice_weight = dice_weight
        self.ce_weight = ce_weight
        self.ce = nn.CrossEntropyLoss(weight=torch.tensor([0.2, 0.8]).to(DEVICE))  # weight foreground higher

    def forward(self, logits, target):
        ce_loss = self.ce(logits, target)
        # Dice
        probs = F.softmax(logits, dim=1)
        pred_fg = probs[:, 1]
        target_fg = (target == 1).float()
        smooth = 1e-6
        intersection = (pred_fg * target_fg).sum(dim=(1, 2))
        dice = (2.0 * intersection + smooth) / (pred_fg.sum(dim=(1, 2)) + target_fg.sum(dim=(1, 2)) + smooth)
        dice_loss = 1.0 - dice.mean()
        return self.dice_weight * dice_loss + self.ce_weight * ce_loss


# ── Training ──
def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss = 0
    total_dice = 0
    n = 0
    for imgs, masks in loader:
        imgs, masks = imgs.to(device), masks.to(device)
        logits = model(imgs)
        loss = criterion(logits, masks)
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        total_loss += loss.item() * imgs.size(0)
        preds = logits.argmax(dim=1)
        total_dice += dice_score(preds, masks).item() * imgs.size(0)
        n += imgs.size(0)
    return total_loss / n, total_dice / n


@torch.no_grad()
def validate(model, loader, device):
    model.eval()
    total_loss = 0
    total_dice = 0
    total_iou = 0
    n = 0
    for imgs, masks in loader:
        imgs, masks = imgs.to(device), masks.to(device)
        logits = model(imgs)
        loss = criterion(logits, masks)
        preds = logits.argmax(dim=1)
        total_loss += loss.item() * imgs.size(0)
        total_dice += dice_score(preds, masks).item() * imgs.size(0)
        total_iou += iou_score(preds, masks).item() * imgs.size(0)
        n += imgs.size(0)
    return total_loss / n, total_dice / n, total_iou / n


if __name__ == '__main__':
    print("=" * 60)
    print("MT Segmentation — ConvNeXt-Tiny + FPN")
    print("=" * 60)

    # Load data
    manifest = pd.read_csv(MANIFEST)
    print("\nLoading G dataset...")
    g_dd = load_from_disk(G_PATH)
    print("Loading C parquet (for image fallback)...")
    c_dfs = []
    for f in sorted(os.listdir(C_PARQUET)):
        if f.endswith('.parquet'):
            c_dfs.append(pd.read_parquet(os.path.join(C_PARQUET, f), columns=['filename', 'image', 'has_segmentation']))
    c_df = pd.concat(c_dfs, ignore_index=True)

    train_ds = MTSegDataset(manifest, g_dd, c_df, MASK_DIR, 'train', IMG_SIZE, augment=True)
    val_ds = MTSegDataset(manifest, g_dd, c_df, MASK_DIR, 'val', IMG_SIZE, augment=False)

    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, pin_memory=True, drop_last=True)
    val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=True)

    # Model
    model = SegModel(BACKBONE, pretrained=True).to(DEVICE)
    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"\nModel: {BACKBONE} + FPN, {n_params/1e6:.1f}M params")

    criterion = DiceCELoss().to(DEVICE)
    optimizer = AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS, eta_min=1e-6)

    # Training loop
    best_dice = 0
    history = []
    for epoch in range(1, EPOCHS + 1):
        t0 = time.time()
        train_loss, train_dice = train_one_epoch(model, train_loader, criterion, optimizer, DEVICE)
        val_loss, val_dice, val_iou = validate(model, val_loader, DEVICE)
        scheduler.step()
        dt = time.time() - t0

        history.append({
            'epoch': epoch, 'train_loss': train_loss, 'train_dice': train_dice,
            'val_loss': val_loss, 'val_dice': val_dice, 'val_iou': val_iou, 'lr': optimizer.param_groups[0]['lr']
        })

        if val_dice > best_dice:
            best_dice = val_dice
            torch.save(model.state_dict(), os.path.join(RESULTS, 'best_model.pt'))

        if epoch % 5 == 0 or epoch == 1 or epoch == EPOCHS:
            print(f"Epoch {epoch:3d}/{EPOCHS} | Loss: {train_loss:.4f} | Train Dice: {train_dice:.4f} | Val Loss: {val_loss:.4f} | Val Dice: {val_dice:.4f} | Val IoU: {val_iou:.4f} | {dt:.1f}s")

    # Save history
    with open(os.path.join(RESULTS, 'history.json'), 'w') as f:
        json.dump(history, f, indent=2)

    # Save final results
    results = {
        'best_val_dice': best_dice,
        'final_val_dice': history[-1]['val_dice'],
        'final_val_iou': history[-1]['val_iou'],
        'config': {
            'backbone': BACKBONE, 'img_size': IMG_SIZE, 'batch_size': BATCH_SIZE,
            'epochs': EPOCHS, 'lr': LR, 'weight_decay': WEIGHT_DECAY, 'seed': SEED,
        },
        'train_images': len(train_ds), 'val_images': len(val_ds),
        'train_patients': len(set(train_ds.patient_ids)),
        'val_patients': len(set(val_ds.patient_ids)),
    }
    with open(os.path.join(RESULTS, 'results.json'), 'w') as f:
        json.dump(results, f, indent=2)

    print(f"\n{'='*60}")
    print(f"BEST VALIDATION DICE: {best_dice:.4f}")
    print(f"Final Val Dice: {history[-1]['val_dice']:.4f}")
    print(f"Final Val IoU:  {history[-1]['val_iou']:.4f}")
    print(f"{'='*60}")

	"""
	MT (Malignant Tumor) Segmentation — ConvNeXt-Tiny + FPN
	Binary segmentation: foreground = malignant tumor region, background = everything else.

	Dataset: G source-C images with CystoDS polygon annotations rasterized to binary masks.
	Splits: inherit G's patient-disjoint train/val split. No test set (source D has no masks).
	Input: 384×384 letterbox.
	"""
	import os, sys, time, json, io, random
	import numpy as np
	import pandas as pd
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from torch.utils.data import Dataset, DataLoader
	from torch.optim import AdamW
	from PIL import Image, ImageDraw
	import timm
	from datasets import load_from_disk

	# ── Paths ──
	BASE = '~/kccitadmin_home/jupyterlab/sandbox/stage2/segmentation'
	G_PATH = '~/kccitadmin_home/jupyterlab/sandbox/datasets/G_hf_dataset'
	C_PARQUET = '~/kccitadmin_home/.cache/huggingface/hub/datasets--milkyroad--C/snapshots/b6315dd34959f8b262a988eaff70238c495f8fc7/data'
	MASK_DIR = os.path.join(BASE, 'data/masks')
	MANIFEST = os.path.join(BASE, 'data/manifest.csv')
	RESULTS = os.path.join(BASE, 'results')
	os.makedirs(RESULTS, exist_ok=True)

	# ── Config ──
	IMG_SIZE = 384
	BATCH_SIZE = 32
	EPOCHS = 80
	LR = 1e-4
	WEIGHT_DECAY = 1e-2
	SEED = 42
	BACKBONE = 'convnext_tiny.fb_in22k_ft_in1k'
	DEVICE = 'cuda'
	NUM_WORKERS = 4

	random.seed(SEED)
	np.random.seed(SEED)
	torch.manual_seed(SEED)
	torch.cuda.manual_seed_all(SEED)

	# ── Dataset ──
	class MTSegDataset(Dataset):
	def __init__(self, manifest_df, g_dd, c_df, mask_dir, split, img_size=384, augment=False):
	self.split = split
	self.img_size = img_size
	self.augment = augment

	# Filter: only MT images with non-empty masks
	df = manifest_df[(manifest_df['split'] == split) & (manifest_df['has_mt_mask'] == 1)].copy()
	self.filenames = df['filename'].tolist()
	self.patient_ids = df['patient_id'].tolist()
	self.stems = df['stem'].tolist()

	# Load G images for these filenames — HF Dataset returns PIL Image directly
	fn_set = set(self.filenames)
	self.g_images = {}
	for split_name in ['train', 'val', 'test']:
	for item in g_dd[split_name]:
	fn = item['original_filename']
	if fn in fn_set:
	img = item['image']
	# Could be PIL Image or dict with bytes
	if isinstance(img, dict) and 'bytes' in img:
	self.g_images[fn] = img['bytes']
	elif isinstance(img, bytes):
	self.g_images[fn] = img
	else:
	# PIL Image — convert to bytes
	import io as _io
	buf = _io.BytesIO()
	img.save(buf, format='PNG')
	self.g_images[fn] = buf.getvalue()

	# Load C images as fallback
	self.c_lookup = {}
	for _, row in c_df.iterrows():
	fn = row['filename']
	if fn in fn_set and fn not in self.g_images:
	img = row['image']
	if isinstance(img, dict) and 'bytes' in img:
	self.c_lookup[fn] = img['bytes']
	elif isinstance(img, bytes):
	self.c_lookup[fn] = img
	else:
	buf = _io.BytesIO()
	img.save(buf, format='PNG')
	self.c_lookup[fn] = buf.getvalue()

	print(f" [{split}] {len(self.filenames)} MT seg images, {df['patient_id'].nunique()} patients")

	def __len__(self):
	return len(self.filenames)

	def _get_image(self, fname):
	if fname in self.g_images:
	return Image.open(io.BytesIO(self.g_images[fname])).convert('RGB')
	elif fname in self.c_lookup:
	return Image.open(io.BytesIO(self.c_lookup[fname])).convert('RGB')
	raise FileNotFoundError(f"Image {fname} not found")

	def __getitem__(self, idx):
	fname = self.filenames[idx]
	stem = self.stems[idx]

	img = self._get_image(fname)
	W, H = img.size # original 352x240

	mask = np.load(os.path.join(MASK_DIR, f"{stem}.npy")).astype(np.float32) # (H, W) binary

	# Letterbox resize
	img, mask = self._letterbox(img, mask, self.img_size)

	if self.augment:
	img, mask = self._augment(img, mask)

	# To tensor
	img_t = torch.from_numpy(np.array(img, dtype=np.float32)).permute(2, 0, 1) / 255.0
	# Normalize with ImageNet stats
	mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
	std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
	img_t = (img_t - mean) / std
	mask_t = torch.from_numpy(mask).long() # (H, W) {0, 1}

	return img_t, mask_t

	def _letterbox(self, img, mask, target_size):
	"""Letterbox: resize keeping aspect ratio, pad with zeros."""
	W, H = img.size
	scale = min(target_size / W, target_size / H)
	new_w = int(W * scale)
	new_h = int(H * scale)

	img_resized = img.resize((new_w, new_h), Image.BILINEAR)
	# Mask: use nearest to keep binary
	mask_pil = Image.fromarray(mask.astype(np.uint8))
	mask_resized = mask_pil.resize((new_w, new_h), Image.NEAREST)

	# Pad to target_size
	pad_w = target_size - new_w
	pad_h = target_size - new_h
	top = pad_h // 2
	bottom = pad_h - top
	left = pad_w // 2
	right = pad_w - left

	img_pad = Image.new('RGB', (target_size, target_size), (0, 0, 0))
	img_pad.paste(img_resized, (left, top))
	mask_pad = Image.new('L', (target_size, target_size), 0)
	mask_pad.paste(mask_resized, (left, top))

	return img_pad, np.array(mask_pad, dtype=np.float32)

	def _augment(self, img, mask):
	"""Joint augmentation: hflip, vflip, rotation, scale, color jitter."""
	# Random horizontal flip
	if random.random() < 0.5:
	img = img.transpose(Image.FLIP_LEFT_RIGHT)
	mask = np.fliplr(mask).copy()
	# Random vertical flip
	if random.random() < 0.5:
	img = img.transpose(Image.FLIP_TOP_BOTTOM)
	mask = np.flipud(mask).copy()
	# Random rotation (-15 to 15 degrees)
	if random.random() < 0.5:
	angle = random.uniform(-15, 15)
	img = img.rotate(angle, fillcolor=(0, 0, 0))
	mask_pil = Image.fromarray(mask.astype(np.uint8))
	mask_pil = mask_pil.rotate(angle, fillcolor=0)
	mask = np.array(mask_pil, dtype=np.float32)
	# Color jitter (brightness/contrast only, applied to image)
	if random.random() < 0.5:
	factor = random.uniform(0.8, 1.2)
	img = Image.fromarray(np.clip(np.array(img, dtype=np.float32) * factor, 0, 255).astype(np.uint8))
	return img, mask


	# ── Model: ConvNeXt-Tiny + FPN ──
	class FPNDecoder(nn.Module):
	"""Simple FPN decoder for binary segmentation."""
	def __init__(self, in_channels=[96, 192, 384, 768], out_channels=256):
	super().__init__()
	self.lateral_convs = nn.ModuleList([
	nn.Conv2d(c, out_channels, 1) for c in in_channels
	])
	self.merge_convs = nn.ModuleList([
	nn.Conv2d(out_channels, out_channels, 3, padding=1) for _ in in_channels
	])
	# Final segmentation head
	self.head = nn.Sequential(
	nn.Conv2d(out_channels * len(in_channels), out_channels, 3, padding=1),
	nn.BatchNorm2d(out_channels),
	nn.ReLU(inplace=True),
	nn.Conv2d(out_channels, out_channels, 3, padding=1),
	nn.BatchNorm2d(out_channels),
	nn.ReLU(inplace=True),
	nn.Conv2d(out_channels, 2, 1), # binary: background + foreground
	)

	def forward(self, features):
	# features: [C2, C3, C4, C5] from coarse to fine
	# Lateral connections
	laterals = [conv(f) for conv, f in zip(self.lateral_convs, features)]
	# Top-down pathway
	for i in range(len(laterals) - 2, -1, -1):
	laterals[i] = laterals[i] + F.interpolate(laterals[i + 1], size=laterals[i].shape[2:], mode='bilinear', align_corners=False)
	# Smooth
	outs = [conv(lat) for conv, lat in zip(self.merge_convs, laterals)]
	# Upsample all to finest resolution
	target_size = outs[0].shape[2:]
	upscaled = [F.interpolate(o, size=target_size, mode='bilinear', align_corners=False) for o in outs[1:]]
	concat = torch.cat([outs[0]] + upscaled, dim=1)
	return self.head(concat)


	class SegModel(nn.Module):
	def __init__(self, backbone_name=BACKBONE, pretrained=True):
	super().__init__()
	self.backbone = timm.create_model(backbone_name, pretrained=pretrained, features_only=True, out_indices=(0, 1, 2, 3))
	self.decoder = FPNDecoder(in_channels=[f['num_chs'] for f in self.backbone.feature_info])

	def forward(self, x):
	features = self.backbone(x)
	out = self.decoder(features)
	# Upsample to input size
	out = F.interpolate(out, size=x.shape[2:], mode='bilinear', align_corners=False)
	return out


	# ── Metrics ──
	def dice_score(pred, target, smooth=1e-6):
	"""Dice for foreground class (class 1)."""
	pred_binary = (pred == 1).float()
	target_binary = (target == 1).float()
	intersection = (pred_binary * target_binary).sum()
	return (2.0 * intersection + smooth) / (pred_binary.sum() + target_binary.sum() + smooth)


	def iou_score(pred, target, smooth=1e-6):
	pred_binary = (pred == 1).float()
	target_binary = (target == 1).float()
	intersection = (pred_binary * target_binary).sum()
	union = pred_binary.sum() + target_binary.sum() - intersection
	return (intersection + smooth) / (union + smooth)


	# ── Loss: Dice + CE ──
	class DiceCELoss(nn.Module):
	def __init__(self, dice_weight=1.0, ce_weight=1.0):
	super().__init__()
	self.dice_weight = dice_weight
	self.ce_weight = ce_weight
	self.ce = nn.CrossEntropyLoss(weight=torch.tensor([0.2, 0.8]).to(DEVICE)) # weight foreground higher

	def forward(self, logits, target):
	ce_loss = self.ce(logits, target)
	# Dice
	probs = F.softmax(logits, dim=1)
	pred_fg = probs[:, 1]
	target_fg = (target == 1).float()
	smooth = 1e-6
	intersection = (pred_fg * target_fg).sum(dim=(1, 2))
	dice = (2.0 * intersection + smooth) / (pred_fg.sum(dim=(1, 2)) + target_fg.sum(dim=(1, 2)) + smooth)
	dice_loss = 1.0 - dice.mean()
	return self.dice_weight * dice_loss + self.ce_weight * ce_loss


	# ── Training ──
	def train_one_epoch(model, loader, criterion, optimizer, device):
	model.train()
	total_loss = 0
	total_dice = 0
	n = 0
	for imgs, masks in loader:
	imgs, masks = imgs.to(device), masks.to(device)
	logits = model(imgs)
	loss = criterion(logits, masks)
	optimizer.zero_grad()
	loss.backward()
	torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
	optimizer.step()
	total_loss += loss.item() * imgs.size(0)
	preds = logits.argmax(dim=1)
	total_dice += dice_score(preds, masks).item() * imgs.size(0)
	n += imgs.size(0)
	return total_loss / n, total_dice / n


	@torch.no_grad()
	def validate(model, loader, device):
	model.eval()
	total_loss = 0
	total_dice = 0
	total_iou = 0
	n = 0
	for imgs, masks in loader:
	imgs, masks = imgs.to(device), masks.to(device)
	logits = model(imgs)
	loss = criterion(logits, masks)
	preds = logits.argmax(dim=1)
	total_loss += loss.item() * imgs.size(0)
	total_dice += dice_score(preds, masks).item() * imgs.size(0)
	total_iou += iou_score(preds, masks).item() * imgs.size(0)
	n += imgs.size(0)
	return total_loss / n, total_dice / n, total_iou / n


	if __name__ == '__main__':
	print("=" * 60)
	print("MT Segmentation — ConvNeXt-Tiny + FPN")
	print("=" * 60)

	# Load data
	manifest = pd.read_csv(MANIFEST)
	print("\nLoading G dataset...")
	g_dd = load_from_disk(G_PATH)
	print("Loading C parquet (for image fallback)...")
	c_dfs = []
	for f in sorted(os.listdir(C_PARQUET)):
	if f.endswith('.parquet'):
	c_dfs.append(pd.read_parquet(os.path.join(C_PARQUET, f), columns=['filename', 'image', 'has_segmentation']))
	c_df = pd.concat(c_dfs, ignore_index=True)

	train_ds = MTSegDataset(manifest, g_dd, c_df, MASK_DIR, 'train', IMG_SIZE, augment=True)
	val_ds = MTSegDataset(manifest, g_dd, c_df, MASK_DIR, 'val', IMG_SIZE, augment=False)

	train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, pin_memory=True, drop_last=True)
	val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=True)

	# Model
	model = SegModel(BACKBONE, pretrained=True).to(DEVICE)
	n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
	print(f"\nModel: {BACKBONE} + FPN, {n_params/1e6:.1f}M params")

	criterion = DiceCELoss().to(DEVICE)
	optimizer = AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
	scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS, eta_min=1e-6)

	# Training loop
	best_dice = 0
	history = []
	for epoch in range(1, EPOCHS + 1):
	t0 = time.time()
	train_loss, train_dice = train_one_epoch(model, train_loader, criterion, optimizer, DEVICE)
	val_loss, val_dice, val_iou = validate(model, val_loader, DEVICE)
	scheduler.step()
	dt = time.time() - t0

	history.append({
	'epoch': epoch, 'train_loss': train_loss, 'train_dice': train_dice,
	'val_loss': val_loss, 'val_dice': val_dice, 'val_iou': val_iou, 'lr': optimizer.param_groups[0]['lr']
	})

	if val_dice > best_dice:
	best_dice = val_dice
	torch.save(model.state_dict(), os.path.join(RESULTS, 'best_model.pt'))

	if epoch % 5 == 0 or epoch == 1 or epoch == EPOCHS:
	print(f"Epoch {epoch:3d}/{EPOCHS} \| Loss: {train_loss:.4f} \| Train Dice: {train_dice:.4f} \| Val Loss: {val_loss:.4f} \| Val Dice: {val_dice:.4f} \| Val IoU: {val_iou:.4f} \| {dt:.1f}s")

	# Save history
	with open(os.path.join(RESULTS, 'history.json'), 'w') as f:
	json.dump(history, f, indent=2)

	# Save final results
	results = {
	'best_val_dice': best_dice,
	'final_val_dice': history[-1]['val_dice'],
	'final_val_iou': history[-1]['val_iou'],
	'config': {
	'backbone': BACKBONE, 'img_size': IMG_SIZE, 'batch_size': BATCH_SIZE,
	'epochs': EPOCHS, 'lr': LR, 'weight_decay': WEIGHT_DECAY, 'seed': SEED,
	},
	'train_images': len(train_ds), 'val_images': len(val_ds),
	'train_patients': len(set(train_ds.patient_ids)),
	'val_patients': len(set(val_ds.patient_ids)),
	}
	with open(os.path.join(RESULTS, 'results.json'), 'w') as f:
	json.dump(results, f, indent=2)

	print(f"\n{'='*60}")
	print(f"BEST VALIDATION DICE: {best_dice:.4f}")
	print(f"Final Val Dice: {history[-1]['val_dice']:.4f}")
	print(f"Final Val IoU: {history[-1]['val_iou']:.4f}")
	print(f"{'='*60}")