{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","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":4534,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":3326,"modelId":986}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\n# -------------------------\n# Cell 1: Imports & Setup\n# -------------------------\nimport os, gc, json, math, random\nfrom pathlib import Path\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\nfrom sklearn.model_selection import train_test_split\nfrom tqdm import tqdm\n\n# NEW: Augmentations\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-08T04:50:27.872323Z","iopub.execute_input":"2025-11-08T04:50:27.872500Z","iopub.status.idle":"2025-11-08T04:51:05.803947Z","shell.execute_reply.started":"2025-11-08T04:50:27.872484Z","shell.execute_reply":"2025-11-08T04:51:05.803107Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# -------------------------\n# Cell 2: Config & Hyperparameters\n# -------------------------\n\nclass CFG:\n    # Seed\n    SEED = 42\n\n    # Paths\n    BASE_DIR  = \"/kaggle/input/recodai-luc-scientific-image-forgery-detection\"\n    AUTH_DIR  = f\"{BASE_DIR}/train_images/authentic\"\n    FORG_DIR  = f\"{BASE_DIR}/train_images/forged\"\n    MASK_DIR  = f\"{BASE_DIR}/train_masks\"\n    TEST_DIR  = f\"{BASE_DIR}/test_images\"\n    SAMPLE_SUB = f\"{BASE_DIR}/sample_submission.csv\"\n    \n    # DINOv2 Path\n    DINO_PATH = \"/kaggle/input/dinov2/pytorch/base/1\"\n    \n    # Model Paths (to save best models)\n    BEST_SEG_MODEL = \"best_seg_model.pth\"\n    BEST_CLS_MODEL = \"best_cls_model.pth\"\n\n    # Params\n    IMG_SIZE = 400  # <-- Increased for more detail\n    BATCH_SEG = 8\n    BATCH_CLS = 16\n    \n    # Training Epochs\n    EPOCHS_SEG = 6  # <-- Increased\n    EPOCHS_CLS = 6  # <-- Increased\n    \n    # Learning Rates & Optimizer\n    LR_SEG = 1e-4\n    LR_CLS = 1e-3\n    WEIGHT_DECAY = 1e-4\n    SCHEDULER_T_MAX = 5 # T_max for CosineAnnealingLR (should match epochs)\n\n# Repro & device\ndef seed_everything(s=CFG.SEED):\n    random.seed(s)\n    np.random.seed(s)\n    torch.manual_seed(s)\n    torch.cuda.manual_seed_all(s)\n    os.environ['PYTHONHASHSEED'] = str(s)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False # Can be True for speed if inputs constant\n\nseed_everything(CFG.SEED)\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-08T04:52:39.272863Z","iopub.execute_input":"2025-11-08T04:52:39.273456Z","iopub.status.idle":"2025-11-08T04:52:39.281342Z","shell.execute_reply.started":"2025-11-08T04:52:39.273430Z","shell.execute_reply":"2025-11-08T04:52:39.280678Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# -------------------------\n# Cell 3: Utils (Metrics, RLE, Loss)\n# -------------------------\n\ndef binarize(x, thr=0.5):\n    return (x > thr).astype(np.uint8)\n\ndef iou_score(pred, gt, eps=1e-7):\n    p = binarize(pred); g = binarize(gt)\n    inter = (p & g).sum()\n    union = (p | g).sum()\n    return float(inter) / (float(union) + eps)\n\ndef dice_score(pred, gt, eps=1e-7):\n    p = binarize(pred); g = binarize(gt)\n    inter = (p & g).sum()\n    return float(2*inter) / (p.sum() + g.sum() + eps)\n\ndef pixel_acc(pred, gt, thr=0.5):\n    p = binarize(pred, thr); g = binarize(gt, thr)\n    return float((p == g).sum()) / float(np.prod(g.shape))\n\ndef rle_encode_numpy(mask):\n    pixels = mask.T.flatten()\n    dots = np.where(pixels == 1)[0]\n    if len(dots) == 0:\n        return \"[]\"\n    run_lengths, prev = [], -2\n    for b in dots:\n        if b > prev + 1:\n            run_lengths.extend((b + 1, 0))\n        run_lengths[-1] += 1\n        prev = b\n    return json.dumps(run_lengths)\n\n# --- NEW: Dice Loss ---\nclass DiceLoss(nn.Module):\n    def __init__(self, smooth=1.0):\n        super(DiceLoss, self).__init__()\n        self.smooth = smooth\n\n    def forward(self, logits, targets):\n        probs = torch.sigmoid(logits)\n        probs_flat = probs.view(-1)\n        targets_flat = targets.view(-1)\n        \n        intersection = (probs_flat * targets_flat).sum()\n        dice = (2. * intersection + self.smooth) / (probs_flat.sum() + targets_flat.sum() + self.smooth)\n        return 1. - dice\n\n# --- NEW: Combo Loss (BCE + Dice) ---\nclass ComboLoss(nn.Module):\n    def __init__(self, bce_weight=0.5, dice_weight=0.5):\n        super(ComboLoss, self).__init__()\n        self.bce_weight = bce_weight\n        self.dice_weight = dice_weight\n        self.bce_loss = nn.BCEWithLogitsLoss()\n        self.dice_loss = DiceLoss()\n\n    def forward(self, logits, targets):\n        bce = self.bce_loss(logits, targets)\n        dice = self.dice_loss(logits, targets)\n        return (self.bce_weight * bce) + (self.dice_weight * dice)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-08T04:52:40.421362Z","iopub.execute_input":"2025-11-08T04:52:40.421940Z","iopub.status.idle":"2025-11-08T04:52:40.431879Z","shell.execute_reply.started":"2025-11-08T04:52:40.421918Z","shell.execute_reply":"2025-11-08T04:52:40.431183Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -------------------------\n# Cell 4: Augmentations (NEW)\n# -------------------------\n\n# Must be careful with augs. Flips and rotations are safe.\n# Avoid crops, shears, or elastic transforms.\ndef get_transforms(img_size):\n    return {\n        \"train\": A.Compose([\n            A.Resize(img_size, img_size),\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.RandomRotate90(p=0.5),\n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # ImageNet stats\n            ToTensorV2(),\n        ]),\n        \"val\": A.Compose([\n            A.Resize(img_size, img_size),\n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n            ToTensorV2(),\n        ]),\n        \"test\": A.Compose([ # For inference\n            A.Resize(img_size, img_size),\n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n            ToTensorV2(),\n        ]),\n    }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-08T04:52:41.640985Z","iopub.execute_input":"2025-11-08T04:52:41.641695Z","iopub.status.idle":"2025-11-08T04:52:41.649103Z","shell.execute_reply.started":"2025-11-08T04:52:41.641657Z","shell.execute_reply":"2025-11-08T04:52:41.647825Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -------------------------\n# Cell 5: Datasets\n# -------------------------\n\n# --- MODIFIED: Segmentation Dataset with Augs ---\nclass ForgerySegDataset(Dataset):\n    def __init__(self, auth_paths, forg_paths, mask_dir, transforms):\n        self.samples = []\n        self.mask_dir = mask_dir\n        self.transforms = transforms\n\n        for p in auth_paths:\n            self.samples.append((p, None))\n\n        for p in forg_paths:\n            stem = Path(p).stem\n            m = os.path.join(mask_dir, stem + \".npy\")\n            self.samples.append((p, m if os.path.exists(m) else None))\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx):\n        img_path, mask_path = self.samples[idx]\n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        h, w, _ = img.shape\n\n        if (mask_path is None):\n            mask = np.zeros((h, w), dtype=np.uint8)\n        else:\n            m = np.load(mask_path)\n            if m.ndim == 3:\n                m = np.max(m, axis=0)\n            mask = m.astype(np.uint8)\n            if mask.shape != (h, w):\n                mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)\n\n        # Apply transforms\n        augmented = self.transforms(image=img, mask=mask)\n        img_t = augmented['image']\n        mask_t = augmented['mask'].unsqueeze(0).float() # [1,H,W]\n        \n        return img_t, mask_t\n\n# --- MODIFIED: Classification Dataset with Augs ---\nclass ForgeryClsDataset(Dataset):\n    def __init__(self, auth_paths, forg_paths, transforms):\n        self.items = [(p,0) for p in auth_paths] + [(p,1) for p in forg_paths]\n        self.transforms = transforms\n        \n    def __len__(self): return len(self.items)\n    \n    def __getitem__(self, idx):\n        p, y = self.items[idx]\n        img = cv2.imread(p)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        \n        # Apply transforms (no mask)\n        augmented = self.transforms(image=img)\n        x = augmented['image']\n        \n        return x, torch.tensor(y, dtype=torch.long)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-08T04:52:42.167259Z","iopub.execute_input":"2025-11-08T04:52:42.167752Z","iopub.status.idle":"2025-11-08T04:52:42.176630Z","shell.execute_reply.started":"2025-11-08T04:52:42.167733Z","shell.execute_reply":"2025-11-08T04:52:42.176033Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# -------------------------\n# Cell 6: File Loading & Splitting\n# -------------------------\n\nauth_imgs = sorted([str(Path(CFG.AUTH_DIR)/f) for f in os.listdir(CFG.AUTH_DIR) if f.lower().endswith((\".png\",\".jpg\",\".jpeg\",\".tif\"))])\nforg_imgs = sorted([str(Path(CFG.FORG_DIR)/f) for f in os.listdir(CFG.FORG_DIR) if f.lower().endswith((\".png\",\".jpg\",\".jpeg\",\".tif\"))])\nprint(f\"Authentic images: {len(auth_imgs)}, Forged images: {len(forg_imgs)}\")\n\n# Split (must be same for both tasks)\ntrain_auth, val_auth = train_test_split(auth_imgs, test_size=0.2, random_state=CFG.SEED)\ntrain_forg, val_forg = train_test_split(forg_imgs, test_size=0.2, random_state=CFG.SEED)\n\n# -------------------------\n# Cell 7: DINOv2 Loader & Model Defs\n# -------------------------\n\nfrom transformers import AutoImageProcessor, AutoModel\n\ntry:\n    processor = AutoImageProcessor.from_pretrained(CFG.DINO_PATH, local_files_only=True)\n    dino_encoder = AutoModel.from_pretrained(CFG.DINO_PATH, local_files_only=True)\nexcept Exception as e:\n    raise RuntimeError(f\"Could not load DINOv2 from {CFG.DINO_PATH}: {e}\")\n\ndino_encoder.eval().to(device)\n\n# --- Deducing grid size (robustly) ---\n# We use the processor's intended size\ncfg = getattr(dino_encoder, \"config\", None)\npatch = getattr(cfg, \"patch_size\", 14)\ninp = getattr(processor, \"size\", {\"shortest_edge\": 224})\n\nif \"shortest_edge\" in inp:\n    proc_size = int(inp[\"shortest_edge\"])\nelse:\n    proc_size = inp.get(\"height\", 224)\n\ngrid_h = grid_w = proc_size // patch\nprint(f\"DINOv2 loaded. Processor size: {proc_size}, Patch size: {patch}, Grid: {grid_h}x{grid_w}\")\n\n# --- Model Definitions (Unchanged from your baseline) ---\nclass DinoTinyDecoder(nn.Module):\n    def __init__(self, in_ch, out_ch=1):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_ch, 256, 3, padding=1, bias=False),\n            nn.BatchNorm2d(256),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(256, 64, 3, padding=1, bias=False),\n            nn.BatchNorm2d(64),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(64, out_ch, 1)\n        )\n    def forward(self, f, out_size):\n        x = nn.functional.interpolate(f, size=out_size, mode=\"bilinear\", align_corners=False)\n        return self.conv(x)\n\nclass DinoMultiTask(nn.Module):\n    def __init__(self, encoder, seg_ch=768, freeze=True):\n        super().__init__()\n        self.encoder = encoder\n        if freeze:\n            for p in self.encoder.parameters():\n                p.requires_grad = False\n        \n        self.seg_head = DinoTinyDecoder(in_ch=seg_ch, out_ch=1)\n        self.cls_head = nn.Linear(seg_ch, 2)\n\n    def forward_features(self, images):\n        # Note: 'images' are expected to be normalized tensors from Dataset\n        B,_,H,W = images.shape\n        \n        # Manually create pixel_values for processor\n        # DINO processor expects [0, 255] images, but our loader gives normalized.\n        # We must *un-normalize* to pass to processor if it expects images.\n        # *Correction*: The HF processor.from_pretrained() when loaded this way\n        # expects *normalized* tensors if we pass `pixel_values`.\n        # Your previous code passed *numpy* arrays, which is different.\n        # Let's stick to the tensor path, as it's cleaner.\n        \n        # The 'processor' object is for preprocessing *raw* images (PIL/numpy).\n        # Our dataset *already* does this (Resize, Normalize, ToTensorV2).\n        # So we can bypass the processor call and feed tensors directly.\n        # BUT: The processor might do things differently (e.g., resize to 224, not 384).\n        # Let's check the processor's expected size. It's `proc_size` (e.g., 224).\n        # Our loader resizes to `IMG_SIZE` (e.g., 384).\n        # We need to resize to the processor's *exact* size.\n\n        # --- MODIFIED forward_features ---\n        # Resize images to the processor's expected size (e.g., 224x224)\n        x = nn.functional.interpolate(images, size=(proc_size, proc_size), mode=\"bilinear\", align_corners=False)\n        \n        with torch.no_grad():\n            feats = self.encoder(pixel_values=x).last_hidden_state  # [B, N, C]\n        \n        B, N, C = feats.shape\n        expected_tokens = grid_h * grid_w\n        \n        if N == expected_tokens + 1:\n            feats_spatial = feats[:, 1:, :]\n            cls_token = feats[:, 0, :]\n        elif N == expected_tokens:\n            feats_spatial = feats\n            cls_token = feats.mean(dim=1)\n        else:\n            # Fallback for unexpected token count\n            s = int(round(math.sqrt(N - 1)))\n            if s*s == (N-1): # CLS token present\n                feats_spatial = feats[:, 1:, :]; cls_token = feats[:, 0, :]\n                s_h, s_w = s, s\n            else: # No CLS token\n                s = int(round(math.sqrt(N)))\n                feats_spatial = feats; cls_token = feats.mean(dim=1)\n                s_h, s_w = s, s\n            \n            print(f\"Warning: Token mismatch. Expected {expected_tokens}, got {N}. Inferred grid {s}x{s}\")\n            fmap = feats_spatial.permute(0,2,1).reshape(B, C, s_h, s_w)\n            return fmap, cls_token\n            \n        fmap = feats_spatial.permute(0,2,1).reshape(B, C, grid_h, grid_w)\n        return fmap, cls_token\n\n    def forward_seg(self, images):\n        fmap, _ = self.forward_features(images)\n        B,_,H,W = images.shape\n        # Upsample to the *input* image size (e.g., 384x384)\n        logits = self.seg_head(fmap, out_size=(H,W))\n        return logits\n\n    def forward_cls(self, images):\n        _, cls_token = self.forward_features(images)\n        return self.cls_head(cls_token)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-08T04:52:42.177764Z","iopub.execute_input":"2025-11-08T04:52:42.178268Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# -------------------------\n# Cell 8: Dataloaders\n# -------------------------\n\ntransforms = get_transforms(CFG.IMG_SIZE)\n\n# Seg Loaders\ntrain_seg_ds = ForgerySegDataset(train_auth, train_forg, CFG.MASK_DIR, transforms=transforms['train'])\nval_seg_ds   = ForgerySegDataset(val_auth,   val_forg,   CFG.MASK_DIR, transforms=transforms['val'])\ntrain_seg_loader = DataLoader(train_seg_ds, batch_size=CFG.BATCH_SEG, shuffle=True, num_workers=2, pin_memory=True)\nval_seg_loader   = DataLoader(val_seg_ds,   batch_size=CFG.BATCH_SEG, shuffle=False, num_workers=2, pin_memory=True)\n\n# Cls Loaders\ntrain_cls_ds = ForgeryClsDataset(train_auth, train_forg, transforms=transforms['train'])\nval_cls_ds   = ForgeryClsDataset(val_auth,   val_forg, transforms=transforms['val'])\ntrain_cls_loader = DataLoader(train_cls_ds, batch_size=CFG.BATCH_CLS, shuffle=True,  num_workers=2, pin_memory=True)\nval_cls_loader   = DataLoader(val_cls_ds,   batch_size=CFG.BATCH_CLS, shuffle=False, num_workers=2, pin_memory=True)\n\nprint(f\"Seg loaders: {len(train_seg_loader)} train, {len(val_seg_loader)} val\")\nprint(f\"Cls loaders: {len(train_cls_loader)} train, {len(val_cls_loader)} val\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# -------------------------\n# Cell 9: Train Segmentation Head\n# -------------------------\n\nprint(\"\\n--- Starting Segmentation Head Training ---\")\nmodel = DinoMultiTask(dino_encoder, seg_ch=768, freeze=True).to(device)\ncrit_seg = ComboLoss(bce_weight=0.5, dice_weight=0.5).to(device) # NEW\nopt_seg  = optim.AdamW(model.seg_head.parameters(), lr=CFG.LR_SEG, weight_decay=CFG.WEIGHT_DECAY)\nsched_seg = CosineAnnealingLR(opt_seg, T_max=CFG.SCHEDULER_T_MAX, eta_min=1e-6)\n\nbest_val_dice = 0.0\n\nfor epoch in range(1, CFG.EPOCHS_SEG + 1):\n    model.train()\n    tr_loss = 0.0\n    for imgs, masks in tqdm(train_seg_loader, desc=f\"[Seg] Epoch {epoch}/{CFG.EPOCHS_SEG}\"):\n        imgs, masks = imgs.to(device), masks.to(device)\n        logits = model.forward_seg(imgs)\n        loss = crit_seg(logits, masks)\n        opt_seg.zero_grad(); loss.backward(); opt_seg.step()\n        tr_loss += loss.item() * imgs.size(0)\n    tr_loss /= len(train_seg_loader.dataset)\n    sched_seg.step() # NEW\n\n    # Validation\n    model.eval()\n    val_loss, m_dice = 0.0, 0.0\n    with torch.no_grad():\n        for imgs, masks in val_seg_loader:\n            imgs, masks = imgs.to(device), masks.to(device)\n            logits = model.forward_seg(imgs)\n            loss = crit_seg(logits, masks)\n            val_loss += loss.item() * imgs.size(0)\n            \n            # Calc dice score\n            probs = torch.sigmoid(logits).cpu().numpy()\n            gts   = masks.cpu().numpy()\n            for p,g in zip(probs, gts):\n                m_dice += dice_score(p[0], g[0]) # Use 0.5 default thr for val\n                \n    n_val = len(val_seg_loader.dataset)\n    val_loss /= n_val\n    m_dice /= n_val\n    \n    print(f\"[Seg] Epoch {epoch} | train_loss={tr_loss:.4f} | val_loss={val_loss:.4f} | val_Dice={m_dice:.3f}\")\n    \n    # NEW: Save best model\n    if m_dice > best_val_dice:\n        best_val_dice = m_dice\n        torch.save(model.state_dict(), CFG.BEST_SEG_MODEL)\n        print(f\"  -> New best seg model saved with Dice: {best_val_dice:.4f}\")\n\ntorch.cuda.empty_cache(); gc.collect()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# -------------------------\n# Cell 10: Train Classification Head\n# -------------------------\n\nprint(\"\\n--- Starting Classification Head Training ---\")\n# Note: We continue training on the *same* model instance.\n# The seg head is already trained, now we train the cls head.\n# We could load the best seg model, but since the backbone is frozen,\n# the heads are independent.\n\ncrit_cls = nn.CrossEntropyLoss()\nopt_cls  = optim.AdamW(model.cls_head.parameters(), lr=CFG.LR_CLS, weight_decay=CFG.WEIGHT_DECAY)\nsched_cls = CosineAnnealingLR(opt_cls, T_max=CFG.SCHEDULER_T_MAX, eta_min=1e-6)\n\nbest_val_acc = 0.0\n\nfor epoch in range(1, CFG.EPOCHS_CLS + 1):\n    model.train()\n    tr_loss = 0.0\n    for imgs, labels in tqdm(train_cls_loader, desc=f\"[Cls] Epoch {epoch}/{CFG.EPOCHS_CLS}\"):\n        imgs, labels = imgs.to(device), labels.to(device)\n        logits = model.forward_cls(imgs)\n        loss = crit_cls(logits, labels)\n        opt_cls.zero_grad(); loss.backward(); opt_cls.step()\n        tr_loss += loss.item() * imgs.size(0)\n    tr_loss /= len(train_cls_loader.dataset)\n    sched_cls.step()\n\n    # Validation\n    model.eval()\n    correct, total = 0, 0\n    with torch.no_grad():\n        for imgs, labels in val_cls_loader:\n            imgs, labels = imgs.to(device), labels.to(device)\n            logits = model.forward_cls(imgs)\n            preds = logits.argmax(dim=1)\n            correct += (preds == labels).sum().item()\n            total   += labels.size(0)\n    \n    val_acc = 100.0 * correct / total\n    print(f\"[Cls] Epoch {epoch} | train_loss={tr_loss:.4f} | val_acc={val_acc:.2f}%\")\n    \n    # NEW: Save best *overall* model\n    if val_acc > best_val_acc:\n        best_val_acc = val_acc\n        # We save the *entire* model state (best seg head + best cls head)\n        # Note: This assumes the best seg head was from the *last* seg epoch\n        # if we don't reload.\n        # A safer way: load best seg, train cls, save cls.\n        \n        # Let's do the safer way:\n        # Load best seg weights *before* training cls.\n        # This part should be run *before* the cls loop:\n        \n        # --- Start of Cell 10 (Revised) ---\n        print(\"\\n--- Loading Best Seg Model & Training Classification Head ---\")\n        try:\n            model.load_state_dict(torch.load(CFG.BEST_SEG_MODEL))\n            print(f\"Loaded best segmentation model from {CFG.BEST_SEG_MODEL}\")\n        except Exception as e:\n            print(f\"Warning: Could not load best seg model. Training cls head on last epoch. {e}\")\n\n        crit_cls = nn.CrossEntropyLoss()\n        opt_cls  = optim.AdamW(model.cls_head.parameters(), lr=CFG.LR_CLS, weight_decay=CFG.WEIGHT_DECAY)\n        sched_cls = CosineAnnealingLR(opt_cls, T_max=CFG.SCHEDULER_T_MAX, eta_min=1e-6)\n        \n        best_val_acc = 0.0\n        \n        for epoch in range(1, CFG.EPOCHS_CLS + 1):\n            model.train() # Set seg_head to train mode (for BN) but grads are off\n            model.cls_head.train()\n            \n            tr_loss = 0.0\n            for imgs, labels in tqdm(train_cls_loader, desc=f\"[Cls] Epoch {epoch}/{CFG.EPOCHS_CLS}\"):\n                imgs, labels = imgs.to(device), labels.to(device)\n                logits = model.forward_cls(imgs)\n                loss = crit_cls(logits, labels)\n                opt_cls.zero_grad(); loss.backward(); opt_cls.step()\n                tr_loss += loss.item() * imgs.size(0)\n            tr_loss /= len(train_cls_loader.dataset)\n            sched_cls.step()\n        \n            # Validation\n            model.eval()\n            correct, total = 0, 0\n            with torch.no_grad():\n                for imgs, labels in val_cls_loader:\n                    imgs, labels = imgs.to(device), labels.to(device)\n                    logits = model.forward_cls(imgs)\n                    preds = logits.argmax(dim=1)\n                    correct += (preds == labels).sum().item()\n                    total   += labels.size(0)\n            \n            val_acc = 100.0 * correct / total\n            print(f\"[Cls] Epoch {epoch} | train_loss={tr_loss:.4f} | val_acc={val_acc:.2f}%\")\n            \n            if val_acc > best_val_acc:\n                best_val_acc = val_acc\n                torch.save(model.state_dict(), CFG.BEST_CLS_MODEL)\n                print(f\"  -> New best cls model saved with Acc: {best_val_acc:.2f}%\")\n\n        torch.cuda.empty_cache(); gc.collect()\n        # --- End of Cell 10 (Revised) ---\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n\n# -------------------------\n# Cell 11: Find Optimal Threshold (NEW)\n# -------------------------\nprint(f\"\\n--- Finding Optimal Segmentation Threshold ---\")\n# Load the best model (which has best seg + best cls heads)\ntry:\n    model.load_state_dict(torch.load(CFG.BEST_CLS_MODEL))\n    print(f\"Loaded best multi-task model from {CFG.BEST_CLS_MODEL}\")\nexcept Exception as e:\n    print(f\"Warning: Could not load {CFG.BEST_CLS_MODEL}. Using last model state. {e}\")\nmodel.eval()\n\n# Get all validation preds and gts\nall_preds = []\nall_gts = []\nwith torch.no_grad():\n    for imgs, masks in tqdm(val_seg_loader, desc=\"Finding Threshold\"):\n        imgs = imgs.to(device)\n        logits = model.forward_seg(imgs)\n        probs = torch.sigmoid(logits).cpu().numpy()\n        gts   = masks.cpu().numpy()\n        \n        for p, g in zip(probs, gts):\n            all_preds.append(p[0]) # p is [1,H,W], g is [1,H,W]\n            all_gts.append(g[0])\n\n# Test thresholds\nthresholds = np.arange(0.2, 0.8, 0.05)\nbest_dice = 0.0\nOPTIMAL_THR = 0.5\n\nfor thr in thresholds:\n    current_dice = 0.0\n    for p, g in zip(all_preds, all_gts):\n        current_dice += dice_score(p, g, eps=1e-7)\n    \n    mean_dice = current_dice / len(all_preds)\n    print(f\"Threshold {thr:.2f} -> Mean Dice: {mean_dice:.4f}\")\n    \n    if mean_dice > best_dice:\n        best_dice = mean_dice\n        OPTIMAL_THR = thr\n\nprint(f\"\\n==> Found Optimal Threshold: {OPTIMAL_THR:.2f} (Dice: {best_dice:.4f})\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# -------------------------\n# Cell 12: Inference & Submission (NEW LOGIC)\n# -------------------------\nprint(\"\\n--- Starting Test Inference ---\")\n\n# Ensure model is in eval mode\nmodel.eval()\ntest_transforms = get_transforms(CFG.IMG_SIZE)['test']\n\n# --- NEW: Prediction functions ---\ndef predict_class(img_pil):\n    img = np.array(img_pil.convert(\"RGB\"))\n    img_t = test_transforms(image=img)['image'].unsqueeze(0).to(device)\n    with torch.no_grad():\n        logits = model.forward_cls(img_t)\n        pred = logits.argmax(dim=1).item()\n    return pred\n\ndef predict_mask_prob(img_pil):\n    img = np.array(img_pil.convert(\"RGB\"))\n    img_t = test_transforms(image=img)['image'].unsqueeze(0).to(device)\n    with torch.no_grad():\n        logits = model.forward_seg(img_t)\n        prob = torch.sigmoid(logits)[0,0].cpu().numpy()\n    return prob\n\n# --- NEW: Two-Stage Inference Loop ---\nrows = []\nif os.path.exists(CFG.TEST_DIR):\n    test_files = sorted(os.listdir(CFG.TEST_DIR))\n    print(f\"Found {len(test_files)} test images.\")\n    \n    for fname in tqdm(test_files, desc=\"Inference\"):\n        case_id = Path(fname).stem\n        path = str(Path(CFG.TEST_DIR)/fname)\n        \n        try:\n            pil = Image.open(path).convert(\"RGB\")\n            ow, oh = pil.size\n\n            # --- STAGE 1: CLASSIFY ---\n            pred_class = predict_class(pil)\n\n            if pred_class == 0:\n                # Predicted 'authentic'\n                annot = \"authentic\"\n            else:\n                # --- STAGE 2: SEGMENT ---\n                # Predicted 'forged', now get the mask\n                prob = predict_mask_prob(pil) # [IMG_SIZE, IMG_SIZE]\n                \n                # Resize to original size\n                mask = cv2.resize(prob, (ow, oh), interpolation=cv2.INTER_NEAREST)\n                binm = (mask > OPTIMAL_THR).astype(np.uint8) # Use optimal thr\n\n                if binm.sum() == 0:\n                    # Fallback: Classifier said forged, but segmentor found nothing.\n                    annot = \"authentic\"\n                else:\n                    annot = rle_encode_numpy(binm)\n\n            rows.append({\"case_id\": case_id, \"annotation\": annot})\n            \n        except Exception as e:\n            print(f\"Error processing {fname}: {e}. Defaulting to 'authentic'.\")\n            rows.append({\"case_id\": case_id, \"annotation\": \"authentic\"})\n\nelse:\n    print(\"Test directory not found. Using sample submission.\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# -------------------------\n# Cell 13: Submission File Creation\n# -------------------------\n\nsub = pd.DataFrame(rows, columns=[\"case_id\",\"annotation\"])\n\nif sub.empty:\n    print(\"No test images processed. Creating submission from sample.\")\n    sub = pd.read_csv(CFG.SAMPLE_SUB)\n    sub[\"annotation\"] = \"authentic\"\nelse:\n    # Aligne with sample_submission\n    if os.path.exists(CFG.SAMPLE_SUB):\n        ss = pd.read_csv(CFG.SAMPLE_SUB)\n        ss[\"case_id\"] = ss[\"case_id\"].astype(str)\n        sub[\"case_id\"] = sub[\"case_id\"].astype(str)\n        \n        # Merge to ensure all test_ids are present\n        sub = ss[[\"case_id\"]].merge(sub, on=\"case_id\", how=\"left\")\n        \n        # Fill any missing (e.g., from errors) with 'authentic'\n        sub[\"annotation\"] = sub[\"annotation\"].fillna(\"authentic\")\n    else:\n        print(\"Sample submission not found. Saving as is.\")\n\nOUT_PATH = \"submission.csv\"\nsub.to_csv(OUT_PATH, index=False)\nprint(f\"\\nWrote submission to: {OUT_PATH}\")\nprint(sub.head())\n\nprint(\"\\n--- Grandmaster Notebook Finished ---\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}