{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":113558,"databundleVersionId":14174843,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":28970,"sourceType":"datasetVersion","datasetId":22588},{"sourceId":13610809,"sourceType":"datasetVersion","datasetId":8643941}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"80691152","cell_type":"markdown","source":"# Scientific Image Forgery — Clean Training & Inference Pipeline\n\nThis notebook consolidates preprocessing, validation (image-level F1), calibration, and inference with consistent transforms. Old conflicting cells were removed.\nKey changes:\n- Unified preprocessing across train/val/test (identical resize/normalize/morphology)\n- Image-level F1 metric (leaderboard) computed on validation\n- Grid-search over (pixel_threshold, area_frac)\n- Save & reuse calibrated params in checkpoint for test-time inference","metadata":{}},{"id":"73072ec4","cell_type":"code","source":"\n# === CONFIG & IMPORTS (NEW) ===\nimport os, math, json, numpy as np, torch, torch.nn.functional as F\nfrom sklearn.metrics import f1_score\nimport torchvision.transforms as T\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Pick ONE size and keep it everywhere (train/val/test)\nIMAGE_SIZE = 518  # or 256, just be consistent across the notebook\n\n# Normalization (keep consistent)\nMEAN = (0.485, 0.456, 0.406)\nSTD  = (0.229, 0.224, 0.225)\n\n# Optional morphology; if you enable it, it MUST be the same in val & test\nUSE_MORPH = False\nMORPH_KERNEL = 3\n\n# Where to save the best model\nCKPT_PATH = \"best_model_calibrated.pth\"\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-10T12:44:19.645865Z","iopub.execute_input":"2025-11-10T12:44:19.646131Z","iopub.status.idle":"2025-11-10T12:44:23.832185Z","shell.execute_reply.started":"2025-11-10T12:44:19.646111Z","shell.execute_reply":"2025-11-10T12:44:23.831267Z"}},"outputs":[],"execution_count":null},{"id":"e965fe30","cell_type":"code","source":"\n# === PREPROCESSING (REPLACE) ===\n# Use the SAME preprocessing everywhere.\n\nimg_transform = T.Compose([\n    T.Resize((IMAGE_SIZE, IMAGE_SIZE), interpolation=T.InterpolationMode.BICUBIC),\n    T.ToTensor(),\n    T.Normalize(MEAN, STD),\n])\n\ndef resize_mask(mask_tensor, size=(IMAGE_SIZE, IMAGE_SIZE)):\n    \"\"\"\n    mask_tensor: FloatTensor in [0,1] or {0,1}, shape HxW or 1xHxW\n    Uses NEAREST for masks to avoid blurring labels.\n    \"\"\"\n    if mask_tensor.ndim == 2:\n        mask_tensor = mask_tensor.unsqueeze(0)  # 1xHxW\n    return F.interpolate(mask_tensor.unsqueeze(0), size=size, mode='nearest').squeeze(0)\n\ndef maybe_morphology(mask_bin):\n    \"\"\"\n    mask_bin: FloatTensor 1xHxW with values {0,1}. Applies opening if USE_MORPH.\n    Keep this identical in validation and test if you enable it.\n    \"\"\"\n    if not USE_MORPH:\n        return mask_bin\n    k = MORPH_KERNEL\n    pad = k // 2\n    kernel = torch.ones((1,1,k,k), device=mask_bin.device)\n    # Erode: pixel stays 1 only if every pixel in the kxk neighborhood is 1\n    eroded = (F.conv2d(mask_bin.unsqueeze(0), kernel, padding=pad) == (k*k)).float()\n    # Dilate: pixel is 1 if any neighbor is 1\n    dilated = (F.conv2d(eroded, kernel, padding=pad) > 0).float()\n    return dilated.squeeze(0)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-10T12:44:23.834431Z","iopub.execute_input":"2025-11-10T12:44:23.835396Z","iopub.status.idle":"2025-11-10T12:44:23.847261Z","shell.execute_reply.started":"2025-11-10T12:44:23.835362Z","shell.execute_reply":"2025-11-10T12:44:23.846114Z"}},"outputs":[],"execution_count":null},{"id":"64a56dd1","cell_type":"markdown","source":"# Recod.ai / LUC — Scientific Image Forgery Detection \n\nThis notebook trains a U-Net segmentation model on GPU T4 to detect copy-move forgeries in biomedical images.\nIt:\n- Loads `train_images/authentic` and `train_images/forged`\n- Merges multiple `.npy` masks per case into a single binary mask\n- Trains with PyTorch/XLA (multi-core)\n- Evaluates with an F1-like metric (approx)\n- Exports a submission CSV using RLE encoding (or \"authentic\" if no mask predicted)","metadata":{}},{"id":"fa523042","cell_type":"markdown","source":"\n## ✨ Update: Data Augmentation (no deformation, size preserved)\n\nChanges made:\n- Added an **AugmentedWrapper** dataset that applies **horizontal / vertical flips** to both image and mask, and **brightness / hue** changes to the **image only** (mask never altered except for flips).\n- Rewired the training pipeline to **double** the effective number of training samples via `ConcatDataset(original, augmented)` → **X2 images** seen per epoch.\n- **No resizing / warping** was added by these changes. Image shapes stay exactly as in the original pipeline after the existing `resize_pair` step.\n","metadata":{}},{"id":"5a104818","cell_type":"code","source":"# No extra installs required; we stick to torchvision to avoid dependency conflicts.\nimport torch, torchvision, sys, numpy\nprint(\"Torch:\", torch.__version__)\nprint(\"Torchvision:\", torchvision.__version__)\nprint(\"Python:\", sys.version)\nprint(\"CUDA available:\", torch.cuda.is_available())","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:44:23.849100Z","iopub.execute_input":"2025-11-10T12:44:23.849391Z","iopub.status.idle":"2025-11-10T12:44:23.872766Z","shell.execute_reply.started":"2025-11-10T12:44:23.849365Z","shell.execute_reply":"2025-11-10T12:44:23.871962Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"50a1c80c","cell_type":"code","source":"import os, glob, random\nfrom pathlib import Path\nimport numpy as np\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms.functional as TF\nfrom torchvision.transforms import InterpolationMode\nimport torchvision.models.segmentation as tvseg\n\nfrom sklearn.model_selection import StratifiedKFold\n\n# ===== NEW PATHS & CONFIG =====\nfrom pathlib import Path\nfrom typing import Optional, Dict, List\nimport random\n\n# RecoDAI LUC Scientific Image Forgery Detection\nRECODAI_AUTH_DIR = Path(\"/kaggle/input/recodai-luc-scientific-image-forgery-detection/train_images/authentic\")\nRECODAI_FORG_DIR = Path(\"/kaggle/input/recodai-luc-scientific-image-forgery-detection/train_images/forged\")\nRECODAI_MASK_DIR = Path(\"/kaggle/input/recodai-luc-scientific-image-forgery-detection/train_masks\")\n\n# Extra train-only datasets\nNUCLEI_ROOT = Path(\"/kaggle/input/nuclei-segmentation-in-microscope-cell-images/Nuclei\")\nAUGMENT_IMG_ROOT = Path(\"/kaggle/input/data-augment\")   # PNGs are somewhere under this root\nAUGMENT_MASK_ROOT = Path(\"/kaggle/input/data-augment\")  # NPys are somewhere under this root\n\n# Split config\nTRAIN_FRACTION = 0.70\nRANDOM_SEED = 42\n\n# Labels\nLABEL_AUTHENTIC = 0\nLABEL_FORGED = 1\n\nrandom.seed(RANDOM_SEED)\n\n\nOUT_DIR    = \"/kaggle/working\"\nos.makedirs(OUT_DIR, exist_ok=True)\n\n# -------------------\n# Training config\n# -------------------\nSEED          = 42\nDEVICE        = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nIMAGE_SIZE    = 518            \nBATCH_SIZE    = 6              \nEPOCHS        = 5             \nLR            = 3e-4\nWEIGHT_DECAY  = 1e-4\nN_FOLDS       = 5\nFOLD_TO_RUN   = 0              \nUSE_AMP       = True\nPRINT_EVERY   = 50\n\nprint(\"Device:\", DEVICE)","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.status.busy":"2025-11-10T12:44:23.873617Z","iopub.execute_input":"2025-11-10T12:44:23.873853Z","iopub.status.idle":"2025-11-10T12:44:23.899524Z","shell.execute_reply.started":"2025-11-10T12:44:23.873831Z","shell.execute_reply":"2025-11-10T12:44:23.898642Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"b727924b","cell_type":"code","source":"# --- DINOv2 config ---\nIMAGE_SIZE = 518     # multiple of 14 (ViT patch size); 518 = 37x14\nNUM_CLASSES = 1      # binary mask (change if you use multi-class)\nBACKBONE_SIZE = 'vitb14'  # choices: 'vits14', 'vitb14', 'vitl14', 'vitg14'\n\n# training knobs (keep your existing ones if you prefer)\nLR = 1e-4\nWEIGHT_DECAY = 1e-4\nUSE_AMP = True\n\n# normalization expected by ImageNet-pretrained backbones (incl. DINOv2)\nIMAGENET_MEAN = (0.485, 0.456, 0.406)\nIMAGENET_STD  = (0.229, 0.224, 0.225)\n\nimport torch\nimport torch.nn.functional as F\nfrom torchvision.transforms import functional as TF\n\ndef dinov2_normalize(img_tensor):\n    \"\"\"\n    img_tensor: float tensor [C,H,W] in [0,1]\n    returns: normalized tensor\n    \"\"\"\n    return TF.normalize(img_tensor, IMAGENET_MEAN, IMAGENET_STD)\n","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:44:23.900378Z","iopub.execute_input":"2025-11-10T12:44:23.900901Z","iopub.status.idle":"2025-11-10T12:44:23.916744Z","shell.execute_reply.started":"2025-11-10T12:44:23.900875Z","shell.execute_reply":"2025-11-10T12:44:23.915888Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"67816d46","cell_type":"code","source":"# ===== HELPERS =====\ndef map_by_stem(paths):\n    \"\"\"Return a dict {stem: path} for quick filename-based lookups.\"\"\"\n    return {p.stem: p for p in paths}\n\ndef to_item(img_path: Path, label: int, mask_path: Optional[Path], src: str) -> Dict:\n    return {\n        \"path\": str(img_path),                              # <— required by your __getitem__\n        \"img\": str(img_path),                               # compatibility\n        \"mask\": (str(mask_path) if mask_path is not None else None),\n        \"mask_path\": (str(mask_path) if mask_path is not None else None),\n        \"label\": int(label),\n        \"src\": src,\n    }\n\n\ndef debug_counts(items: List[Dict], title=\"Items\"):\n    from collections import Counter\n    c = Counter([it[\"label\"] for it in items])\n    print(f\"{title}: total={len(items)}  authentic={c.get(0,0)}  forged={c.get(1,0)}\")\n","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:44:23.917570Z","iopub.execute_input":"2025-11-10T12:44:23.917843Z","iopub.status.idle":"2025-11-10T12:44:23.936910Z","shell.execute_reply.started":"2025-11-10T12:44:23.917821Z","shell.execute_reply":"2025-11-10T12:44:23.936017Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"68954c01","cell_type":"code","source":"def seed_everything(seed=42):\n    import os, random\n    import numpy as np\n    import torch\n    random.seed(seed)\n    np.random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    if torch.cuda.is_available():\n        torch.manual_seed(seed)\n        torch.cuda.manual_seed_all(seed)\n        torch.backends.cudnn.deterministic = False\n        torch.backends.cudnn.benchmark = True\n\nseed_everything(SEED)","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:44:23.937763Z","iopub.execute_input":"2025-11-10T12:44:23.937965Z","iopub.status.idle":"2025-11-10T12:44:23.954996Z","shell.execute_reply.started":"2025-11-10T12:44:23.937949Z","shell.execute_reply":"2025-11-10T12:44:23.954278Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"9d61fe5f","cell_type":"code","source":"def load_mask_npy(mask_path: str) -> np.ndarray:\n    \"\"\"\n    Load .npy mask which may be:\n      - (H, W): single binary mask\n      - (N, H, W): multiple instance masks (we OR them)\n      - list/tuple/dict: we try to extract arrays and OR them\n    Returns uint8 array {0,1}\n    \"\"\"\n    m = np.load(mask_path, allow_pickle=True)\n\n    # unpack common container types\n    if isinstance(m, (list, tuple)):\n        arrs = []\n        for item in m:\n            item = np.asarray(item)\n            if item.ndim == 2:\n                arrs.append(item.astype(np.uint8))\n            elif item.ndim == 3:\n                arrs.append((item.sum(axis=0) > 0).astype(np.uint8))\n        if len(arrs) == 0:\n            raise ValueError(f\"Unsupported mask list/tuple in {mask_path}\")\n        m = np.stack(arrs, axis=0)\n\n    if isinstance(m, dict):\n        # If your dict uses a different key, edit here\n        key = \"masks\" if \"masks\" in m else list(m.keys())[0]\n        m = np.asarray(m[key])\n\n    m = np.asarray(m)\n    if m.ndim == 2:\n        mask = (m > 0).astype(np.uint8)\n    elif m.ndim == 3:\n        mask = (m.sum(axis=0) > 0).astype(np.uint8)\n    else:\n        raise ValueError(f\"Unsupported mask ndim {m.ndim} for {mask_path}\")\n    return mask\n\n\ndef resize_pair(img_pil: Image.Image, mask_np: np.ndarray, size: int) -> tuple[Image.Image, Image.Image]:\n    \"\"\"\n    Resize image and mask to a square 'size' while keeping nearest for mask.\n    \"\"\"\n    img_resized  = img_pil.resize((size, size), resample=Image.BILINEAR)\n    mask_pil     = Image.fromarray(mask_np)\n    mask_resized = mask_pil.resize((size, size), resample=Image.NEAREST)\n    return img_resized, mask_resized\n\n\ndef random_flip_rotate(img: Image.Image, mask: Image.Image) -> tuple[Image.Image, Image.Image]:\n    \"\"\"\n    Light augmentations that are mask-safe:\n    - random horizontal flip\n    - random vertical flip\n    - small rotate [-10,10] degrees (expand=False), bilinear for img, nearest for mask\n    \"\"\"\n    import random\n    if random.random() < 0.5:\n        img  = TF.hflip(img)\n        mask = TF.hflip(mask)\n    if random.random() < 0.2:\n        img  = TF.vflip(img)\n        mask = TF.vflip(mask)\n    angle = random.uniform(-10, 10)\n    img  = TF.rotate(img, angle=angle, interpolation=InterpolationMode.BILINEAR, expand=False)\n    mask = TF.rotate(mask, angle=angle, interpolation=InterpolationMode.NEAREST,   expand=False)\n    return img, mask\n","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:44:23.958301Z","iopub.execute_input":"2025-11-10T12:44:23.958547Z","iopub.status.idle":"2025-11-10T12:44:23.982970Z","shell.execute_reply.started":"2025-11-10T12:44:23.958529Z","shell.execute_reply":"2025-11-10T12:44:23.982111Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"7f126181","cell_type":"code","source":"def build_items(base_dir: str):\n    \"\"\"\n    Build list of samples from a base directory that contains:\n      - authentic/ (images)\n      - forged/ (images)\n      - masks_npy/ (masks)\n    Returns:\n      - path: image path\n      - case_id: stem\n      - label: 1 if forged, 0 if authentic\n      - mask_path: path to .npy if forged else None\n    \"\"\"\n    items = []\n    # Define specific paths based on the new structure\n    mask_dir = f\"{base_dir}/masks_npy\" \n    \n    for cls in [\"authentic\", \"forged\"]:\n        img_dir = f\"{base_dir}/{cls}\" \n        if not os.path.exists(img_dir):\n            continue\n        for p in glob.glob(os.path.join(img_dir, \"*\")):\n            case_id = Path(p).stem\n            mask_path = None\n            if cls == \"forged\":\n                # The mask_dir is now defined inside this function\n                cand = os.path.join(mask_dir, f\"{case_id}.npy\") \n                mask_path = cand if os.path.exists(cand) else None\n            items.append({\n                \"path\": p,\n                \"case_id\": case_id,\n                \"label\": 1 if cls == \"forged\" else 0,\n                \"mask_path\": mask_path\n            })\n    return items\n\n\n# New name so it won't conflict with the older ForgeryDataset in memory\nfrom PIL import Image\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset\nfrom torchvision.transforms.functional import resize as tv_resize, InterpolationMode\n\ndef _to_single_channel_mask(mask_np: np.ndarray) -> torch.Tensor:\n    \"\"\"Convert any (H,W), (H,W,C), or (C,H,W) mask to (1,H,W) binary tensor.\"\"\"\n    if mask_np.dtype != np.float32:\n        mask_np = mask_np.astype(np.float32)\n\n    if mask_np.ndim == 2:\n        m = torch.from_numpy(mask_np).unsqueeze(0)        # (1,H,W)\n    elif mask_np.ndim == 3:\n        # handle (H,W,C) or (C,H,W)\n        if mask_np.shape[0] <= 4:\n            m = torch.from_numpy(mask_np)                 # (C,H,W)\n        else:\n            m = torch.from_numpy(mask_np).permute(2,0,1)  # (H,W,C)->(C,H,W)\n        if m.shape[0] > 1:\n            m = m.max(dim=0, keepdim=True)[0]            # collapse channels\n    else:\n        mask_np = np.squeeze(mask_np).astype(np.float32)\n        return _to_single_channel_mask(mask_np)\n\n    return (m > 0.5).float()                              # (1,H,W)\n\nclass ForgeryDatasetV2(Dataset):\n    def __init__(self, items, transforms=None, image_size=512, is_train=True):\n        \"\"\"\n        items: list of dicts with keys 'path' (or 'img'), optional 'mask'/'mask_path', and 'label'\n        transforms: optional callable applied to the image tensor (C,H,W)\n        image_size: int or (H,W) to resize image & mask\n        is_train: available if you want to branch later (not used here)\n        \"\"\"\n        self.items = items\n        self.transforms = transforms\n        self.is_train = is_train\n        if isinstance(image_size, int):\n            image_size = (image_size, image_size)\n        self.image_size = tuple(image_size)\n\n    def __len__(self):\n        return len(self.items)\n\n    def __getitem__(self, idx):\n        it = self.items[idx]\n        p = it.get(\"path\") or it.get(\"img\") or it.get(\"image\") or it.get(\"image_path\")\n        if p is None:\n            raise KeyError(f\"Missing image path at idx={idx}: {it}\")\n\n        # image -> tensor\n        img = Image.open(p).convert(\"RGB\")\n        img_np = np.array(img)                                  # (H,W,3)\n        img_t  = torch.from_numpy(img_np).permute(2,0,1).float() / 255.0  # (3,H,W)\n\n        # mask -> (1,H,W)\n        mpath = it.get(\"mask_path\") or it.get(\"mask\")\n        mask_np = None\n        if mpath is not None:\n            try:\n                mask_np = np.load(mpath)\n            except Exception as e:\n                print(f\"[WARN] Could not load mask at {mpath}: {e}. Using empty mask.\")\n        if mask_np is None:\n            h, w = img_np.shape[:2]\n            mask_np = np.zeros((h, w), dtype=np.float32)\n\n        mask_t = _to_single_channel_mask(mask_np)               # (1,H,W)\n\n        # resize both to common size\n        img_t  = tv_resize(img_t, size=self.image_size, interpolation=InterpolationMode.BILINEAR, antialias=True)\n        mask_t = F.interpolate(mask_t.unsqueeze(0), size=self.image_size, mode=\"nearest\").squeeze(0)  # (1,H,W)\n\n        if self.transforms is not None:\n            img_t = self.transforms(img_t)\n\n        label   = int(it.get(\"label\", 0))\n        case_id = p\n        return img_t, mask_t, case_id, label\n","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:44:23.983772Z","iopub.execute_input":"2025-11-10T12:44:23.984005Z","iopub.status.idle":"2025-11-10T12:44:24.002411Z","shell.execute_reply.started":"2025-11-10T12:44:23.983986Z","shell.execute_reply":"2025-11-10T12:44:24.001515Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"363ca0a2","cell_type":"code","source":"# --- Data augmentation wrapper (no deformation; flips + color only) ---\nfrom torch.utils.data import Dataset\nimport torchvision.transforms.functional as TF\nfrom PIL import Image\nimport numpy as np\nimport torch, random\n\nIMAGENET_MEAN = (0.485, 0.456, 0.406)\nIMAGENET_STD  = (0.229, 0.224, 0.225)\n\nclass AugmentedWrapper(Dataset):\n    \"\"\"Wrap a base Dataset that yields (img_t, mask_t, case_id, label)\n    and return an augmented view:\n      - Horizontal and/or vertical flips are applied to both image and mask.\n      - Brightness and hue adjustments are applied to the image only.\n      - No geometric resizing/cropping/warping is introduced here.\n    \"\"\"\n    def __init__(self, base_ds, \n                 hflip_p: float = 0.5, vflip_p: float = 0.5,\n                 brightness: tuple = (0.9, 1.1),  # multiplicative\n                 hue: tuple = (-0.03, 0.03)):     # fraction in [-0.5, 0.5]\n        self.base_ds   = base_ds\n        self.hflip_p   = hflip_p\n        self.vflip_p   = vflip_p\n        self.brightness= brightness\n        self.hue       = hue\n\n    def __len__(self):\n        return len(self.base_ds)\n\n    def __getitem__(self, idx):\n        img_t, mask_t, case_id, label = self.base_ds[idx]  # img_t: [3,H,W] float 0..1, mask_t: [1,H,W] float {0,1}\n        \n        # Convert to PIL for functional transforms that require PIL\n        img_pil  = TF.to_pil_image(img_t)  # maintains size\n        mask_np  = mask_t.squeeze(0).cpu().numpy().astype(np.uint8)\n        mask_pil = Image.fromarray(mask_np, mode=\"L\")      # single-channel mask\n        \n        # --- Flips ---\n        if random.random() < self.hflip_p:\n            img_pil  = TF.hflip(img_pil)\n            mask_pil = TF.hflip(mask_pil)\n        if random.random() < self.vflip_p:\n            img_pil  = TF.vflip(img_pil)\n            mask_pil = TF.vflip(mask_pil)\n\n        # --- Color (image only) ---\n        # Keep adjustments mild to respect copy–move appearance; no saturation/contrast changes per request.\n        b = random.uniform(*self.brightness)\n        img_pil = TF.adjust_brightness(img_pil, b)\n        h = random.uniform(*self.hue)\n        img_pil = TF.adjust_hue(img_pil, h)\n\n        # Back to tensors (preserve size)\n        img_out  = TF.to_tensor(img_pil).to(dtype=img_t.dtype)\n        mask_out = torch.from_numpy(np.array(mask_pil, dtype=np.float32)).unsqueeze(0)  # [1,H,W], {0,1}\n\n        img_np = np.array(img)                    # (H,W,3) uint8\n        img_t  = torch.from_numpy(img_np).permute(2,0,1).float() / 255.0  # (3,H,W)\n        \n        mask_t = torch.from_numpy(mask.astype(np.float32))                 # (H,W)\n        mask_t = mask_t.unsqueeze(0)                                       # (1,H,W)  <-- ADD THIS\n        \n        label  = int(it.get(\"label\", 0))\n        case_id = p\n        img_t = TF.normalize(img_t, IMAGENET_MEAN, IMAGENET_STD)\n        return img_t, mask_t, case_id, label","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:44:24.003368Z","iopub.execute_input":"2025-11-10T12:44:24.003747Z","iopub.status.idle":"2025-11-10T12:44:24.021814Z","shell.execute_reply.started":"2025-11-10T12:44:24.003725Z","shell.execute_reply":"2025-11-10T12:44:24.020952Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"00a36010","cell_type":"code","source":"# ===== BUILD 70/30 TRAIN/VAL FROM RECODAI =====\n# Gather images\nrecodai_auth_imgs = sorted(RECODAI_AUTH_DIR.glob(\"*.png\"))\nrecodai_forg_imgs = sorted(RECODAI_FORG_DIR.glob(\"*.png\"))\n\n# Gather masks (npy) and map by stem for quick lookup\nrecodai_masks_map = map_by_stem(sorted(RECODAI_MASK_DIR.glob(\"*.npy\")))\n\n# Build labeled list (image, label, mask_path or None)\nrecodai_all = []\n\n# Authentic: normally no mask; set to None (your dataset class can handle None or you can create zero masks later)\nfor p in recodai_auth_imgs:\n    recodai_all.append(to_item(p, LABEL_AUTHENTIC, None, src=\"recodai\"))\n\n# Forged: try to match mask by stem\nmissing_masks = 0\nfor p in recodai_forg_imgs:\n    m = recodai_masks_map.get(p.stem, None)\n    if m is None:\n        missing_masks += 1\n    recodai_all.append(to_item(p, LABEL_FORGED, m, src=\"recodai\"))\n\nif missing_masks:\n    print(f\"[WARN] RecoDAI forged images missing masks: {missing_masks} (will set mask=None for those)\")\n\n# Deterministic shuffle before split\nrandom.shuffle(recodai_all)\n\n# Stratified 70/30 split by label\n# We do this manually to avoid a hard dependency on sklearn.\nauth_pool = [it for it in recodai_all if it[\"label\"] == LABEL_AUTHENTIC]\nforg_pool = [it for it in recodai_all if it[\"label\"] == LABEL_FORGED]\n\ndef split_pool(pool, frac):\n    n_train = int(round(len(pool) * frac))\n    return pool[:n_train], pool[n_train:]\n\nauth_train, auth_val = split_pool(auth_pool, TRAIN_FRACTION)\nforg_train, forg_val = split_pool(forg_pool, TRAIN_FRACTION)\n\ntrain_items = auth_train + forg_train\nval_items   = auth_val   + forg_val\n\n# Shuffle each split for good measure\nrandom.shuffle(train_items)\nrandom.shuffle(val_items)\n\ndebug_counts(train_items, \"RecoDAI TRAIN\")\ndebug_counts(val_items,   \"RecoDAI VAL\")\n","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:44:24.022690Z","iopub.execute_input":"2025-11-10T12:44:24.022940Z","iopub.status.idle":"2025-11-10T12:44:24.137132Z","shell.execute_reply.started":"2025-11-10T12:44:24.022918Z","shell.execute_reply":"2025-11-10T12:44:24.136326Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"0c503200","cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nfrom typing import List, Tuple, Optional\nfrom scipy import ndimage\n\ndef load_mask_instances(mask_path: str) -> List[np.ndarray]:\n    \"\"\"\n    Load instance masks (list of [H,W] boolean arrays).\n    - If the .npy contains (N,H,W) -> split along axis 0\n    - If it's (H,W) -> split into connected components as separate instances\n    - If it's a container (list/tuple/dict), try to extract arrays and split accordingly\n    \"\"\"\n    m = np.load(mask_path, allow_pickle=True)\n\n    def _to_list_of_masks(arr: np.ndarray) -> List[np.ndarray]:\n        arr = np.asarray(arr)\n        if arr.ndim == 2:\n            # Connected components as instances\n            lab, n = ndimage.label(arr > 0)\n            inst = [(lab == k) for k in range(1, n + 1)]\n            return [x.astype(np.uint8) for x in inst] if n > 0 else []\n        elif arr.ndim == 3:\n            return [((arr[k] > 0).astype(np.uint8)) for k in range(arr.shape[0])]\n        else:\n            raise ValueError(f\"Unsupported array ndim {arr.ndim}\")\n\n    if isinstance(m, (list, tuple)):\n        out = []\n        for item in m:\n            out.extend(_to_list_of_masks(np.asarray(item)))\n        return out\n\n    if isinstance(m, dict):\n        key = \"masks\" if \"masks\" in m else list(m.keys())[0]\n        return _to_list_of_masks(np.asarray(m[key]))\n\n    # ndarray\n    return _to_list_of_masks(np.asarray(m))\n\n\ndef overlay_instances_on_image(\n    img_rgb: np.ndarray,\n    inst_masks: List[np.ndarray],\n    alpha: float = 0.35\n) -> np.ndarray:\n    \"\"\"\n    Draw semi-transparent colored overlays for each instance mask on top of the RGB image.\n    Returns an RGB uint8 array.\n    \"\"\"\n    out = img_rgb.astype(np.float32).copy()\n    H, W, _ = out.shape\n\n    # Generate a fixed set of random colors (to be consistent across runs, use seed)\n    rng = np.random.default_rng(1234)\n    colors = rng.integers(low=0, high=255, size=(max(1, len(inst_masks)), 3), dtype=np.uint8)\n\n    for i, m in enumerate(inst_masks):\n        if m.shape != (H, W):\n            # Resize instance to match (rare, but safety)\n            m = np.array(Image.fromarray((m > 0).astype(np.uint8)).resize((W, H), Image.NEAREST))\n        mask = (m > 0).astype(np.uint8)\n        color = colors[i % len(colors)].astype(np.float32)\n\n        # Blend: out = (1-alpha)*out + alpha*color on masked pixels\n        out[mask == 1] = (1 - alpha) * out[mask == 1] + alpha * color\n\n    return np.clip(out, 0, 255).astype(np.uint8)\n\n\ndef show_forged_example(\n    item: dict,\n    image_size: Optional[int] = None,\n    title: Optional[str] = None\n) -> None:\n    \"\"\"\n    Visualize one forged sample with instance overlays.\n    - item: dict from items list (must be forged and have mask_path)\n    - image_size: if provided, resize for display; otherwise show native\n    \"\"\"\n    assert item[\"mask_path\"] is not None, \"This helper expects a forged item with a mask_path.\"\n    img = Image.open(item[\"path\"]).convert(\"RGB\")\n    w, h = img.size\n\n    insts = load_mask_instances(item[\"mask_path\"])\n\n    if image_size is not None:\n        img = img.resize((image_size, image_size), Image.BILINEAR)\n        # Resize instances for display\n        insts_resized = []\n        for m in insts:\n            m = np.array(Image.fromarray((m > 0).astype(np.uint8)).resize((image_size, image_size), Image.NEAREST))\n            insts_resized.append(m)\n        insts = insts_resized\n\n    img_np = np.array(img)\n    over = overlay_instances_on_image(img_np, insts, alpha=0.38)\n\n    plt.figure(figsize=(10, 5))\n    plt.subplot(1, 2, 1)\n    plt.imshow(img_np)\n    plt.axis(\"off\")\n    plt.title(title or f\"Forged image — {Path(item['path']).name}\")\n\n    plt.subplot(1, 2, 2)\n    plt.imshow(over)\n    plt.axis(\"off\")\n    plt.title(f\"Overlay with {len(insts)} instance(s)\")\n    plt.show()\n\n\ndef show_forged_grid(\n    forged_list: List[dict],\n    n: int = 6,\n    image_size: int = 512,\n    cols: int = 3\n) -> None:\n    \"\"\"\n    Show a grid of forged examples with overlays.\n    \"\"\"\n    n = min(n, len(forged_list))\n    sel = random.sample(forged_list, n)\n\n    rows = (n + cols - 1) // cols\n    plt.figure(figsize=(5 * cols, 5 * rows))\n    for i, it in enumerate(sel, 1):\n        img = Image.open(it[\"path\"]).convert(\"RGB\").resize((image_size, image_size), Image.BILINEAR)\n        insts = load_mask_instances(it[\"mask_path\"])\n        insts = [np.array(Image.fromarray((m > 0).astype(np.uint8)).resize((image_size, image_size), Image.NEAREST)) for m in insts]\n        over  = overlay_instances_on_image(np.array(img), insts, alpha=0.38)\n\n        ax = plt.subplot(rows, cols, i)\n        ax.imshow(over)\n        ax.axis(\"off\")\n        ax.set_title(f\"{Path(it['path']).name}\\n{len(insts)} instance(s)\")\n    plt.tight_layout()\n    plt.show()","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-11-10T12:44:24.138089Z","iopub.execute_input":"2025-11-10T12:44:24.138363Z","iopub.status.idle":"2025-11-10T12:44:24.166978Z","shell.execute_reply.started":"2025-11-10T12:44:24.138342Z","shell.execute_reply":"2025-11-10T12:44:24.166091Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"2a32596b","cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n# DINOv2-based lightweight segmentation head; returns dict with \"out\" to be drop-in compatible.\ndef _load_dinov2(backbone='vitb14'):\n    hub_id = {\n        'vits14': 'dinov2_vits14',\n        'vitb14': 'dinov2_vitb14',\n        'vitl14': 'dinov2_vitl14',\n        'vitg14': 'dinov2_vitg14',\n    }[backbone]\n    m = torch.hub.load('facebookresearch/dinov2', hub_id)\n    m.eval()\n    return m\n\nclass DINOv2Seg(nn.Module):\n    def __init__(self, backbone='vitb14', out_ch=1):\n        super().__init__()\n        self.vit = _load_dinov2(backbone)\n        self.embed_dim = {'vits14':384,'vitb14':768,'vitl14':1024,'vitg14':1536}[backbone]\n        self.decode = nn.Sequential(\n            nn.Conv2d(self.embed_dim, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True),\n            nn.Conv2d(256, out_ch, 1)\n        )\n\n    def _tokens_from_vit(self, x):\n        \"\"\"\n        Return patch tokens as [B, N, D] (no CLS), regardless of the exact API.\n        Supports torch.hub dinov2 and timm variants.\n        \"\"\"\n        # 1) Preferred: forward_features with return_all_tokens (timm & some hub builds)\n        ff = getattr(self.vit, \"forward_features\", None)\n        if callable(ff):\n            try:\n                out = ff(x, return_all_tokens=True)  # timm ViT supports this\n                # timm sometimes returns a tensor [B, 1+N, D] or a dict\n                if isinstance(out, dict):\n                    if \"x_norm_patchtokens\" in out:         # dinov2-style dict\n                        tokens = out[\"x_norm_patchtokens\"]   # [B, N, D]\n                        return tokens\n                    if \"tokens\" in out:\n                        t = out[\"tokens\"]                    # [B, 1+N, D]\n                        return t[:, 1:, :]\n                elif out.dim() == 3:\n                    # [B, 1+N, D] or [B, N, D]\n                    return out[:, 1:, :] if out.size(1) > 0 and out.size(1) != int(out.size(1)**0.5)**2 else out\n            except TypeError:\n                # Some builds don't accept return_all_tokens\n                out = ff(x)\n                if isinstance(out, dict) and \"x_norm_patchtokens\" in out:\n                    return out[\"x_norm_patchtokens\"]\n\n        # 2) Fallback: get_intermediate_layers (available on many ViTs, incl. dinov2)\n        gil = getattr(self.vit, \"get_intermediate_layers\", None)\n        if callable(gil):\n            # returns list of tuples (cls, tokens) or tensors; take the last block\n            inter = gil(x, n=1, return_class_token=True)\n            last = inter[-1]\n            if isinstance(last, (list, tuple)) and len(last) == 2:\n                cls, tokens = last\n                return tokens                         # [B, N, D]\n            if torch.is_tensor(last) and last.dim() == 3:\n                # [B, 1+N, D]\n                return last[:, 1:, :]\n\n        # 3) Absolute fallback: plain forward → usually [B, D] pooled → not usable\n        y = self.vit(x)\n        if torch.is_tensor(y) and y.dim() == 3:\n            return y[:, 1:, :]\n        raise RuntimeError(\"Could not extract patch tokens from DINOv2 backbone; got shape \"\n                           f\"{tuple(y.shape) if torch.is_tensor(y) else type(y)}\")\n\n    def forward(self, x):\n        tokens = self._tokens_from_vit(x)     # [B, N, D]\n        B, N, D = tokens.shape\n        S = int(N ** 0.5)                     # assume square token grid\n        feat = tokens.transpose(1, 2).reshape(B, D, S, S)   # [B, D, S, S]\n        logits = self.decode(feat)\n        logits = F.interpolate(logits, size=x.shape[-2:], mode=\"bilinear\", align_corners=False)\n        return {\"out\": logits}\n\n\ndef build_model():\n    # Choose backbone: 'vits14' (small), 'vitb14' (base), 'vitl14', 'vitg14'\n    return DINOv2Seg(backbone='vitb14', out_ch=1)\n\nclass DiceLoss(nn.Module):\n    def __init__(self, eps=1e-6):\n        super().__init__()\n        self.eps = eps\n\n    def forward(self, logits, targets):\n        probs = torch.sigmoid(logits)\n        num = 2 * (probs * targets).sum() + self.eps\n        den = probs.sum() + targets.sum() + self.eps\n        return 1 - num / den\n\n\ndef bce_dice_loss(logits, targets, alpha=0.5, pos_weight=4.0):\n    if targets.dim() == 3:\n        targets = targets.unsqueeze(1)\n    bce  = torch.nn.functional.binary_cross_entropy_with_logits(\n        logits, targets, pos_weight=torch.tensor([pos_weight], device=logits.device)\n    )\n    dice = DiceLoss()(logits, targets)\n    return alpha * bce + (1 - alpha) * dice\n","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:44:24.168004Z","iopub.execute_input":"2025-11-10T12:44:24.168282Z","iopub.status.idle":"2025-11-10T12:44:24.189019Z","shell.execute_reply.started":"2025-11-10T12:44:24.168260Z","shell.execute_reply":"2025-11-10T12:44:24.188294Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"81e5f3e7","cell_type":"code","source":"@torch.no_grad()\ndef batch_f1(logits, targets, thr=0.5):\n    \"\"\"\n    Quick pixel F1 proxy (not the official image-wise oF1).\n    Good enough to monitor training.\n    \"\"\"\n    probs = torch.sigmoid(logits)\n    preds = (probs > thr).float()\n    tp = (preds * targets).sum().item()\n    fp = (preds * (1 - targets)).sum().item()\n    fn = ((1 - preds) * targets).sum().item()\n    precision = tp / (tp + fp + 1e-6)\n    recall    = tp / (tp + fn + 1e-6)\n    return 2 * precision * recall / (precision + recall + 1e-6)\n\n\ndef train_one_epoch(model, loader, optimizer, device, scaler=None, print_every=50):\n    model.train()\n    running_loss, running_f1 = 0.0, 0.0\n\n    for it, (imgs, masks, _, _) in enumerate(loader):\n        imgs, masks = imgs.to(device), masks.to(device)\n        if masks.dim() == 3:\n            masks = masks.unsqueeze(1)   # (B,1,H,W)\n\n        optimizer.zero_grad(set_to_none=True)\n        # ✅ Proper indentation + new AMP API\n        with torch.amp.autocast(\"cuda\", enabled=(scaler is not None)):\n            out = model(imgs)[\"out\"]                # DeepLab returns dict\n            loss = bce_dice_loss(out, masks)\n\n        if scaler is not None:\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n        else:\n            loss.backward()\n            optimizer.step()\n\n        running_loss += loss.item()\n        running_f1   += batch_f1(out.detach(), masks)\n\n        if (it + 1) % print_every == 0:\n            n = it + 1\n            print(f\"[train] it {n:04d} | loss {running_loss/n:.4f} | f1 {running_f1/n:.4f}\")\n\n    n = len(loader)\n    return running_loss / n, running_f1 / n\n\n\n@torch.no_grad()\ndef valid_one_epoch(model, loader, device, thr=0.5):\n    model.eval()\n    running_loss, running_f1 = 0.0, 0.0\n\n    for imgs, masks, _, _ in loader:\n        imgs, masks = imgs.to(device), masks.to(device)\n        if masks.dim() == 3:\n            masks = masks.unsqueeze(1)\n\n        out  = model(imgs)[\"out\"]\n        loss = bce_dice_loss(out, masks)\n        running_loss += loss.item()\n        running_f1   += batch_f1(out, masks, thr=thr)\n\n    n = len(loader)\n    return running_loss / n, running_f1 / n\n\n@torch.no_grad()\ndef find_best_threshold(model, loader, device, thr_grid=None):\n    model.eval()\n    if thr_grid is None:\n        thr_grid = np.linspace(0.05, 0.95, 19)\n    all_logits = []\n    all_masks  = []\n    for imgs, masks, _, _ in loader:\n        imgs = imgs.to(device)\n        out  = model(imgs)[\"out\"]  # (B,1,H,W)\n        all_logits.append(out.cpu())\n        all_masks.append(masks.cpu())\n    logits = torch.cat(all_logits, dim=0)\n    targets = torch.cat(all_masks, dim=0)\n    best_thr, best_f1 = 0.5, -1.0\n    for thr in thr_grid:\n        f1 = batch_f1(logits, targets, thr=float(thr))\n        if f1 > best_f1:\n            best_f1, best_thr = float(f1), float(thr)\n    print(f\"[val] best F1={best_f1:.4f} @ thr={best_thr:.2f}\")\n    return best_thr, best_f1","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:44:24.189701Z","iopub.execute_input":"2025-11-10T12:44:24.190171Z","iopub.status.idle":"2025-11-10T12:44:24.211679Z","shell.execute_reply.started":"2025-11-10T12:44:24.190152Z","shell.execute_reply":"2025-11-10T12:44:24.210853Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"5eca9402","cell_type":"code","source":"from tqdm.auto import tqdm\n\n# ===== ADD TRAIN-ONLY DATA =====\n\n# 1) Nuclei images: put ALL as Authentic in TRAIN\nnuclei_imgs = sorted(NUCLEI_ROOT.rglob(\"*.png\"))\nnuclei_items = []\nfor p in tqdm(nuclei_imgs, desc=\"Nuclei → Authentic\", unit=\"img\"):\n    nuclei_items.append(to_item(p, LABEL_AUTHENTIC, None, src=\"nuclei\"))\n\n# 2) Data-augment set: all PNGs as Forged + match NPys by stem if available\naugment_imgs = sorted(AUGMENT_IMG_ROOT.rglob(\"*.png\"))\naugment_masks_map = map_by_stem(sorted(AUGMENT_MASK_ROOT.rglob(\"*.npy\")))\n\naugment_items = []\naug_missing = 0\nfor p in tqdm(augment_imgs, desc=\"Augment → Forged (match masks)\", unit=\"img\"):\n    m = augment_masks_map.get(p.stem, None)\n    if m is None:\n        aug_missing += 1\n    augment_items.append(to_item(p, LABEL_FORGED, m, src=\"augment\"))\n\nif aug_missing:\n    print(f\"[WARN] Augment forged images missing masks: {aug_missing} (mask=None)\")\n\n# Append both groups ONLY to TRAIN (show progress as well)\nfor it in tqdm(nuclei_items, desc=\"Append nuclei to TRAIN\", unit=\"item\"):\n    train_items.append(it)\nfor it in tqdm(augment_items, desc=\"Append augment to TRAIN\", unit=\"item\"):\n    train_items.append(it)\n\n# Shuffle final train\nrandom.shuffle(train_items)\n\ndebug_counts(nuclei_items,  \"Added Nuclei (TRAIN-only)\")\ndebug_counts(augment_items, \"Added Augment (TRAIN-only)\")\ndebug_counts(train_items,   \"FINAL TRAIN\")\ndebug_counts(val_items,     \"FINAL VAL (unchanged)\")\n","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:44:24.212381Z","iopub.execute_input":"2025-11-10T12:44:24.212619Z","iopub.status.idle":"2025-11-10T12:45:58.218199Z","shell.execute_reply.started":"2025-11-10T12:44:24.212603Z","shell.execute_reply":"2025-11-10T12:45:58.217340Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"53183ee6","cell_type":"code","source":"# === FIX PATHS + RESCAN AUGMENT & NUCLEI, APPEND TO TRAIN ===\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\nimport random\n\n# 1) Correct roots\nNUCLEI_ROOT       = Path(\"/kaggle/input/nuclei-segmentation-in-microscope-cell-images/Nuclei/Nuclei\")\nAUGMENT_IMG_ROOT  = Path(\"/kaggle/input/data-augment\")\nAUGMENT_MASK_ROOT = Path(\"/kaggle/input/data-augment\")\n\n# 2) Find files\nnuclei_imgs   = sorted(NUCLEI_ROOT.rglob(\"*.png\"))\naugment_imgs  = sorted(AUGMENT_IMG_ROOT.rglob(\"*.png\"))\naugment_masks = sorted(AUGMENT_MASK_ROOT.rglob(\"*.npy\"))\naugment_masks_map = {p.stem: p for p in augment_masks}\n\nprint(f\"[scan] nuclei png:  {len(nuclei_imgs)}\")\nprint(f\"[scan] augment png: {len(augment_imgs)}\")\nprint(f\"[scan] augment npy: {len(augment_masks)}\")\n\n# 3) Build items for each source (ensure to_item writes 'path'/'mask_path')\ndef to_item(img_path, label, mask_path, src):\n    return {\n        \"path\": str(img_path),\n        \"img\": str(img_path),\n        \"mask_path\": (str(mask_path) if mask_path is not None else None),\n        \"mask\": (str(mask_path) if mask_path is not None else None),\n        \"label\": int(label),\n        \"src\": src,\n    }\n\nLABEL_AUTHENTIC = 0\nLABEL_FORGED    = 1\n\nnuclei_items = [to_item(p, LABEL_AUTHENTIC, None, src=\"nuclei\")\n                for p in tqdm(nuclei_imgs, desc=\"Build nuclei items\", unit=\"img\")]\n\naugment_items = []\naug_missing = 0\nfor p in tqdm(augment_imgs, desc=\"Build augment items (+mask match)\", unit=\"img\"):\n    m = augment_masks_map.get(p.stem, None)\n    if m is None:\n        aug_missing += 1\n    augment_items.append(to_item(p, LABEL_FORGED, m, src=\"augment\"))\nif aug_missing:\n    print(f\"[WARN] augment images missing masks: {aug_missing}\")\n\n# 4) De-duplicate before appending (avoid double-adding if you re-run)\nexisting = set()\nfor it in train_items:\n    existing.add(it.get(\"path\") or it.get(\"img\"))\n\nadded_nuclei = [it for it in nuclei_items  if it[\"path\"] not in existing]\nadded_aug    = [it for it in augment_items if it[\"path\"] not in existing]\n\ntrain_items.extend(added_nuclei)\nfor it in added_nuclei: existing.add(it[\"path\"])\ntrain_items.extend(added_aug)\nfor it in added_aug:    existing.add(it[\"path\"])\n\nrandom.shuffle(train_items)\n\nprint(f\"[append] nuclei added: {len(added_nuclei)} / {len(nuclei_items)}\")\nprint(f\"[append] augment added: {len(added_aug)} / {len(augment_items)}\")\n\n# 5) Normalize keys (so downstream code can rely on 'path'/'mask_path')\ndef _normalize_items(items):\n    fixed = []\n    for it in items:\n        img_path  = it.get(\"path\") or it.get(\"img\") or it.get(\"image\") or it.get(\"image_path\")\n        mask_path = it.get(\"mask_path\") or it.get(\"mask\")\n        fixed.append({\n            \"path\": str(img_path) if img_path is not None else None,\n            \"img\":  str(img_path) if img_path is not None else None,\n            \"mask\": (str(mask_path) if mask_path is not None else None),\n            \"mask_path\": (str(mask_path) if mask_path is not None else None),\n            \"label\": int(it.get(\"label\", 0)),\n            \"src\": it.get(\"src\", \"unknown\"),\n        })\n    return fixed\n\ntrain_items = _normalize_items(train_items)\n\n# 6) Show counts by src to verify we really have augment/nuclei now\nfrom collections import Counter\nc = Counter(it.get(\"src\",\"unknown\") for it in train_items)\nprint(\"Train by src:\", dict(c))\n\n# A quick peek at examples (optional)\nprint(\"Example augment item:\",\n      next((it for it in train_items if it[\"src\"]==\"augment\"), None))\nprint(\"Example nuclei item:\",\n      next((it for it in train_items if it[\"src\"]==\"nuclei\"), None))\n","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:45:58.219356Z","iopub.execute_input":"2025-11-10T12:45:58.219758Z","iopub.status.idle":"2025-11-10T12:46:05.875444Z","shell.execute_reply.started":"2025-11-10T12:45:58.219730Z","shell.execute_reply":"2025-11-10T12:46:05.874590Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"d36aff34","cell_type":"code","source":"IMAGENET_MEAN = (0.485, 0.456, 0.406)\nIMAGENET_STD  = (0.229, 0.224, 0.225)\n\ndef dinov2_norm(img_t):\n    return TF.normalize(img_t, IMAGENET_MEAN, IMAGENET_STD)\n\n# ===== Build train/valid datasets using REAL augment path + optional on-the-fly augs =====\nfrom torch.utils.data import ConcatDataset, DataLoader\n\n# ---- Robust source tagging (prefer existing 'src'; else detect by path) ----\nAUG_PATH_TOKEN    = \"/kaggle/input/data-augment\"\nNUCLEI_PATH_TOKEN = \"/kaggle/input/nuclei-segmentation-in-microscope-cell-images/nuclei/nuclei\"  # lowercase compare\nRECODAI_TOKEN     = \"/kaggle/input/recodai-luc-scientific-image-forgery-detection\"\n\ndef _src_of(it):\n    # 1) honor pre-set src if present\n    if \"src\" in it and it[\"src\"]:\n        return it[\"src\"]\n    # 2) detect by path\n    p = (it.get(\"path\") or it.get(\"img\") or \"\").lower()\n    if AUG_PATH_TOKEN in p:\n        return \"augment\"\n    if NUCLEI_PATH_TOKEN in p:\n        return \"nuclei\"\n    if RECODAI_TOKEN in p:\n        return \"recodai\"\n    return \"unknown\"\n\n# ---- Split training items by source ----\ntrain_by_src = {\"augment\": [], \"nuclei\": [], \"recodai\": [], \"unknown\": []}\nfor it in train_items:\n    train_by_src[_src_of(it)].append(it)\n\nprint(\"Train split by src:\")\nfor k in [\"recodai\", \"nuclei\", \"augment\", \"unknown\"]:\n    print(f\"  {k:8s}: {len(train_by_src[k])}\")\n\n# ---- Build datasets ----\n# Base = real data (no pre-generated augmentations)\n_base_pool = train_by_src[\"recodai\"] + train_by_src[\"nuclei\"] + train_by_src[\"unknown\"]\n_base_train_ds = ForgeryDatasetV2(\n    _base_pool,\n    image_size=IMAGE_SIZE,\n    is_train=False  # deterministic branch (no random augs)\n)\n\n# Augment-from-disk = ONLY the files under /kaggle/input/data-augment\n_aug_disk_pool = train_by_src[\"augment\"]\nif len(_aug_disk_pool) == 0:\n    print(\"[WARN] No augment (disk) items found. Did the 'ADD TRAIN-ONLY DATA' cell run with the corrected paths?\")\n_aug_disk_ds = ForgeryDatasetV2(\n    _aug_disk_pool,\n    image_size=IMAGE_SIZE,\n    is_train=False  # these are already augmented renders\n)\n\n# (Optional) On-the-fly augmentation branch (stochastic)\n# NOTE: Make sure your AugmentedWrapper accepts **tensor** images; adapt if it expects PIL.\nDO_ON_THE_FLY = True\nif DO_ON_THE_FLY:\n    _aug_onthefly_ds = AugmentedWrapper(\n        ForgeryDatasetV2(_base_pool, image_size=IMAGE_SIZE, is_train=True)\n    )\n    train_ds = ConcatDataset([_base_train_ds, _aug_disk_ds, _aug_onthefly_ds])\nelse:\n    train_ds = ConcatDataset([_base_train_ds, _aug_disk_ds])\n\nprint(\"Dataset sizes:\")\nprint(\"  base:            \", len(_base_train_ds))\nprint(\"  augment (disk):  \", len(_aug_disk_ds))\nif DO_ON_THE_FLY:\n    print(\"  augment (OTF):   \", len(_aug_onthefly_ds))\nprint(\"  TOTAL train:     \", len(train_ds))\n\n# ---- Build loaders ----\nvalid_ds = ForgeryDatasetV2(val_items, image_size=IMAGE_SIZE, is_train=False, transforms=dinov2_norm)\nvalid_loader = DataLoader(valid_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True)\n\ntrain_loader = DataLoader(\n    train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True, drop_last=True\n)","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:46:05.876496Z","iopub.execute_input":"2025-11-10T12:46:05.876878Z","iopub.status.idle":"2025-11-10T12:46:05.897319Z","shell.execute_reply.started":"2025-11-10T12:46:05.876859Z","shell.execute_reply":"2025-11-10T12:46:05.896517Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"5f5aa395","cell_type":"code","source":"# immediately before your training loop\nprint(\"len(train_ds) =\", len(train_ds))\nprint(\"len(valid_ds) =\", len(valid_ds))","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:46:05.898262Z","iopub.execute_input":"2025-11-10T12:46:05.898641Z","iopub.status.idle":"2025-11-10T12:46:05.915336Z","shell.execute_reply.started":"2025-11-10T12:46:05.898612Z","shell.execute_reply":"2025-11-10T12:46:05.914516Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"ec029fb0","cell_type":"code","source":"train_dataset = ForgeryDatasetV2(\n    train_items,\n    transforms=train_transforms if 'train_transforms' in globals() else None,\n    image_size=IMAGE_SIZE if 'IMAGE_SIZE' in globals() else (512,512),\n    is_train=True\n)\n\nvalid_dataset = ForgeryDatasetV2(\n    val_items,\n    transforms=valid_transforms if 'valid_transforms' in globals() else None,\n    image_size=IMAGE_SIZE if 'IMAGE_SIZE' in globals() else (512,512),\n    is_train=False\n)\n\ntrain_loader = torch.utils.data.DataLoader(\n    train_dataset,\n    batch_size=BATCH_SIZE,\n    shuffle=True,\n    num_workers=2,\n    pin_memory=True,\n    drop_last=True\n)\n\nvalid_loader = torch.utils.data.DataLoader(\n    valid_dataset,\n    batch_size=VAL_BATCH_SIZE if 'VAL_BATCH_SIZE' in globals() else BATCH_SIZE,\n    shuffle=False,\n    num_workers=2,\n    pin_memory=True\n)\n","metadata":{"execution":{"iopub.status.busy":"2025-11-10T12:46:05.916069Z","iopub.execute_input":"2025-11-10T12:46:05.916315Z","iopub.status.idle":"2025-11-10T12:46:05.931166Z","shell.execute_reply.started":"2025-11-10T12:46:05.916299Z","shell.execute_reply":"2025-11-10T12:46:05.930389Z"},"trusted":true},"outputs":[],"execution_count":null},{"id":"2af4024a","cell_type":"code","source":"\n# === VALIDATION HELPERS (ROBUST) — REPLACE ===\nimport torch\nimport torch.nn.functional as F\nimport numpy as np\nfrom sklearn.metrics import f1_score\n\ndef _batch_to_tensors(batch, device):\n    \"\"\"\n    Returns: imgs [B,3,H,W], labels [B] (0/1), optional masks [B,1,H,W] or None.\n    Tries to infer from common batch shapes/keys.\n    \"\"\"\n    imgs = None\n    labels = None\n    masks = None\n\n    # Case A: dict-like batch\n    if isinstance(batch, dict):\n        # images\n        if \"image\" in batch:\n            imgs = batch[\"image\"]\n        elif \"images\" in batch:\n            imgs = batch[\"images\"]\n        else:\n            raise KeyError(\"Could not find images in batch (expected keys: 'image' or 'images').\")\n\n        # masks (optional)\n        if \"mask\" in batch:\n            masks = batch[\"mask\"]\n        elif \"masks\" in batch:\n            masks = batch[\"masks\"]\n\n        # image-level labels (optional)\n        if \"image_label\" in batch:\n            labels = batch[\"image_label\"]\n        elif \"labels\" in batch:\n            labels = batch[\"labels\"]\n\n    # Case B: tuple/list batch\n    elif isinstance(batch, (list, tuple)):\n        # Common patterns:\n        #  (images,) | (images, masks) | (images, labels) | (images, masks, labels)\n        if len(batch) == 1:\n            imgs = batch[0]\n        elif len(batch) == 2:\n            imgs, second = batch\n            if torch.is_tensor(second):\n                if second.ndim >= 3:\n                    masks = second\n                else:\n                    labels = second\n            else:\n                labels = second\n        elif len(batch) >= 3:\n            imgs, masks, labels = batch[0], batch[1], batch[2]\n        else:\n            raise TypeError(f\"Unsupported batch structure with len={len(batch)}\")\n    else:\n        raise TypeError(f\"Unsupported batch type: {type(batch)}\")\n\n    imgs = imgs.to(device)\n\n    if masks is not None:\n        if masks.ndim == 3:\n            masks = masks.unsqueeze(1)\n        masks = masks.to(device).float()\n\n    if labels is None:\n        if masks is None:\n            raise ValueError(\"No image-level labels found and no masks to derive them from.\")\n        labels = (masks.view(masks.size(0), -1).max(dim=1).values > 0.5).long()\n\n    labels = labels.detach().cpu().numpy().astype(np.uint8)\n    return imgs, labels, masks\n\n@torch.no_grad()\ndef collect_val_probs(model, val_loader, device):\n    \"\"\"\n    Returns:\n      prob_masks: FloatTensor [N,1,IMAGE_SIZE,IMAGE_SIZE] in [0,1]\n      labels:     np.array [N] (0/1)\n    \"\"\"\n    model.eval()\n    all_probs = []\n    all_labels = []\n\n    for batch in val_loader:\n        imgs, labels, _ = _batch_to_tensors(batch, device)\n\n        logits = model(imgs)                       # expect [B,1,h',w'] (segmentation)\n        if isinstance(logits, dict) and \"logits\" in logits:\n            logits = logits[\"logits\"]\n        if logits.ndim == 3:                       # [B,h',w'] -> [B,1,h',w']\n            logits = logits.unsqueeze(1)\n\n        probs = torch.sigmoid(logits)\n\n        if probs.shape[-2:] != (IMAGE_SIZE, IMAGE_SIZE):\n            probs = F.interpolate(probs, size=(IMAGE_SIZE, IMAGE_SIZE), mode='bilinear', align_corners=False)\n\n        all_probs.append(probs.cpu())\n        all_labels.append(labels)\n\n    prob_masks = torch.cat(all_probs, dim=0)\n    labels = np.concatenate(all_labels, axis=0)\n    return prob_masks, labels\n\ndef image_level_preds(prob_masks, pixel_thr, area_frac):\n    with torch.no_grad():\n        bin_masks = (prob_masks >= pixel_thr).float()\n        if USE_MORPH:\n            outs = []\n            for i in range(bin_masks.shape[0]):\n                outs.append(maybe_morphology(bin_masks[i:i+1]))\n            bin_masks = torch.stack(outs, dim=0)\n        area = bin_masks.mean(dim=[1,2,3]).cpu().numpy()\n    return (area >= area_frac).astype(np.uint8)\n\ndef image_f1(prob_masks, labels, pixel_thr, area_frac):\n    preds = image_level_preds(prob_masks, pixel_thr, area_frac)\n    return f1_score(labels, preds)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-10T12:46:45.588415Z","iopub.execute_input":"2025-11-10T12:46:45.588751Z","iopub.status.idle":"2025-11-10T12:46:45.605583Z","shell.execute_reply.started":"2025-11-10T12:46:45.588729Z","shell.execute_reply":"2025-11-10T12:46:45.604733Z"}},"outputs":[],"execution_count":null},{"id":"a4fc8a56","cell_type":"code","source":"\n# === GRID SEARCH (NEW) ===\nimport numpy as np\n\ndef grid_search_f1(prob_masks, labels, thr_values=None, area_values=None):\n    if thr_values is None:\n        thr_values = np.linspace(0.20, 0.80, 13)      # 0.20..0.80 step ≈0.05\n    if area_values is None:\n        area_values = np.geomspace(5e-4, 2e-2, 10)    # 0.05% .. 2% of pixels\n\n    best = {\"f1\": -1.0, \"thr\": None, \"area\": None}\n    for t in thr_values:\n        for a in area_values:\n            f1 = image_f1(prob_masks, labels, t, a)\n            if f1 > best[\"f1\"]:\n                best.update({\"f1\": float(f1), \"thr\": float(t), \"area\": float(a)})\n    print(f\"[GRID] best F1={best['f1']:.4f} at thr={best['thr']:.3f}, area={best['area']:.5f}\")\n    return best\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-10T12:46:48.993581Z","iopub.execute_input":"2025-11-10T12:46:48.994133Z","iopub.status.idle":"2025-11-10T12:46:49.000368Z","shell.execute_reply.started":"2025-11-10T12:46:48.994108Z","shell.execute_reply":"2025-11-10T12:46:48.999548Z"}},"outputs":[],"execution_count":null},{"id":"8ae57b58","cell_type":"code","source":"# === TRAIN/VALIDATE & SAVE-BEST (FIXED: no scheduler kwarg) ===\nmodel = build_model().to(DEVICE)\n\n# Optimizer & optional scheduler\noptimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n\nUSE_SCHEDULER = True  # set False if you don't want LR scheduling\nif USE_SCHEDULER:\n    from torch.optim.lr_scheduler import CosineAnnealingLR\n    scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=1e-6)\n\nbest_score = -1.0\nnum_epochs = 2  # change as you like\n\nfor epoch in range(num_epochs):\n    print(f\"\\nEpoch {epoch+1}/{num_epochs}\")\n    # ---- TRAIN ----\n    model.train()\n    # IMPORTANT: do NOT pass 'scheduler=' here, since your train_one_epoch doesn't accept it\n    train_loss = train_one_epoch(model, train_loader, optimizer, DEVICE)\n    print(f\"train_loss: {train_loss:.4f}\")\n\n    # ---- VALIDATE & CALIBRATE ----\n    val_probs, val_labels = collect_val_probs(model, valid_loader, DEVICE)\n    search = grid_search_f1(val_probs, val_labels)\n    val_score = search['f1']\n    print(f\"[VAL] F1={val_score:.4f} at thr={search['thr']:.3f}, area={search['area']:.5f}\")\n\n    preds = image_level_preds(val_probs, search['thr'], search['area'])\n    print(f\"[VAL] % predicted forged: {100*preds.mean():.2f}%  (GT {100*val_labels.mean():.2f}%)\")\n\n    # ---- SAVE IF BEST (by leaderboard metric) ----\n    if val_score > best_score:\n        best_score = val_score\n        to_save = {\n            \"state_dict\": model.state_dict(),\n            \"image_size\": IMAGE_SIZE,\n            \"mean\": MEAN, \"std\": STD,\n            \"use_morph\": USE_MORPH, \"morph_kernel\": MORPH_KERNEL,\n            \"pixel_thr\": search['thr'],\n            \"area_frac\": search['area'],\n            \"val_f1\": best_score,\n            \"arch\": getattr(model, \"name\", \"model\"),\n            \"epoch\": epoch,\n        }\n        torch.save(to_save, CKPT_PATH)\n        print(f\"✅ Saved new best to {CKPT_PATH} | F1={best_score:.4f}\")\n\n    # ---- STEP SCHEDULER (outside train_one_epoch) ----\n    if USE_SCHEDULER:\n        scheduler.step()\n        # (optional) show current LR\n        if hasattr(optimizer, \"param_groups\"):\n            print(\"lr:\", optimizer.param_groups[0][\"lr\"])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-10T12:50:07.500076Z","iopub.execute_input":"2025-11-10T12:50:07.500672Z"}},"outputs":[],"execution_count":null},{"id":"d12b36c4","cell_type":"code","source":"\n# === SANITY CHECK (NEW) ===\n# Load and evaluate on validation with saved thresholds\nckpt = torch.load(CKPT_PATH, map_location=\"cpu\")\n\n# Rebuild/Load model\n# Expect a build_model() function defined earlier in the notebook\ntry:\n    model = build_model()\nexcept NameError as e:\n    raise RuntimeError(\"build_model() must be defined before running sanity check.\") from e\n\nmodel.load_state_dict(ckpt[\"state_dict\"])\nmodel.to(DEVICE).eval()\n\n# Restore preprocessing/thresholds\nIMAGE_SIZE = int(ckpt[\"image_size\"])\nMEAN, STD = tuple(ckpt[\"mean\"]), tuple(ckpt[\"std\"])\nUSE_MORPH = bool(ckpt[\"use_morph\"])\nMORPH_KERNEL = int(ckpt[\"morph_kernel\"])\nPIXEL_THR = float(ckpt[\"pixel_thr\"])\nAREA_FRAC = float(ckpt[\"area_frac\"])\n\nimg_transform = T.Compose([\n    T.Resize((IMAGE_SIZE, IMAGE_SIZE), interpolation=T.InterpolationMode.BICUBIC),\n    T.ToTensor(),\n    T.Normalize(MEAN, STD),\n])\n\nwith torch.no_grad():\n    val_probs, val_labels = collect_val_probs(model, valid_loader, DEVICE)\n    preds = image_level_preds(val_probs, PIXEL_THR, AREA_FRAC)\n    print(f\"Sanity — GT forged rate:   {100*val_labels.mean():.2f}%\")\n    print(f\"Sanity — Pred forged rate: {100*preds.mean():.2f}%\")\n    print(f\"Sanity — Val F1 (fixed):   {image_f1(val_probs, val_labels, PIXEL_THR, AREA_FRAC):.4f}\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"ffd26a3f","cell_type":"code","source":"\n# === INFERENCE / TEST (REPLACE) ===\n# Load the best checkpoint and rebuild preprocessing exactly.\n\nckpt = torch.load(CKPT_PATH, map_location=\"cpu\")\n# Expect a build_model() defined earlier\nmodel = build_model()\nmodel.load_state_dict(ckpt[\"state_dict\"])\nmodel.to(DEVICE).eval()\n\nIMAGE_SIZE = ckpt[\"image_size\"]\nMEAN, STD = tuple(ckpt[\"mean\"]), tuple(ckpt[\"std\"])\nUSE_MORPH = bool(ckpt[\"use_morph\"])\nMORPH_KERNEL = int(ckpt[\"morph_kernel\"])\nPIXEL_THR = float(ckpt[\"pixel_thr\"])\nAREA_FRAC = float(ckpt[\"area_frac\"])\n\n# rebuild transforms from saved params (identical to train/val)\nimg_transform = T.Compose([\n    T.Resize((IMAGE_SIZE, IMAGE_SIZE), interpolation=T.InterpolationMode.BICUBIC),\n    T.ToTensor(),\n    T.Normalize(MEAN, STD),\n])\n\n@torch.no_grad()\ndef infer_batch(imgs):  # imgs: tensor [B,3,H,W] already transformed with img_transform\n    logits = model(imgs.to(DEVICE))  # [B,1,h',w']\n    if isinstance(logits, dict) and \"logits\" in logits:\n        logits = logits[\"logits\"]\n    if logits.ndim == 3:\n        logits = logits.unsqueeze(1)\n    probs = torch.sigmoid(logits)\n\n    if probs.shape[-2:] != (IMAGE_SIZE, IMAGE_SIZE):\n        probs = F.interpolate(probs, size=(IMAGE_SIZE, IMAGE_SIZE), mode='bilinear', align_corners=False)\n\n    bin_masks = (probs >= PIXEL_THR).float()\n    if USE_MORPH:\n        outs = []\n        for i in range(bin_masks.shape[0]):\n            outs.append(maybe_morphology(bin_masks[i:i+1]))\n        bin_masks = torch.stack(outs, dim=0)\n\n    # image-level decision: forged if area >= AREA_FRAC\n    area = bin_masks.mean(dim=[1,2,3])  # [B]\n    y_pred = (area >= AREA_FRAC).long()  # 1=forged, 0=authentic\n    return y_pred.cpu().numpy(), probs.cpu()  # (labels, prob masks)\n\n# Example loop over your test loader:\nimage_level_predictions = []\nfor batch in test_loader:\n    if isinstance(batch, dict):\n        imgs = batch[\"image\"]\n    elif isinstance(batch, (list, tuple)):\n        imgs = batch[0]\n    else:\n        raise TypeError(\"Unsupported test batch type; expected dict or tuple/list.\")\n    y_pred, _ = infer_batch(imgs)\n    image_level_predictions.extend(y_pred.tolist())\n\n# TODO: Format & write your submission from 'image_level_predictions'\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"ca3da0c2","cell_type":"code","source":"\n# === DEBUG BATCH STRUCTURE (OPTIONAL) ===\nb = next(iter(valid_loader))\nprint(type(b))\nif isinstance(b, dict):\n    print(\"dict keys:\", b.keys())\nelif isinstance(b, (list, tuple)):\n    print(\"tuple/list length:\", len(b))\n    for i, x in enumerate(b):\n        if torch.is_tensor(x):\n            print(f\"  idx {i}: tensor shape {tuple(x.shape)}, dtype {x.dtype}\")\n        else:\n            print(f\"  idx {i}: type {type(x)}\")\nelse:\n    print(\"unknown batch type\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"8dea8205-c341-4c02-9607-3356ba12659f","cell_type":"code","source":"# === EXPORT MODEL (NEW) ===\n# Save the fully trained & calibrated model checkpoint to share or reuse\n\nexport_path = \"dinov2_forgery_detector_final.pth\"\n\n# Load the best calibrated checkpoint (from training)\nckpt = torch.load(CKPT_PATH, map_location=\"cpu\")\n\n# If not already loaded in memory, rebuild and load model weights\ntry:\n    model\nexcept NameError:\n    model = build_model()\n    model.load_state_dict(ckpt[\"state_dict\"])\n    model.to(DEVICE).eval()\n\n# Create export dictionary (same structure as training save)\nexport_ckpt = {\n    \"arch\": getattr(model, \"name\", \"dinov2_forgery_detector\"),\n    \"state_dict\": model.state_dict(),\n    \"image_size\": int(ckpt.get(\"image_size\", IMAGE_SIZE)),\n    \"mean\": tuple(ckpt.get(\"mean\", MEAN)),\n    \"std\": tuple(ckpt.get(\"std\", STD)),\n    \"use_morph\": bool(ckpt.get(\"use_morph\", USE_MORPH)),\n    \"morph_kernel\": int(ckpt.get(\"morph_kernel\", MORPH_KERNEL)),\n    \"pixel_thr\": float(ckpt.get(\"pixel_thr\", PIXEL_THR)),\n    \"area_frac\": float(ckpt.get(\"area_frac\", AREA_FRAC)),\n    \"val_f1\": float(ckpt.get(\"val_f1\", 0.0)),\n}\n\n# Save export\ntorch.save(export_ckpt, export_path)\nprint(f\"✅ Model exported to {export_path}\")\n\n# Optional: check file size\nimport os\nprint(f\"File size: {os.path.getsize(export_path)/1e6:.2f} MB\")\n\n# Optional: verify reloading works\ntest_load = torch.load(export_path, map_location=\"cpu\")\nprint(\"Reload OK — keys:\", list(test_load.keys()))\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"id":"3aa51ab2","cell_type":"code","source":"import torch\nimport torch.nn as nn\n\n# 1) Reload best checkpoint cleanly ow\nexport_model = build_model().to(DEVICE)\nstate = torch.load(f\"{OUT_DIR}/dinov2_vitb14_linear_fold{FOLD_TO_RUN}.pt\", map_location=DEVICE)\n\nexport_model.load_state_dict(state[\"state_dict\"], strict=True)\nexport_model.eval()\n\n# 2) Wrap the model so forward returns only the logits tensor (works for dict OR tensor)\nclass LogitsOut(nn.Module):\n    def __init__(self, m: nn.Module):\n        super().__init__()\n        self.m = m.eval()\n        for p in self.m.parameters():\n            p.requires_grad_(False)\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        y = self.m(x)\n        if isinstance(y, dict) and \"out\" in y:\n            return y[\"out\"]            # torchvision seg + our DINOv2Seg return {\"out\": logits}\n        if torch.is_tensor(y):\n            return y                   # if a plain tensor is returned\n        raise RuntimeError(f\"Unexpected model output type: {type(y)}\")\n\nwrapper = LogitsOut(export_model).to(DEVICE)\n\n# 3) Script first (preferred); if that fails, fall back to trace\nexample = torch.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE, device=DEVICE)\nts_path = f\"{OUT_DIR}/model_fold{FOLD_TO_RUN}_ts.pt\"  # generic name\n\ntry:\n    with torch.no_grad():\n        scripted = torch.jit.script(wrapper)\n        _ = scripted(example)\n        scripted.save(ts_path)\n        print(\"Saved TorchScript (script):\", ts_path)\nexcept Exception as e:\n    print(\"Script failed, falling back to trace. Reason:\", e)\n    with torch.no_grad():\n        traced = torch.jit.trace(wrapper, example, strict=False)\n        _ = traced(example)\n        traced.save(ts_path)\n        print(\"Saved TorchScript (trace):\", ts_path)\n\n# 4) Optional ONNX export (trace-friendly; ViTs may need trace to succeed)\nonnx_path = f\"{OUT_DIR}/model_fold{FOLD_TO_RUN}.onnx\"\ntry:\n    import onnx  # ensures onnx is installed\n    with torch.no_grad():\n        torch.onnx.export(\n            wrapper,\n            example,\n            onnx_path,\n            input_names=[\"input\"],\n            output_names=[\"logits\"],\n            opset_version=17,  # 17 or 18; try 18 if 17 complains\n            do_constant_folding=True,\n            dynamic_axes={\"input\": {0: \"batch\"}, \"logits\": {0: \"batch\"}},\n        )\n    print(\"Saved ONNX:\", onnx_path)\nexcept Exception as e:\n    print(\"ONNX export skipped:\", e)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}