{"cells":[{"cell_type":"code","execution_count":null,"id":"main","metadata":{},"outputs":[],"source":"#!/usr/bin/env python3\n\"\"\"\nKaggle Inference - Scientific Image Forgery Detection\nModel: Trial 14 Retrained (IoU=0.4803)\nArchitecture: n_layers=4, base_channels=64, use_attention=True\n\"\"\"\n\nimport os\nimport glob\nimport math\nimport json\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom transformers import AutoImageProcessor, AutoModel\nfrom tqdm import tqdm\n\n# Config\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nIMG_SIZE = 518\nUSE_TTA = True\nUSE_MULTISCALE = True\n\n# Post-processing (optimized)\nALPHA_GRAD = 0.10\nTHRESHOLD_MULTIPLIER = 0.245\nMORPH_CLOSE_KERNEL = 7\nMORPH_OPEN_KERNEL = 7\nMIN_AREA = 350\nMIN_CONFIDENCE = 0.28\nMAX_AREA_RATIO = 0.72\n\n# Model architecture (MUST match training)\nMODEL_CONFIG = {\n    'base_channels': 64,\n    'n_layers': 4,\n    'use_attention': True,\n    'dropout': 0.218\n}\n\nprint(f\"Device: {DEVICE}\")\n\n\ndef rle_encode(mask):\n    pixels = mask.T.flatten()\n    dots = np.where(pixels == 1)[0]\n    if len(dots) == 0:\n        return \"authentic\"\n    run_lengths = []\n    prev = -2\n    for d in dots:\n        if d > prev + 1:\n            run_lengths.extend((d + 1, 0))\n        run_lengths[-1] += 1\n        prev = d\n    return json.dumps([int(x) for x in run_lengths])\n\n\nclass ConvBlock(nn.Module):\n    def __init__(self, in_ch, out_ch, dropout=0.0):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_ch, out_ch, 3, padding=1),\n            nn.BatchNorm2d(out_ch),\n            nn.ReLU(inplace=True),\n            nn.Dropout2d(dropout) if dropout > 0 else nn.Identity(),\n            nn.Conv2d(out_ch, out_ch, 3, padding=1),\n            nn.BatchNorm2d(out_ch),\n            nn.ReLU(inplace=True),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass AttentionBlock(nn.Module):\n    def __init__(self, channels):\n        super().__init__()\n        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n        self.fc = nn.Sequential(\n            nn.Linear(channels, channels // 8),\n            nn.ReLU(inplace=True),\n            nn.Linear(channels // 8, channels),\n            nn.Sigmoid()\n        )\n\n    def forward(self, x):\n        b, c, _, _ = x.size()\n        y = self.avg_pool(x).view(b, c)\n        y = self.fc(y).view(b, c, 1, 1)\n        return x * y\n\n\nclass FlexibleDecoder(nn.Module):\n    def __init__(self, in_channels=768, base_channels=64, use_attention=True, n_layers=4, dropout=0.218):\n        super().__init__()\n        self.use_attention = use_attention\n        self.n_layers = n_layers\n\n        channels = [in_channels]\n        for i in range(n_layers):\n            channels.append(base_channels * (2 ** (n_layers - 1 - i)))\n\n        self.ups = nn.ModuleList()\n        self.convs = nn.ModuleList()\n        self.attns = nn.ModuleList() if use_attention else None\n\n        for i in range(n_layers):\n            self.ups.append(nn.ConvTranspose2d(channels[i], channels[i+1], 2, stride=2))\n            self.convs.append(ConvBlock(channels[i+1], channels[i+1], dropout))\n            if use_attention:\n                self.attns.append(AttentionBlock(channels[i+1]))\n\n        self.final = nn.Sequential(\n            nn.Conv2d(channels[-1], channels[-1] // 2, 3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(channels[-1] // 2, 1, 1)\n        )\n\n    def forward(self, features, target_size):\n        x = features\n        for i in range(self.n_layers):\n            x = self.ups[i](x)\n            if self.use_attention:\n                x = self.attns[i](x)\n            x = self.convs[i](x)\n        x = F.interpolate(x, size=target_size, mode='bilinear', align_corners=False)\n        x = self.final(x)\n        return x\n\n\nclass DinoV2Segmenter(nn.Module):\n    def __init__(self, dino_model, decoder):\n        super().__init__()\n        self.dino = dino_model\n        self.decoder = decoder\n\n    def extract_features(self, x):\n        with torch.no_grad():\n            outputs = self.dino(x)\n            features = outputs.last_hidden_state[:, 1:, :]\n        B, N, C = features.shape\n        h = w = int(math.sqrt(N))\n        features = features.permute(0, 2, 1).reshape(B, C, h, w)\n        return features\n\n    def forward(self, x, target_size=None):\n        H, W = x.shape[2], x.shape[3]\n        if target_size is None:\n            target_size = (H, W)\n        features = self.extract_features(x)\n        logits = self.decoder(features, target_size)\n        return logits\n\n\ndef find_test_images():\n    comp_dir = '/kaggle/input/recodai-luc-scientific-image-forgery-detection'\n    for path in [os.path.join(comp_dir, 'test_images'), os.path.join(comp_dir, 'test'), comp_dir]:\n        if os.path.exists(path):\n            images = glob.glob(os.path.join(path, '*.png'))\n            if images:\n                print(f\"Found {len(images)} images in {path}\")\n                return sorted(images)\n    images = glob.glob(os.path.join(comp_dir, '**', '*.png'), recursive=True)\n    if images:\n        return sorted(images)\n    raise FileNotFoundError(f\"No test images found\")\n\n\ndef predict_single(model, processor, image, device, target_size):\n    inputs = processor(images=image, return_tensors=\"pt\")\n    pixel_values = inputs['pixel_values'].to(device)\n    with torch.no_grad():\n        outputs = model(pixel_values, target_size=target_size)\n        pred = torch.sigmoid(outputs).squeeze().cpu().numpy()\n    return pred\n\n\ndef predict_with_tta(model, processor, image, device, target_size):\n    predictions = []\n    weights = [1.5, 1.0, 1.0, 0.8, 0.8, 0.8]\n    \n    # Original\n    predictions.append(predict_single(model, processor, image, device, target_size))\n    # H-flip\n    pred = predict_single(model, processor, image.transpose(Image.FLIP_LEFT_RIGHT), device, target_size)\n    predictions.append(np.fliplr(pred))\n    # V-flip\n    pred = predict_single(model, processor, image.transpose(Image.FLIP_TOP_BOTTOM), device, target_size)\n    predictions.append(np.flipud(pred))\n    # Rotations\n    for angle, k in [(90, -1), (180, 2), (270, 1)]:\n        pred = predict_single(model, processor, image.rotate(angle, expand=False), device, target_size)\n        predictions.append(np.rot90(pred, k))\n    \n    weighted_pred = sum(p * w for p, w in zip(predictions, weights))\n    return weighted_pred / sum(weights)\n\n\ndef predict_multiscale(model, processor, image, device):\n    scales = [448, 518, 588]\n    scale_weights = [1.0, 1.5, 1.0]\n    predictions = []\n    \n    for scale, weight in zip(scales, scale_weights):\n        if USE_TTA:\n            pred = predict_with_tta(model, processor, image, device, (scale, scale))\n        else:\n            pred = predict_single(model, processor, image, device, (scale, scale))\n        if scale != IMG_SIZE:\n            pred = cv2.resize(pred, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_LINEAR)\n        predictions.append(pred * weight)\n    \n    return sum(predictions) / sum(scale_weights)\n\n\ndef postprocess(preds, original_size):\n    # Gradient enhancement\n    gx = cv2.Sobel(preds.astype(np.float32), cv2.CV_32F, 1, 0, ksize=3)\n    gy = cv2.Sobel(preds.astype(np.float32), cv2.CV_32F, 0, 1, ksize=3)\n    grad_mag = np.sqrt(gx**2 + gy**2)\n    grad_norm = grad_mag / (grad_mag.max() + 1e-6)\n    \n    enhanced = (1 - ALPHA_GRAD) * preds + ALPHA_GRAD * grad_norm\n    enhanced = cv2.GaussianBlur(enhanced, (5, 5), 0)\n    \n    # Threshold\n    thr = np.mean(enhanced) + THRESHOLD_MULTIPLIER * np.std(enhanced)\n    mask = (enhanced > thr).astype(np.uint8)\n    \n    # Morphology\n    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((MORPH_CLOSE_KERNEL, MORPH_CLOSE_KERNEL), np.uint8))\n    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((MORPH_OPEN_KERNEL, MORPH_OPEN_KERNEL), np.uint8))\n    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((11, 11), np.uint8))\n    \n    mask = cv2.resize(mask, original_size, interpolation=cv2.INTER_NEAREST)\n    return mask, preds\n\n\ndef infer_image(model, processor, image, device):\n    orig_size = image.size\n    if image.mode != 'RGB':\n        image = image.convert('RGB')\n    \n    if USE_MULTISCALE:\n        preds = predict_multiscale(model, processor, image, device)\n    elif USE_TTA:\n        preds = predict_with_tta(model, processor, image, device, (IMG_SIZE, IMG_SIZE))\n    else:\n        preds = predict_single(model, processor, image, device, (IMG_SIZE, IMG_SIZE))\n    \n    mask, raw_preds = postprocess(preds, orig_size)\n    \n    area = int(mask.sum())\n    total_pixels = orig_size[0] * orig_size[1]\n    \n    if area > 0:\n        preds_resized = cv2.resize(preds, orig_size, interpolation=cv2.INTER_LINEAR)\n        mean_confidence = float(preds_resized[mask == 1].mean())\n    else:\n        mean_confidence = 0.0\n    \n    if area < MIN_AREA or mean_confidence < MIN_CONFIDENCE:\n        return \"authentic\", None\n    if area > total_pixels * MAX_AREA_RATIO:\n        return \"authentic\", None\n    \n    return \"forged\", mask\n\n\ndef predict():\n    print(\"Loading model (Trial 14 Retrained - IoU=0.4803)...\")\n    \n    # Load from Kaggle dataset (OFFLINE - no internet needed)\n    model_input = '/kaggle/input/forgery-dinov2-model'\n    backbone_path = os.path.join(model_input, 'dinov2_backbone_trial14.pth')\n    decoder_path = os.path.join(model_input, 'dinov2_decoder_trial14.pth')\n    dinov2_local = os.path.join(model_input, 'dinov2-base')\n\n    # Load DINOv2 from local dataset (offline)\n    print(f\"Loading DINOv2 from {dinov2_local}\")\n    processor = AutoImageProcessor.from_pretrained(dinov2_local, local_files_only=True)\n    dino_model = AutoModel.from_pretrained(dinov2_local, local_files_only=True)\n    dino_model = dino_model.to(DEVICE)\n    \n    # Load fine-tuned backbone\n    if os.path.exists(backbone_path):\n        print(f\"Loading fine-tuned backbone from {backbone_path}\")\n        state_dict = torch.load(backbone_path, map_location=DEVICE, weights_only=True)\n        dino_model.load_state_dict(state_dict, strict=False)\n    \n    dino_model.eval()\n    \n    # Create decoder with EXACT architecture\n    decoder = FlexibleDecoder(\n        in_channels=768,\n        base_channels=MODEL_CONFIG['base_channels'],\n        use_attention=MODEL_CONFIG['use_attention'],\n        n_layers=MODEL_CONFIG['n_layers'],\n        dropout=MODEL_CONFIG['dropout']\n    )\n    \n    if os.path.exists(decoder_path):\n        print(f\"Loading decoder from {decoder_path}\")\n        decoder.load_state_dict(torch.load(decoder_path, map_location=DEVICE, weights_only=True))\n    \n    decoder = decoder.to(DEVICE)\n    decoder.eval()\n    \n    model = DinoV2Segmenter(dino_model, decoder)\n    model.eval()\n    \n    print(\"Model loaded successfully!\")\n    print(f\"Config: n_layers={MODEL_CONFIG['n_layers']}, base_channels={MODEL_CONFIG['base_channels']}\")\n    \n    # Find test images\n    test_images = find_test_images()\n    print(f\"Processing {len(test_images)} test images\")\n    print(f\"Settings: TTA={USE_TTA}, Multiscale={USE_MULTISCALE}\")\n    \n    results = []\n    for img_path in tqdm(test_images):\n        img_name = os.path.basename(img_path)\n        case_id = int(os.path.splitext(img_name)[0])\n        \n        image = Image.open(img_path)\n        label, mask = infer_image(model, processor, image, DEVICE)\n        \n        if label == \"authentic\" or mask is None:\n            annotation = \"authentic\"\n        else:\n            annotation = rle_encode(mask.astype(np.uint8))\n        \n        results.append({'case_id': case_id, 'annotation': annotation})\n    \n    submission = pd.DataFrame(results)\n    submission = submission.sort_values('case_id').reset_index(drop=True)\n    submission.to_csv('submission.csv', index=False)\n    \n    n_authentic = sum(1 for r in results if r['annotation'] == 'authentic')\n    n_forged = len(results) - n_authentic\n    print(f\"\\nSubmission: {n_authentic} authentic, {n_forged} forged ({n_forged/len(results)*100:.1f}%)\")\n    print(submission.head(10))\n    \n    return submission\n\n\nif __name__ == '__main__':\n    predict()"}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"}},"nbformat":4,"nbformat_minor":5}