""" 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}")