{"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":14456136,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install segmentation-models-pytorch","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T09:56:13.644162Z","iopub.execute_input":"2025-11-19T09:56:13.644395Z","iopub.status.idle":"2025-11-19T09:57:25.592218Z","shell.execute_reply.started":"2025-11-19T09:56:13.644372Z","shell.execute_reply":"2025-11-19T09:57:25.591469Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 1. Configuration & Imports","metadata":{}},{"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda.amp import autocast, GradScaler\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport segmentation_models_pytorch as smp\nfrom sklearn.model_selection import train_test_split\nfrom tqdm import tqdm\n\nclass CFG:\n    # Paths\n    base_path = \"/kaggle/input/recodai-luc-scientific-image-forgery-detection\"\n    train_img_dir = f\"{base_path}/train_images\"\n    train_mask_dir = f\"{base_path}/train_masks\"\n    model_save_path = \"best_b4_srm_model.pth\"\n    \n    # Model Config\n    encoder = 'efficientnet-b4' \n    weights = 'imagenet'\n    img_size = 768   # High resolution maintained as requested\n    \n    # Training Config\n    # B4 @ 768px is heavy. We use small batch + high accumulation\n    batch_size = 2   \n    accum_iter = 8   # Effective batch size = 16\n    epochs = 15\n    lr = 1e-4\n    weight_decay = 1e-5\n    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n    num_workers = 2","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T09:57:27.395615Z","iopub.execute_input":"2025-11-19T09:57:27.396251Z","iopub.status.idle":"2025-11-19T09:57:38.749084Z","shell.execute_reply.started":"2025-11-19T09:57:27.396219Z","shell.execute_reply":"2025-11-19T09:57:38.748457Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2. The GPU-Accelerated Model","metadata":{}},{"cell_type":"code","source":"class SRMConv2d(nn.Module):\n    def __init__(self):\n        super().__init__()\n        # Standard Steganalysis filters (Spam 12, 11, 14 variants)\n        filters = [\n            [[0, 0, 0, 0, 0], [0, -1, 2, -1, 0], [0, 2, -4, 2, 0], [0, -1, 2, -1, 0], [0, 0, 0, 0, 0]],\n            [[-1, 2, -2, 2, -1], [2, -6, 8, -6, 2], [-2, 8, -12, 8, -2], [2, -6, 8, -6, 2], [-1, 2, -2, 2, -1]],\n            [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 1, -2, 1, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]\n        ]\n        # Shape: (Out=3, In=1, K=5, K=5)\n        kernel = torch.tensor(filters, dtype=torch.float32).unsqueeze(1)\n        \n        # Register as buffer so it saves with state_dict but creates no gradients\n        self.register_buffer('weight', kernel)\n        \n    def forward(self, x):\n        \"\"\"\n        Input: (B, 3, H, W) normalized RGB\n        Output: (B, 3, H, W) Noise Features\n        \"\"\"\n        # RGB to Grayscale (Standard Rec. 601)\n        # x[:, 0:1] is Red, x[:, 1:2] is Green...\n        gray = x[:, 0:1, :, :] * 0.299 + x[:, 1:2, :, :] * 0.587 + x[:, 2:3, :, :] * 0.114\n        \n        # Convolution on GPU\n        return F.conv2d(gray, self.weight, padding=2)\n\nclass ForgeryDetectorB4(nn.Module):\n    def __init__(self, encoder_name, encoder_weights):\n        super().__init__()\n        self.srm_layer = SRMConv2d()\n        \n        # U-Net with 6 input channels (3 RGB + 3 SRM)\n        self.unet = smp.Unet(\n            encoder_name=encoder_name, \n            encoder_weights=encoder_weights, \n            in_channels=6, \n            classes=1,\n            decoder_attention_type='scse'\n        )\n        \n    def forward(self, x):\n        # 1. Compute noise on GPU\n        noise = self.srm_layer(x)\n        \n        # 2. Concatenate\n        combined = torch.cat([x, noise], dim=1) \n        \n        # 3. Predict\n        return self.unet(combined)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T09:57:56.410394Z","iopub.execute_input":"2025-11-19T09:57:56.410913Z","iopub.status.idle":"2025-11-19T09:57:56.419738Z","shell.execute_reply.started":"2025-11-19T09:57:56.410890Z","shell.execute_reply":"2025-11-19T09:57:56.418859Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 3. The Optimized Dataset & Transforms","metadata":{}},{"cell_type":"code","source":"class ScientificForgeryDataset(Dataset):\n    def __init__(self, df, transform=None):\n        self.df = df\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        \n        # Load RGB\n        image = cv2.imread(row['image_path'])\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        # Load Mask\n        h, w = image.shape[:2]\n        if row['label'] == 0:\n            mask = np.zeros((h, w), dtype=np.float32)\n        else:\n            try:\n                mask = np.load(row['mask_path'])\n                if mask.ndim == 3: mask = np.max(mask, axis=0)\n                mask = (mask > 0).astype(np.float32)\n            except:\n                mask = np.zeros((h, w), dtype=np.float32)\n\n        # Augment\n        if self.transform:\n            augmented = self.transform(image=image, mask=mask)\n            image = augmented['image']\n            mask = augmented['mask']\n            \n        # Returns (3, H, W) tensor. Model handles the rest.\n        return image, mask.unsqueeze(0)\n\ndef get_transforms(data):\n    if data == 'train':\n        return A.Compose([\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.RandomRotate90(p=0.5),\n            # Heavy geometric augmentations are good for forgery\n            A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.05, rotate_limit=15, p=0.5),\n            A.PadIfNeeded(min_height=CFG.img_size, min_width=CFG.img_size, border_mode=0, value=0),\n            A.RandomCrop(height=CFG.img_size, width=CFG.img_size),\n            A.Normalize(),\n            ToTensorV2(),\n        ])\n    elif data == 'valid':\n        return A.Compose([\n            A.PadIfNeeded(min_height=CFG.img_size, min_width=CFG.img_size, border_mode=0, value=0),\n            A.CenterCrop(height=CFG.img_size, width=CFG.img_size),\n            A.Normalize(),\n            ToTensorV2(),\n        ])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T09:58:12.475109Z","iopub.execute_input":"2025-11-19T09:58:12.475421Z","iopub.status.idle":"2025-11-19T09:58:12.483853Z","shell.execute_reply.started":"2025-11-19T09:58:12.475399Z","shell.execute_reply":"2025-11-19T09:58:12.483008Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 4. Advanced Training Engine","metadata":{}},{"cell_type":"code","source":"class WeightedBCEDiceLoss(nn.Module):\n    def __init__(self, pos_weight=50):\n        super().__init__()\n        self.pos_weight = pos_weight\n        self.dice = smp.losses.DiceLoss(mode='binary')\n\n    def forward(self, logits, targets):\n        bce_weight = torch.tensor([self.pos_weight], device=logits.device)\n        bce_loss = F.binary_cross_entropy_with_logits(\n            logits, targets, pos_weight=bce_weight\n        )\n        probs = torch.sigmoid(logits)\n        dice_loss = self.dice(probs, targets)\n        return (bce_loss * 0.5) + (dice_loss * 0.5)\n\ndef train_one_epoch(model, loader, criterion, optimizer, scaler, device):\n    model.train()\n    running_loss = 0.0\n    \n    pbar = tqdm(loader, desc=\"Train\")\n    for i, (images, masks) in enumerate(pbar):\n        images, masks = images.to(device), masks.to(device)\n        \n        with autocast():\n            outputs = model(images)\n            loss = criterion(outputs, masks)\n            loss = loss / CFG.accum_iter\n            \n        scaler.scale(loss).backward()\n        \n        if ((i + 1) % CFG.accum_iter == 0) or (i + 1 == len(loader)):\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad()\n            \n        running_loss += (loss.item() * CFG.accum_iter)\n        pbar.set_postfix(loss=loss.item() * CFG.accum_iter)\n        \n    return running_loss / len(loader)\n\ndef validate(model, loader, criterion, device):\n    model.eval()\n    running_loss = 0.0\n    \n    with torch.no_grad():\n        for images, masks in tqdm(loader, desc=\"Valid\"):\n            images, masks = images.to(device), masks.to(device)\n            with autocast():\n                outputs = model(images)\n                loss = criterion(outputs, masks)\n            running_loss += loss.item()\n            \n    return running_loss / len(loader)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T09:58:32.096651Z","iopub.execute_input":"2025-11-19T09:58:32.097371Z","iopub.status.idle":"2025-11-19T09:58:32.105686Z","shell.execute_reply.started":"2025-11-19T09:58:32.097344Z","shell.execute_reply":"2025-11-19T09:58:32.105019Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 5. Advanced Inference (TTA + Morphology)","metadata":{}},{"cell_type":"code","source":"def predict_robust(model, image_tensor, device, threshold=0.5):\n    \"\"\"\n    Inputs: image_tensor (1, 3, H, W)\n    Returns: binary_mask (H, W)\n    \"\"\"\n    model.eval()\n    input_img = image_tensor.to(device)\n    \n    # --- TTA: Original + HFlip + VFlip ---\n    preds = []\n    with torch.no_grad():\n        # 1. Original\n        logits = model(input_img)\n        preds.append(torch.sigmoid(logits).cpu().numpy()[0, 0])\n        \n        # 2. Horizontal Flip\n        logits_h = model(torch.flip(input_img, [3]))\n        pred_h = torch.sigmoid(logits_h).cpu().numpy()[0, 0]\n        preds.append(np.fliplr(pred_h)) # Flip back\n        \n        # 3. Vertical Flip\n        logits_v = model(torch.flip(input_img, [2]))\n        pred_v = torch.sigmoid(logits_v).cpu().numpy()[0, 0]\n        preds.append(np.flipud(pred_v)) # Flip back\n    \n    # Average Predictions\n    avg_pred = np.mean(preds, axis=0)\n    \n    # --- Post-Processing (Morphology) ---\n    binary_mask = (avg_pred > threshold).astype(np.uint8)\n    \n    # Kernel size relative to 768px image\n    kernel = np.ones((5, 5), np.uint8)\n    \n    # Opening (remove small noise) -> Closing (fill small holes)\n    processed = cv2.morphologyEx(binary_mask, cv2.MORPH_OPEN, kernel)\n    processed = cv2.morphologyEx(processed, cv2.MORPH_CLOSE, kernel)\n    \n    return processed","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T09:58:48.061695Z","iopub.execute_input":"2025-11-19T09:58:48.061989Z","iopub.status.idle":"2025-11-19T09:58:48.068353Z","shell.execute_reply.started":"2025-11-19T09:58:48.061969Z","shell.execute_reply":"2025-11-19T09:58:48.067518Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 6. Main Execution Block","metadata":{}},{"cell_type":"code","source":"import glob\n\nif __name__ == \"__main__\":\n    # --- 1. Populate DataFrames (The Missing Logic) ---\n    data = []\n    \n    print(f\"Scanning images in {CFG.train_img_dir}...\")\n    \n    # A. Load Authentic Images (Label 0, No Mask)\n    auth_files = glob.glob(f\"{CFG.train_img_dir}/authentic/*.png\")\n    print(f\"Found {len(auth_files)} authentic images.\")\n    for path in auth_files:\n        data.append({\n            \"image_path\": path, \n            \"mask_path\": None, \n            \"label\": 0\n        })\n        \n    # B. Load Forged Images (Label 1, Has Mask)\n    forged_files = glob.glob(f\"{CFG.train_img_dir}/forged/*.png\")\n    print(f\"Found {len(forged_files)} forged images.\")\n    \n    for path in forged_files:\n        # Match image to mask by filename ID\n        file_id = os.path.basename(path).split('.')[0]\n        mask_path = f\"{CFG.train_mask_dir}/{file_id}.npy\"\n        \n        # Only add if mask actually exists\n        if os.path.exists(mask_path):\n            data.append({\n                \"image_path\": path, \n                \"mask_path\": mask_path, \n                \"label\": 1\n            })\n            \n    train_df = pd.DataFrame(data)\n    print(f\"Total dataset size: {len(train_df)}\")\n    \n    # Safety Check\n    if len(train_df) == 0:\n        raise ValueError(\"No images found! Check your CFG.train_img_dir paths.\")\n\n    # --- 2. Split Data ---\n    # Stratify ensures we keep the same ratio of forged/authentic in train and valid\n    train_df, valid_df = train_test_split(\n        train_df, \n        test_size=0.15, \n        stratify=train_df['label'],\n        random_state=42\n    )\n    \n    # --- 3. Loaders ---\n    train_ds = ScientificForgeryDataset(train_df, transform=get_transforms('train'))\n    valid_ds = ScientificForgeryDataset(valid_df, transform=get_transforms('valid'))\n    \n    train_loader = DataLoader(train_ds, batch_size=CFG.batch_size, shuffle=True, \n                              num_workers=CFG.num_workers, pin_memory=True)\n    valid_loader = DataLoader(valid_ds, batch_size=CFG.batch_size, shuffle=False, \n                              num_workers=CFG.num_workers, pin_memory=True)\n    \n    # --- 4. Initialize Custom GPU-Accelerated Model ---\n    model = ForgeryDetectorB4(CFG.encoder, CFG.weights).to(CFG.device)\n    \n    # --- 5. Optimization ---\n    optimizer = torch.optim.AdamW(model.parameters(), lr=CFG.lr, weight_decay=CFG.weight_decay)\n    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=2, factor=0.5)\n    criterion = WeightedBCEDiceLoss(pos_weight=50)\n    scaler = GradScaler()\n    \n    # --- 6. Training Loop ---\n    best_loss = float('inf')\n    \n    print(f\"Starting training with {CFG.encoder} @ {CFG.img_size}px\")\n    \n    for epoch in range(CFG.epochs):\n        print(f\"\\nEpoch {epoch+1}/{CFG.epochs}\")\n        \n        train_loss = train_one_epoch(model, train_loader, criterion, optimizer, scaler, CFG.device)\n        valid_loss = validate(model, valid_loader, criterion, CFG.device)\n        \n        scheduler.step(valid_loss)\n        \n        print(f\"Train Loss: {train_loss:.4f} | Valid Loss: {valid_loss:.4f}\")\n        \n        if valid_loss < best_loss:\n            best_loss = valid_loss\n            torch.save(model.state_dict(), CFG.model_save_path)\n            print(\">>> Saved Best Model\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T10:00:30.809218Z","iopub.execute_input":"2025-11-19T10:00:30.809913Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}