{"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":"gpu","dataSources":[{"sourceId":113558,"databundleVersionId":14174843,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":631747,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":475685,"modelId":491594}],"dockerImageVersionId":31153,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Simple EDA\n\nThis section re-uses https://www.kaggle.com/code/jirkaborovec/forgery-detection-eda-visual-annotations","metadata":{}},{"cell_type":"code","source":"import os\nimport glob\nimport warnings\n\nPATH_DATASET = \"/kaggle/input/recodai-luc-scientific-image-forgery-detection\"\nauthentic_images = glob.glob(os.path.join(PATH_DATASET, 'train_images', 'authentic', '*.png'))\nforged_images = glob.glob(os.path.join(PATH_DATASET, 'train_images', 'forged', '*.png'))\nforged_images += glob.glob(os.path.join(PATH_DATASET, 'supplemental_images', '*.png'))\n\nwarnings.simplefilter(action='ignore', category=FutureWarning)\nprint(f\"Found {len(authentic_images)} authentic images.\")\nprint(f\"Found {len(forged_images)} forged images.\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:30:21.871829Z","iopub.execute_input":"2025-12-21T20:30:21.872069Z","iopub.status.idle":"2025-12-21T20:30:21.941341Z","shell.execute_reply.started":"2025-12-21T20:30:21.872040Z","shell.execute_reply":"2025-12-21T20:30:21.940654Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Assuming the masks are in a 'train_masks' directory within the data_dir\nmask_dir = os.path.join(PATH_DATASET, 'train_masks')\n\n# Find all .npy files in the train_masks directory and store in a dictionary\nmask_files_dict = {}\nmask_files = glob.glob(os.path.join(mask_dir, '*.npy'))\nmask_files += glob.glob(os.path.join(PATH_DATASET, 'supplemental_masks', '*.npy'))\nfor mask_path in mask_files:\n    basename = os.path.basename(mask_path)\n    filename_without_extension, _ = os.path.splitext(basename) # Remove extension\n    mask_files_dict[filename_without_extension] = mask_path\n\nprint(f\"Found {len(mask_files_dict)} mask files and stored in a dictionary with keys as filenames without extensions.\")\nmask_dict = [f\"{k}: {v}\" for k, v in list(mask_files_dict.items())]\nprint(\"\\n\".join(mask_dict[:5]))","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-12-21T20:30:21.942719Z","iopub.execute_input":"2025-12-21T20:30:21.943200Z","iopub.status.idle":"2025-12-21T20:30:21.987336Z","shell.execute_reply.started":"2025-12-21T20:30:21.943175Z","shell.execute_reply":"2025-12-21T20:30:21.986568Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_mask(mask_path: str):\n    mask_raw = np.load(mask_path)\n    # Sum across the first dimension and binarize: 1 if any channel has a value > 0, 0 otherwise.\n    mask = np.zeros_like(mask_raw[0, :, :], dtype=np.uint8)\n    for c in range(mask_raw.shape[0]):\n        mask[mask_raw[c, :, :] > 0] = c + 1\n    return mask\n\n# Define a list of colors for the different mask levels (excluding background 0)\n# You can customize this list with more colors if you expect more levels\nmask_colors = ['red', 'blue', 'green', 'purple', 'orange', 'brown', 'pink', 'gray', 'olive', 'cyan']","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:30:21.988127Z","iopub.execute_input":"2025-12-21T20:30:21.988450Z","iopub.status.idle":"2025-12-21T20:30:21.993674Z","shell.execute_reply.started":"2025-12-21T20:30:21.988423Z","shell.execute_reply":"2025-12-21T20:30:21.992964Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Overlap authentic and forged cases","metadata":{}},{"cell_type":"code","source":"# Get just the filenames without the path\nauthentic_filenames = [os.path.basename(img_path) for img_path in authentic_images]\nforged_filenames = [os.path.basename(img_path) for img_path in forged_images]\n\n# Find the intersection of the two sets of filenames\noverlapping_filenames = list(set(authentic_filenames).intersection(forged_filenames))\n\nprint(f\"Found {len(overlapping_filenames)} overlapping filenames in authentic and forged folders.\")","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-12-21T20:30:21.994331Z","iopub.execute_input":"2025-12-21T20:30:21.994586Z","iopub.status.idle":"2025-12-21T20:30:22.011443Z","shell.execute_reply.started":"2025-12-21T20:30:21.994558Z","shell.execute_reply":"2025-12-21T20:30:22.010670Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create dictionaries mapping filename without extension to full path for quicker lookup\nauthentic_image_dict = {os.path.splitext(os.path.basename(img_path))[0]: img_path for img_path in authentic_images}\nforged_image_dict = {os.path.splitext(os.path.basename(img_path))[0]: img_path for img_path in forged_images}\n\n# Find filenames that exist in both authentic and forged sets (using keys without extensions)\noverlapping_filenames_without_extension = list(set(authentic_image_dict.keys()).intersection(forged_image_dict.keys()))\n\n# Create pairs of (authentic_path, forged_path, mask_path) for overlapping filenames\nmatching_pairs_with_mask = []\nfor filename_without_extension in overlapping_filenames_without_extension:\n    authentic_path = authentic_image_dict[filename_without_extension]\n    forged_path = forged_image_dict[filename_without_extension]\n    # Check if a mask exists for this forged image (mask_files_dict already uses keys without extension)\n    if filename_without_extension in mask_files_dict:\n        mask_path = mask_files_dict[filename_without_extension]\n        matching_pairs_with_mask.append((authentic_path, forged_path, mask_path))\n    else:\n        print(f\"Warning: No mask found for forged image with filename (without extension): {filename_without_extension}\")\n\nprint(f\"Found {len(matching_pairs_with_mask)} matching image-mask pairs with the same filename (without extension).\")","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-12-21T20:30:22.013271Z","iopub.execute_input":"2025-12-21T20:30:22.013476Z","iopub.status.idle":"2025-12-21T20:30:22.045247Z","shell.execute_reply.started":"2025-12-21T20:30:22.013460Z","shell.execute_reply":"2025-12-21T20:30:22.044480Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import random\nimport matplotlib.pyplot as plt\nimport matplotlib.image as mpimg\nimport numpy as np\n\n# Determine the number of pairs to show (up to 12 rows)\nnum_pairs_to_show = min(12, len(matching_pairs_with_mask))\n\n# Select a random subset of matching pairs\nrandom_matching_pairs_with_mask = random.sample(matching_pairs_with_mask, num_pairs_to_show)\n\n# Create the grid (num_pairs_to_show rows, 3 columns)\n\nfor i in range(num_pairs_to_show):\n    auth_img_path, forged_img_path, mask_path = random_matching_pairs_with_mask[i]\n    auth_img = mpimg.imread(auth_img_path)\n    forged_img = mpimg.imread(forged_img_path)\n    # Load the mask as multilabel\n    mask = load_mask(mask_path)\n    levels = np.unique(mask)[:-1] + 0.5\n\n    fig, axes = plt.subplots(1, 2, figsize=(14, 7)) # Adjust figsize as needed\n    # Display authentic image in the first column\n    axes[0].imshow(auth_img)\n    # Find and draw contours on the second column axes using the mask\n    axes[0].contour(mask, levels=levels, colors=mask_colors, linewidths=1)\n    axes[0].axis('off')\n    axes[0].set_title(f\"Authentic + Mask: {os.path.basename(auth_img_path)}\", fontsize=8)\n\n    # Display forged image with mask contour in the second column\n    axes[1].imshow(forged_img) # Display the forged image\n    # Find and draw contours on the second column axes using the mask\n    axes[1].contour(mask, levels=levels, colors=mask_colors, linewidths=1)\n    axes[1].axis('off')\n    axes[1].set_title(f\"Forged + Mask: {os.path.basename(forged_img_path)}\", fontsize=8)\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"_kg_hide-input":false,"execution":{"iopub.status.busy":"2025-12-21T20:30:22.046118Z","iopub.execute_input":"2025-12-21T20:30:22.046330Z","iopub.status.idle":"2025-12-21T20:30:32.860583Z","shell.execute_reply.started":"2025-12-21T20:30:22.046313Z","shell.execute_reply":"2025-12-21T20:30:32.859917Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Discover how many instances are per image","metadata":{}},{"cell_type":"code","source":"from tqdm.auto import tqdm\n\nall_mask_instances = []\nfor filename_without_extension, mask_path in tqdm(mask_files_dict.items()):\n    mask = np.load(mask_path)\n    all_mask_instances.append(mask.shape[0])\n\nprint(f\"Loaded shapes for {len(all_mask_instances)} masks.\")\nprint(\"Mask shapes:\", set(all_mask_instances))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:30:32.861373Z","iopub.execute_input":"2025-12-21T20:30:32.861810Z","iopub.status.idle":"2025-12-21T20:30:57.285239Z","shell.execute_reply.started":"2025-12-21T20:30:32.861788Z","shell.execute_reply":"2025-12-21T20:30:57.284478Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import collections\nimport seaborn as sns\n\n# Count the occurrences of each instance count\ninstance_counts = collections.Counter(all_mask_instances)\nsorted_instance_counts = dict(sorted(instance_counts.items()))\n\n# Create a bar plot of the instance counts\nplt.figure(figsize=(8, 3))\nsns.barplot(x=list(sorted_instance_counts.keys()), y=list(sorted_instance_counts.values()))\nplt.title(\"Mask Instance Counts\")\nplt.xlabel(\"Number of Instances in Mask\")\nplt.ylabel(\"Occurances\")\nplt.grid(axis='y', alpha=0.75)\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:30:57.286012Z","iopub.execute_input":"2025-12-21T20:30:57.286787Z","iopub.status.idle":"2025-12-21T20:30:58.321069Z","shell.execute_reply.started":"2025-12-21T20:30:57.286767Z","shell.execute_reply":"2025-12-21T20:30:58.320364Z"},"_kg_hide-input":false},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Explore the ratios of object to image size","metadata":{}},{"cell_type":"code","source":"# Initialize a list to store all area ratios\nall_area_ratios = []\n\n# Iterate through mask files\nfor filename_without_extension, mask_path in tqdm(mask_files_dict.items()):\n    # Load the raw mask data (not using the load_mask function as we need individual layers)\n    mask_raw = np.load(mask_path)\n    # Get image dimensions from the mask shape (assuming mask and image have same dimensions)\n    num_instances, height, width = mask_raw.shape\n    total_image_area = height * width\n    # List to store ratios for the current mask\n    mask_area_ratios = []\n\n    # Iterate through mask instances (layers)\n    for instance_layer in mask_raw:\n        # Calculate segmented area for the instance\n        segmented_area = np.sum(instance_layer > 0)\n        # Calculate area ratio\n        area_ratio = segmented_area / total_image_area\n        # Store the ratio\n        mask_area_ratios.append(area_ratio)\n\n    # Extend the main list with ratios from the current mask\n    all_area_ratios.extend(mask_area_ratios)\n\nprint(f\"Calculated area ratios for {len(all_area_ratios)} instances across all masks.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:30:58.322939Z","iopub.execute_input":"2025-12-21T20:30:58.323310Z","iopub.status.idle":"2025-12-21T20:31:05.947003Z","shell.execute_reply.started":"2025-12-21T20:30:58.323292Z","shell.execute_reply":"2025-12-21T20:31:05.946211Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Analyze and visualize the distribution of area ratios\nplt.figure(figsize=(10, 3))\nsns.histplot(all_area_ratios, bins=50, kde=True) # Using 50 bins to show the distribution shape\nplt.title(\"Distribution of Segmented Area Ratios\")\nplt.xlabel(\"Area Ratio (Segmented Area / Total Image Area)\")\nplt.ylabel(\"Frequency\")\nplt.grid(axis='y', alpha=0.75)\nplt.show()\n\n# Print some basic statistics about the area ratios\nprint(\"\\nBasic statistics for area ratios:\")\nprint(f\"Mean: {np.mean(all_area_ratios):.4f}\")\nprint(f\"Median: {np.median(all_area_ratios):.4f}\")\nprint(f\"Standard Deviation: {np.std(all_area_ratios):.4f}\")\nprint(f\"Min: {np.min(all_area_ratios):.4f}\")\nprint(f\"Max: {np.max(all_area_ratios):.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:31:05.947910Z","iopub.execute_input":"2025-12-21T20:31:05.948442Z","iopub.status.idle":"2025-12-21T20:31:06.304619Z","shell.execute_reply.started":"2025-12-21T20:31:05.948420Z","shell.execute_reply":"2025-12-21T20:31:06.303974Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Feature-Based Self-Similarity Analysis for Copy-Move Detection\n\n### Copy-Move Forgery Detection Solution\n\n#### Problem\nDetect copy-move forgeries in biomedical images where parts of an image are copied and pasted elsewhere in the same image to manipulate results.\n\n#### Solution Overview\n\n##### Architecture\n**Self-Similarity Detection with Background Suppression**\n```\nInput Image (H×W×3)\n    ↓\nTIMM Backbone (ResNet/EfficientNet/etc.)\n    ↓\nFeature Extraction (C×H'×W')\n    ↓\nSelf-Similarity Matrix (compare all spatial locations)\n    ↓\nBackground Suppressor (filter uniform/repetitive regions)\n    ↓\nCopy-Move Detection (find distant similar regions)\n    ↓\nRefinement Network\n    ↓\nOutput Heatmap (H×W) [0-1 probability]\n```\n\n##### Key Components\n\n**1. TIMM Feature Extractor**\n- Uses pretrained backbones (ResNet50, EfficientNet, ConvNeXt, etc.)\n- Extracts normalized deep features\n- Transfer learning from ImageNet\n\n**2. Self-Similarity Matrix**\n- Computes cosine similarity between all spatial locations\n- High similarity = potential copied regions\n- Formula: `similarity[i,j] = feature[i] · feature[j]`\n\n**3. Background Suppressor** (Unsupervised)\n- **Problem**: Uniform backgrounds (white/black) create false matches\n- **Solution**: Compute feature complexity using:\n  - Feature variance across channels\n  - Local diversity (patch variations)\n  - Gradient magnitude\n- **Output**: Suppress matches in low-complexity regions\n- **No labels needed**: Works automatically\n\n**4. Copy-Move Detection**\n- Filters out self-matches and nearby locations\n- Only keeps high similarity at distant locations\n- Spatial distance threshold prevents matching adjacent pixels\n- Max similarity score = forgery probability\n\n**5. Refinement Network**\n- Small CNN to refine the heatmap\n- Upsamples to original resolution\n- Applies sigmoid for [0,1] probability\n\n#### Training Strategy\n\n##### Framework: PyTorch Lightning\n- Automatic training loops\n- Built-in callbacks (checkpointing, early stopping)\n- Mixed precision training (16-bit)\n- TorchMetrics for evaluation\n\n##### Loss Function\n**Combined Loss = α × Dice + (1-α) × BCE + β × Complexity**\n\n- **Dice Loss**: Handles class imbalance (small forged regions)\n- **BCE**: Binary cross-entropy for pixel-wise classification\n- **Complexity Regularization**: Encourages high complexity in forged regions\n\n##### Metrics (TorchMetrics)\n- **IoU (Jaccard)**: Intersection over Union\n- **F1 Score**: Harmonic mean of precision/recall\n- **Precision**: Accuracy of forgery predictions\n- **Recall**: Coverage of actual forgeries\n- **Accuracy**: Overall correctness\n\n##### Data Augmentation\n- Horizontal/Vertical flips\n- Rotation (±90°)\n- Brightness/Contrast adjustment\n- Noise and blur\n- **Preserves spatial relationships** (important for copy-move detection)\n\n#### Why This Works\n\n**1. Self-Similarity is the Key Signal**\nCopy-move creates identical patterns at different locations - this is unnatural and detectable through feature similarity.\n\n**2. Background Suppression Prevents False Positives**\nScientific images often have uniform backgrounds that would otherwise match everywhere. By measuring feature complexity, we ignore these regions without needing labels.\n\n**3. Deep Features are Robust**\n- Invariant to small transformations (rotation, scaling, brightness)\n- Capture semantic content, not just pixels\n- Pretrained on ImageNet provides strong initialization\n\n**4. Spatial Distance Filtering**\nOnly flag similarities at distant locations - prevents matching a pixel with itself or immediate neighbors.\n\n**5. Heatmap Output > Binary Masks**\n- Provides uncertainty/confidence\n- Easier to threshold based on use case\n- More interpretable for analysis\n\n#### Technical Advantages\n\n✅ **No foreground/background labels needed** - fully unsupervised suppression  \n✅ **Flexible backbone** - easy to swap TIMM models  \n✅ **End-to-end trainable** - loss directly optimizes detection  \n✅ **Scalable** - PyTorch Lightning handles distributed training  \n✅ **Production ready** - proper logging, checkpointing, metrics  \n\n#### Limitations\n\n❌ **Computational cost**: O(N²) similarity matrix for N pixels  \n❌ **Small forgeries**: May miss very small copied regions  \n❌ **Heavily compressed images**: JPEG artifacts can interfere  \n❌ **Non-rigid transformations**: Assumes copy is similar in appearance  \n\n#### Output\n\n**Training**: Best model checkpoint saved based on validation IoU  \n**Inference**: \n- `predictions/heatmaps/*.npy` - Probability maps [0-1]\n- `predictions/masks/*.npy` - Binary masks (thresholded)\n- `submission.csv` - Kaggle submission format\n\n#### Performance Tips\n\n1. **Increase image size** (512→768) for better detail\n2. **Try different backbones** (efficientnet_b4, convnext_tiny)\n3. **Adjust distance_threshold** (3→5) for larger images\n4. **Ensemble multiple models** for robustness\n5. **Lower threshold** (0.5→0.3) to catch more forgeries (higher recall)","metadata":{}},{"cell_type":"code","source":"# ! pip install -q pytorch_lightning torchmetrics timm albumentations","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:31:06.305423Z","iopub.execute_input":"2025-12-21T20:31:06.305848Z","iopub.status.idle":"2025-12-21T20:31:06.309915Z","shell.execute_reply.started":"2025-12-21T20:31:06.305822Z","shell.execute_reply":"2025-12-21T20:31:06.309179Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader, Subset\nimport pytorch_lightning as pl\nfrom pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping, LearningRateMonitor\nfrom pytorch_lightning.loggers import CSVLogger\nimport torchmetrics\nfrom torchmetrics import MetricCollection\nfrom torchmetrics.classification import BinaryPrecision, BinaryRecall, BinaryF1Score, BinaryJaccardIndex, BinaryAccuracy\nimport timm\n\nimport cv2\nimport pandas as pd\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom typing import Tuple, Optional\nfrom pathlib import Path","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:31:06.311452Z","iopub.execute_input":"2025-12-21T20:31:06.311709Z","iopub.status.idle":"2025-12-21T20:31:54.426825Z","shell.execute_reply.started":"2025-12-21T20:31:06.311685Z","shell.execute_reply":"2025-12-21T20:31:54.425982Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"PATH_MODELS = \"/kaggle/input/timm-resnets-weights/pytorch/default/2\"\nTRAIN_IMAGES_SUBDIR = 'train_images'\nTRAIN_MASKS_SUBDIR = 'train_masks'\nTEST_IMAGES_DIR = f'{PATH_DATASET}/test_images'\nLEARN_BG_SUPPRESSION = True\nIMAGE_SIZE = 512\nBATCH_SIZE = 14\nNUM_EPOCHS = 45\nLEARNING_RATE = 1e-3\nWEIGHT_DECAY = 1e-4\nBACKBONE = 'resnet50'\nTHRESHOLD = 0.5\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\nSAVE_DIR = 'checkpoints'\nOUTPUT_DIR = 'predictions'\nTRAIN_VAL_SPLIT = 0.8\nNUM_WORKERS = 4","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:31:54.427718Z","iopub.execute_input":"2025-12-21T20:31:54.427946Z","iopub.status.idle":"2025-12-21T20:31:54.456742Z","shell.execute_reply.started":"2025-12-21T20:31:54.427929Z","shell.execute_reply":"2025-12-21T20:31:54.456006Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Data and DataModule","metadata":{}},{"cell_type":"code","source":"def load_mask(mask_path: str):\n    \"\"\"Load mask from .npy file with correct format\"\"\"\n    mask_raw = np.load(mask_path)\n    mask = np.zeros_like(mask_raw[0, :, :], dtype=np.uint8)\n    for c in range(mask_raw.shape[0]):\n        mask[mask_raw[c, :, :] > 0] = c + 1\n    return mask\n\n\ndef get_train_transforms(image_size=512):\n    \"\"\"Training augmentations that preserve self-similarity\"\"\"\n    return A.Compose([\n        A.Resize(image_size, image_size),\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n        A.Rotate(limit=30, p=0.5),\n        A.RandomBrightnessContrast(p=0.2), # Increased probability slightly\n        A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=10, val_shift_limit=10, p=0.2),\n        A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.05, rotate_limit=0, p=0.2),\n        # Geometric transformations that preserve relative positions and shapes\n        # A.ElasticTransform(alpha=1, sigma=50, alpha_affine=50, p=0.1),\n        A.GridDistortion(p=0.1),\n        # Removed noise and blur as they can break self-similarity\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])\n\n\ndef get_val_transforms(image_size=512):\n    \"\"\"Validation transforms\"\"\"\n    return A.Compose([\n        A.Resize(image_size, image_size),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:31:54.457653Z","iopub.execute_input":"2025-12-21T20:31:54.458146Z","iopub.status.idle":"2025-12-21T20:31:54.486372Z","shell.execute_reply.started":"2025-12-21T20:31:54.458126Z","shell.execute_reply":"2025-12-21T20:31:54.485735Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CopyMoveDataset(Dataset):\n    \"\"\"Dataset for copy-move forgery detection\"\"\"\n\n    def __init__(self, image_dir, mask_dir=None, transform=None, image_size=512):\n        # Directory containing image folders (authentic and forged)\n        self.image_dir = Path(image_dir)\n        # Directory containing mask files (optional)\n        self.mask_dir = Path(mask_dir) if mask_dir else None\n        # Albumentations transform to apply\n        self.transform = transform\n        # Target image size for resizing\n        self.image_size = image_size\n        # List to store paths of all image files\n        self.image_files = []\n\n        authentic_dir = self.image_dir / 'authentic' # Path to authentic images\n        forged_dir = self.image_dir / 'forged' # Path to forged images\n\n        if authentic_dir.exists():\n            self.image_files.extend(sorted(list(authentic_dir.glob('*.png'))))\n        if forged_dir.exists():\n            self.image_files.extend(sorted(list(forged_dir.glob('*.png'))))\n\n        if self.transform is None:\n            # Default transform if none is provided\n            self.transform = A.Compose([\n                A.Resize(image_size, image_size), # Resize image\n                A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n                ToTensorV2() # Convert to PyTorch tensor\n            ])\n\n    def __len__(self):\n        return len(self.image_files) # Return the total number of image files\n\n    def __getitem__(self, idx):\n        img_path = self.image_files[idx] # Get image path by index\n        image = cv2.imread(str(img_path)) # Read image using OpenCV\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert image from BGR to RGB\n\n        is_authentic = 'authentic' in str(img_path) # Check if the image is authentic\n\n        if self.mask_dir and not is_authentic:\n            mask_path = self.mask_dir / (img_path.stem + '.npy') # Construct mask path\n            if mask_path.exists():\n                mask = load_mask(str(mask_path)) # Load mask using the load_mask function\n                mask = (mask > 0).astype(np.float32) # Binarize the mask\n            else:\n                # Create an empty mask if no mask file found\n                mask = np.zeros(image.shape[:2], dtype=np.float32)\n        else:\n            # Create an empty mask for authentic images\n            mask = np.zeros(image.shape[:2], dtype=np.float32)\n\n        # Ensure mask has the same spatial dimensions as the image before transform\n        if mask.shape[:2] != image.shape[:2]:\n            mask = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)\n\n        if self.transform:\n            transformed = self.transform(image=image, mask=mask) # Apply transformations\n            image = transformed['image']\n            mask = transformed['mask']\n\n        if len(mask.shape) == 2:\n            mask = mask.unsqueeze(0) # Add a channel dimension to the mask if it's missing\n\n        return image, mask, str(img_path.name) # Return image, mask, and image filename","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:31:54.487258Z","iopub.execute_input":"2025-12-21T20:31:54.487529Z","iopub.status.idle":"2025-12-21T20:31:54.506829Z","shell.execute_reply.started":"2025-12-21T20:31:54.487512Z","shell.execute_reply":"2025-12-21T20:31:54.506195Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TransformDataset(Dataset):\n    \"\"\"Wrapper to apply transforms to subset\"\"\"\n\n    def __init__(self, subset, transform):\n        self.subset = subset # Subset of the original dataset\n        self.transform = transform # Albumentations transform to apply\n\n    def __len__(self):\n        return len(self.subset) # Return the length of the subset\n\n    def __getitem__(self, idx):\n        # Get the original index in the full dataset from the subset\n        img_path = self.subset.dataset.image_files[self.subset.indices[idx]]\n        image = cv2.imread(str(img_path)) # Read image using OpenCV\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert image from BGR to RGB\n\n        is_authentic = 'authentic' in str(img_path) # Check if the image is authentic\n        if not is_authentic and self.subset.dataset.mask_dir:\n            # Construct mask path based on the original dataset's mask directory and image filename\n            mask_path = self.subset.dataset.mask_dir / (img_path.stem + '.npy')\n            if mask_path.exists():\n                # Load mask using the load_mask function\n                mask = load_mask(str(mask_path))\n                # Binarize the mask\n                mask = (mask > 0).astype(np.float32)\n            else:\n                # Create an empty mask if no mask file found\n                mask = np.zeros(image.shape[:2], dtype=np.float32)\n        else:\n            # Create an empty mask for authentic images\n            mask = np.zeros(image.shape[:2], dtype=np.float32)\n\n        if self.transform:\n            # Apply transformations\n            transformed = self.transform(image=image, mask=mask)\n            image = transformed['image']\n            mask = transformed['mask']\n\n        if len(mask.shape) == 2:\n            # Add a channel dimension to the mask if it's missing\n            mask = mask.unsqueeze(0)\n\n        return image, mask, str(img_path.name) # Return image, mask, and image filename","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:31:54.507658Z","iopub.execute_input":"2025-12-21T20:31:54.507918Z","iopub.status.idle":"2025-12-21T20:31:54.524376Z","shell.execute_reply.started":"2025-12-21T20:31:54.507897Z","shell.execute_reply":"2025-12-21T20:31:54.523547Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CopyMoveDataModule(pl.LightningDataModule):\n    \"\"\"PyTorch Lightning DataModule\"\"\"\n\n    def __init__(\n        self,\n        data_dir,\n        train_images_subdir='train_images',\n        train_masks_subdir='train_masks',\n        image_size=512,\n        batch_size=4,\n        num_workers=4,\n        train_val_split=0.8\n    ):\n        super().__init__()\n        self.data_dir = Path(data_dir) # Root directory of the dataset\n        self.train_images_dir = self.data_dir / train_images_subdir\n        self.train_masks_dir = self.data_dir / train_masks_subdir\n        self.image_size = image_size # Target image size\n        self.batch_size = batch_size # Batch size for dataloaders\n        self.num_workers = num_workers # Number of workers for dataloaders\n        # Ratio for splitting training and validation data\n        self.train_val_split = train_val_split\n\n    def setup(self, stage=None):\n        # Create the full dataset\n        full_dataset = CopyMoveDataset(\n            self.train_images_dir,\n            self.train_masks_dir,\n            image_size=self.image_size\n        )\n\n        # Determine the sizes of training and validation sets\n        train_size = int(self.train_val_split * len(full_dataset))\n        val_size = len(full_dataset) - train_size\n        # Get indices for training and validation sets\n        train_indices = list(range(train_size))\n        val_indices = list(range(train_size, len(full_dataset)))\n\n        # Create Subset datasets for training and validation\n        self.train_dataset = TransformDataset(\n            Subset(full_dataset, train_indices),\n            # Apply training transformations\n            get_train_transforms(self.image_size)\n        )\n        self.val_dataset = TransformDataset(\n            Subset(full_dataset, val_indices),\n            # Apply validation transformations\n            get_val_transforms(self.image_size)\n        )\n        print(f\"Train: {len(self.train_dataset)}, Val: {len(self.val_dataset)}\")\n\n    def train_dataloader(self):\n        # Create DataLoader for the training set\n        return DataLoader(\n            self.train_dataset,\n            batch_size=self.batch_size,\n            shuffle=True, # Shuffle training data\n            num_workers=self.num_workers,\n            # Pin memory for faster data transfer to GPU\n            pin_memory=True,\n            # Keep workers alive between epochs\n            persistent_workers=True if self.num_workers > 0 else False\n        )\n\n    def val_dataloader(self):\n        # Create DataLoader for the validation set\n        return DataLoader(\n            self.val_dataset,\n            batch_size=self.batch_size,\n            shuffle=False, # Do not shuffle validation data\n            num_workers=self.num_workers,\n            # Pin memory for faster data transfer to GPU\n            pin_memory=True,\n            # Keep workers alive between epochs\n            persistent_workers=True if self.num_workers > 0 else False\n        )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:31:54.525201Z","iopub.execute_input":"2025-12-21T20:31:54.525548Z","iopub.status.idle":"2025-12-21T20:31:54.543317Z","shell.execute_reply.started":"2025-12-21T20:31:54.525518Z","shell.execute_reply":"2025-12-21T20:31:54.542589Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create data module\ndata_module = CopyMoveDataModule(\n    data_dir=PATH_DATASET,\n    train_images_subdir=TRAIN_IMAGES_SUBDIR,\n    train_masks_subdir=TRAIN_MASKS_SUBDIR,\n    image_size=IMAGE_SIZE,\n    batch_size=BATCH_SIZE,\n    num_workers=NUM_WORKERS,\n    train_val_split=TRAIN_VAL_SPLIT\n)\n\ndata_module.setup()\n\n# You can now access the dataloaders\ntrain_loader = data_module.train_dataloader()\nval_loader = data_module.val_dataloader()\n\nprint(f\"Number of training batches: {len(train_loader)}\")\nprint(f\"Number of validation batches: {len(val_loader)}\")\n\n# Example of getting a batch (optional)\ntrain_images, train_masks, train_filenames = next(iter(train_loader))\nprint(f\"Shape of training images batch: {train_images.shape}\")\nprint(f\"Shape of training masks batch: {train_masks.shape}\")\nprint(f\"Shape of training filenames batch: {len(train_filenames)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:31:54.544120Z","iopub.execute_input":"2025-12-21T20:31:54.544555Z","iopub.status.idle":"2025-12-21T20:31:55.985229Z","shell.execute_reply.started":"2025-12-21T20:31:54.544524Z","shell.execute_reply":"2025-12-21T20:31:55.983670Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_dataloader = data_module.train_dataloader()\nimages, masks, filenames = next(iter(train_dataloader))\n\n# Determine the number of samples to show (up to the batch size or a smaller number)\nnum_samples_to_show = min(4, images.shape[0])\n\n# Display the samples # 2 columns: Image and Mask\nfig, axes = plt.subplots(num_samples_to_show, 2, figsize=(10, num_samples_to_show * 5))\n\nfor i in range(num_samples_to_show):\n    # Denormalize the image for display\n    img = images[i].permute(1, 2, 0).cpu().numpy()\n    mean = np.array([0.485, 0.456, 0.406])\n    std = np.array([0.229, 0.224, 0.225])\n    img = std * img + mean\n    img = np.clip(img, 0, 1) # Clip values to be within [0, 1]\n\n    mask = masks[i].squeeze().cpu().numpy() # Remove channel dimension for display\n\n    # Display the image\n    axes[i, 0].imshow(img)\n    axes[i, 0].set_title(f\"Image: {filenames[i]}\", fontsize=8)\n    axes[i, 0].axis('off')\n\n    # Display the mask\n    axes[i, 1].imshow(mask, cmap='gray') # Use grayscale cmap for binary mask\n    axes[i, 1].set_title(f\"Mask: {filenames[i]}\", fontsize=8)\n    axes[i, 1].axis('off')\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:31:55.986990Z","iopub.execute_input":"2025-12-21T20:31:55.990381Z","iopub.status.idle":"2025-12-21T20:32:01.375322Z","shell.execute_reply.started":"2025-12-21T20:31:55.990330Z","shell.execute_reply":"2025-12-21T20:32:01.371723Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Model architecture\n\nThe model used for copy-move forgery detection is a custom architecture named `CopyMoveHeatmapDetector`. This model is designed to take an image as input and output a heatmap indicating potential forged regions. The architecture consists of the following main components:\n\n1.  **TIMMFeatureExtractor**: This module uses a pre-trained convolutional neural network (specified by the `backbone` parameter, e.g., ResNet50) to extract rich features from the input image. It extracts features from a specific layer of the backbone, which are then normalized.\n\n2.  **BackgroundSuppressor**: This module aims to reduce the influence of background regions on the self-similarity calculation. It can either use a learnable network (`informativeness_net`) or compute feature complexity based on variance, local diversity, and gradient magnitude. The output is a complexity map used to weigh the self-similarity matrix.\n\n3.  **Self-Similarity Computation**: The extracted features are used to compute a self-similarity matrix. This matrix represents the similarity between every pair of feature vectors in the image.\n\n4.  **Spatial Distance Filtering**: A spatial distance mask is applied to the self-similarity matrix. This filter removes matches between feature vectors that are too close to each other in the image, focusing on potential copy-move regions that are spatially separated.\n\n5.  **Forgery Detection**: The filtered similarity matrix is then processed to obtain initial forgery scores, which are reshaped into a heatmap.\n\n6.  **Refinement Module**: A small convolutional network is used to refine the initial heatmap, further enhancing the detection of forged areas.\n\n7.  **Upsampling**: The refined heatmap and the complexity map are upsampled to the desired output size.\n\nHere is a simplified diagram of the architecture:\n\n```\nInput Image\n      |\n      V\nTIMMFeatureExtractor\n      |\n      +-----------------------+\n      V                       V\n  Features           (for learnable suppression)\n      |                       |\n      +-----> BackgroundSuppressor\n      |                       |\n      V                       V\nSelf-Similarity <---+  Complexity Map\n      |               |\n      V               |\nSpatial Distance Filter\n      |\n      V\nFiltered Similarity\n      |\n      V\nForgery Detection\n      |\n      V\nInitial Heatmap\n      |\n      V\nRefinement Module\n      |\n      V\nFinal Heatmap\n```","metadata":{}},{"cell_type":"code","source":"# # Define the model name you want to download\n# model_name_to_download = 'resnet50' # You can change this to other timm models\n\n# # Define the directory and filename for saving the weights\n# save_dir = 'pretrained_models'\n# model_filename = f'{model_name_to_download}_weights.pth'\n# save_path = os.path.join(save_dir, model_filename)\n\n# # Create the directory if it doesn't exist\n# os.makedirs(save_dir, exist_ok=True)\n\n# # Download the pretrained model\n# print(f\"Downloading {model_name_to_download} with pretrained weights...\")\n# model = timm.create_model(model_name_to_download, pretrained=True)\n# print(\"Download successful.\")\n\n# # Save the state dictionary\n# print(f\"Saving model weights to {save_path}...\")\n# torch.save(model.state_dict(), save_path)\n# print(\"Saving complete.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:32:01.389965Z","iopub.execute_input":"2025-12-21T20:32:01.391685Z","iopub.status.idle":"2025-12-21T20:32:01.399370Z","shell.execute_reply.started":"2025-12-21T20:32:01.391644Z","shell.execute_reply":"2025-12-21T20:32:01.398623Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class FeatureExtractor(nn.Module):\n    \"\"\"Feature extractor using TIMM models with optional local weights loading\"\"\"\n\n    def __init__(\n        self,\n        model_name: str = 'resnet50', # Name of the TIMM model to use\n        # Whether to use pretrained weights from timm (ignored if weights_path is provided)\n        pretrained: bool = True,\n        features_only: bool = True, # Whether to return only features\n        out_indices: Tuple[int] = (3,), # Indices of the feature maps to return\n        weights_path: Optional[str] = None # Path to local pretrained weights file\n    ):\n        super().__init__()\n\n        # Determine whether to use timm's pretrained weights based on weights_path\n        use_timm_pretrained = pretrained and (weights_path is None)\n\n        self.backbone = timm.create_model(\n            model_name,\n            # Use timm pretrained only if no weights_path\n            pretrained=use_timm_pretrained,\n            features_only=features_only,\n            out_indices=out_indices\n        )\n\n        if weights_path is not None:\n            # Load weights from the local file if provided\n            print(f\"Loading weights from {weights_path}\")\n            self.backbone.load_state_dict(torch.load(weights_path))\n            print(\"Successfully loaded weights from local file.\")\n\n        # Get feature dimensions\n        # Create a dummy input to get feature dimensions\n        dummy_input = torch.randn(1, 3, 224, 224)\n        # Pass dummy input through the backbone to get feature dimensions\n        # Need to handle potential list/tuple output from features_only=True\n        with torch.no_grad():\n            dummy_features = self.backbone(dummy_input)\n        if isinstance(dummy_features, (list, tuple)):\n             # Assuming the last feature map is the one we'll use\n            self.feature_dim = dummy_features[-1].shape[1]\n        else:\n            self.feature_dim = dummy_features.shape[1]\n\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        features = self.backbone(x) # Extract features using the backbone\n        if isinstance(features, (list, tuple)):\n            features = features[-1] # Use the last feature map if multiple are returned\n        features = F.normalize(features, p=2, dim=1) # Normalize features\n        return features # Return extracted features","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:32:01.400006Z","iopub.execute_input":"2025-12-21T20:32:01.401455Z","iopub.status.idle":"2025-12-21T20:32:01.417201Z","shell.execute_reply.started":"2025-12-21T20:32:01.401430Z","shell.execute_reply":"2025-12-21T20:32:01.416413Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BackgroundSuppressor(nn.Module):\n    \"\"\"Unsupervised background suppression based on feature complexity\"\"\"\n\n    def __init__(self):\n        super().__init__()\n\n    def compute_feature_complexity(self, features: torch.Tensor) -> torch.Tensor:\n        pass\n\n    def forward(\n        self,\n        features: torch.Tensor,\n        similarity_matrix: torch.Tensor\n    ) -> Tuple[torch.Tensor, torch.Tensor]:\n        B, C, H, W = features.shape\n        N = H * W\n        # compute the complexity based on the selected method\n        complexity = self.compute_feature_complexity(features) # Compute complexity\n        complexity_flat = complexity.view(B, 1, N) # Flatten complexity map\n        # Transpose for matrix multiplication\n        complexity_product = torch.bmm(\n            complexity_flat.transpose(1, 2),\n            complexity_flat\n        )\n        # Apply suppression to the similarity matrix\n        suppressed_similarity = similarity_matrix * complexity_product\n        # Return suppressed similarity and complexity map\n        return suppressed_similarity, complexity\n\n\nclass LearnableBackgroundSuppressor(BackgroundSuppressor):\n    \"\"\"Unsupervised background suppression based on feature complexity\"\"\"\n\n    def __init__(self, feature_dim: int):\n        super().__init__()\n        self.informativeness_net = nn.Sequential(\n            nn.Conv2d(feature_dim, feature_dim // 4, 1), # 1x1 convolution to reduce channels\n            nn.BatchNorm2d(feature_dim // 4), # Batch normalization\n            nn.ReLU(inplace=True), # ReLU activation\n            nn.Conv2d(feature_dim // 4, feature_dim // 8, 1), # Another 1x1 convolution\n            nn.BatchNorm2d(feature_dim // 8), # Batch normalization\n            nn.ReLU(inplace=True), # ReLU activation\n            nn.Conv2d(feature_dim // 8, 1, 1), # 1x1 convolution to output a single channel\n            nn.Sigmoid() # Sigmoid activation to get values between 0 and 1\n        )\n\n    def compute_feature_complexity(self, features: torch.Tensor) -> torch.Tensor:\n        return self.informativeness_net(features) # Use learnable network\n\n\nclass StatisticalBackgroundSuppressor(BackgroundSuppressor):\n    \"\"\"Unsupervised background suppression based on feature complexity\"\"\"\n\n    def __init__(self):\n        super().__init__()\n\n    def compute_feature_complexity(self, features: torch.Tensor) -> torch.Tensor:\n        B, C, H, W = features.shape\n\n        # Method 1: Feature variance\n        variance = torch.var(features, dim=1, keepdim=True) # Calculate variance across channels\n\n        # Method 2: Local feature diversity\n        kernel_size = 3\n        padding = kernel_size // 2\n        unfold = nn.Unfold(kernel_size=kernel_size, padding=padding) # Unfold features into patches\n        patches = unfold(features)\n        patches = patches.view(B, C, kernel_size * kernel_size, H * W) # Reshape patches\n        local_diversity = torch.var(patches, dim=2) # Calculate variance within patches\n        local_diversity = local_diversity.view(B, C, H, W) # Reshape back to spatial dimensions\n        local_diversity = torch.mean(local_diversity, dim=1, keepdim=True) # Average across channels\n\n        # Method 3: Gradient magnitude\n        sobel_x = torch.tensor(\n            [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=features.dtype, device=features.device)\n        sobel_y = sobel_x.t()\n        # Convolve with Sobel filter for x-gradients\n        grad_x = F.conv2d(\n            features, sobel_x.view(1, 1, 3, 3).repeat(C, 1, 1, 1), padding=1, groups=C)\n        # Convolve with Sobel filter for y-gradients\n        grad_y = F.conv2d(\n            features, sobel_y.view(1, 1, 3, 3).repeat(C, 1, 1, 1), padding=1, groups=C)\n        gradient_mag = torch.sqrt(grad_x ** 2 + grad_y ** 2) # Calculate gradient magnitude\n        gradient_mag = torch.mean(gradient_mag, dim=1, keepdim=True) # Average across channels\n\n        # Combine all metrics\n        # Average the complexity metrics\n        complexity = (variance + local_diversity + gradient_mag) / 3.0\n        # Normalize to [0, 1]\n        complexity = (complexity - complexity.min()) / (complexity.max() - complexity.min() + 1e-8)\n        return complexity","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:32:01.418085Z","iopub.execute_input":"2025-12-21T20:32:01.418355Z","iopub.status.idle":"2025-12-21T20:32:01.445440Z","shell.execute_reply.started":"2025-12-21T20:32:01.418332Z","shell.execute_reply":"2025-12-21T20:32:01.444688Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CopyMoveHeatmapDetector(nn.Module):\n    \"\"\"Complete copy-move forgery detection network with heatmap output\"\"\"\n\n    def __init__(\n        self,\n        backbone: str = 'resnet50', # Name of the backbone model (e.g., resnet50)\n        pretrained: bool = True, # Whether to use a pretrained backbone\n        output_size: Tuple[int, int] = (512, 512), # Desired output size of the heatmap\n        # Minimum spatial distance between potential copy-move regions\n        distance_threshold: int = 3,\n        # Whether to use a learnable background suppressor\n        learnable_bg_suppression: bool = False,\n        # Path to local pretrained weights file for the backbone\n        weights_path: Optional[str] = None\n    ):\n        super().__init__()\n\n        self.output_size = output_size\n        self.distance_threshold = distance_threshold\n\n        self.feature_extractor = FeatureExtractor(\n            model_name=backbone,\n            pretrained=pretrained,\n            features_only=True,\n            # Extract features from the last stage of the backbone\n            out_indices=(3,),\n            # Pass the weights_path to the feature extractor\n            weights_path=weights_path\n        )\n        feature_dim=self.feature_extractor.feature_dim\n        if learnable_bg_suppression:\n            self.bg_suppressor = LearnableBackgroundSuppressor(feature_dim)\n        else:\n            self.bg_suppressor = StatisticalBackgroundSuppressor()\n\n        # Refinement module to process the initial heatmap\n        self.refine = nn.Sequential(\n            nn.Conv2d(1, 32, 3, padding=1), # Initial convolution\n            nn.BatchNorm2d(32), # Batch normalization\n            nn.ReLU(inplace=True), # ReLU activation\n            nn.Conv2d(32, 16, 3, padding=1), # Second convolution\n            nn.BatchNorm2d(16), # Batch normalization\n            nn.ReLU(inplace=True), # ReLU activation\n            nn.Conv2d(16, 1, 1), # Final 1x1 convolution to output a single channel\n            # Removed Sigmoid here as BCEWithLogitsLoss is used\n        )\n\n    def compute_self_similarity(\n        self,\n        features: torch.Tensor\n    ) -> Tuple[torch.Tensor, Tuple[int, int]]:\n        B, C, H, W = features.shape\n        N = H * W\n\n        features_flat = features.view(B, C, N) # Flatten spatial dimensions\n        # Compute self-similarity matrix\n        similarity = torch.bmm(\n            features_flat.transpose(1, 2), # Transpose to get shape (B, N, C)\n            features_flat # Shape (B, C, N)\n        )\n        return similarity, (H, W) # Return similarity matrix and spatial shape\n\n    def create_spatial_distance_mask(\n        self,\n        H: int,\n        W: int,\n        device: torch.device,\n        threshold: int\n    ) -> torch.Tensor:\n        # Create coordinate grid\n        y_coords = torch.arange(H, device=device).view(-1, 1).repeat(1, W)\n        x_coords = torch.arange(W, device=device).view(1, -1).repeat(H, 1)\n        coords = torch.stack([y_coords, x_coords], dim=-1).view(-1, 2).float() # Flatten coordinates\n        # Compute pairwise spatial distances\n        spatial_dist = torch.cdist(coords, coords, p=2)\n        # Create mask where distance is greater than threshold\n        mask = (spatial_dist > threshold).float()\n        return mask\n\n    def detect_copy_move(\n        self,\n        similarity_matrix: torch.Tensor,\n        spatial_shape: Tuple[int, int]\n    ) -> torch.Tensor:\n        B, N, _ = similarity_matrix.shape\n        H, W = spatial_shape\n        # Create and expand spatial mask to the batch size\n        spatial_mask = self.create_spatial_distance_mask(\n            H, W,\n            similarity_matrix.device,\n            self.distance_threshold\n        )\n        spatial_mask = spatial_mask.unsqueeze(0).expand(B, -1, -1)\n        # Apply spatial distance filter to similarity matrix\n        filtered_similarity = similarity_matrix * spatial_mask\n        # Get the maximum similarity score for each pixel (potential forgery score)\n        forgery_scores, _ = torch.max(filtered_similarity, dim=2)\n        # Reshape forgery scores back to spatial dimensions (heatmap)\n        forgery_heatmap = forgery_scores.view(B, H, W).unsqueeze(1)\n        return forgery_heatmap\n\n    def forward(\n        self,\n        x: torch.Tensor\n    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:\n        original_size = x.shape[2:] # Store original image size\n        # Extract features\n        features = self.feature_extractor(x)\n        # Compute self-similarity\n        similarity, spatial_shape = self.compute_self_similarity(features)\n        # Apply background suppression\n        suppressed_similarity, complexity = self.bg_suppressor(features, similarity)\n        # Detect copy-move regions\n        forgery_heatmap = self.detect_copy_move(suppressed_similarity, spatial_shape)\n\n        # Upsample initial heatmap for refinement\n        heatmap = F.interpolate(\n            forgery_heatmap,\n            # Upsample to half original size\n            size=(original_size[0] // 2, original_size[1] // 2),\n            mode='bilinear',\n            align_corners=False\n        )\n        heatmap = self.refine(heatmap) # Refine the heatmap\n\n        # Final upsampling to the target output size\n        heatmap = F.interpolate(\n            heatmap,\n            size=self.output_size,\n            mode='bilinear',\n            align_corners=False\n        )\n\n        # Upsample complexity map to the target output size\n        complexity_upsampled = F.interpolate(\n            complexity,\n            size=self.output_size,\n            mode='bilinear',\n            align_corners=False\n        )\n        # Return final heatmap and upsampled complexity map\n        return heatmap, complexity_upsampled","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:32:01.446465Z","iopub.execute_input":"2025-12-21T20:32:01.446963Z","iopub.status.idle":"2025-12-21T20:32:01.466784Z","shell.execute_reply.started":"2025-12-21T20:32:01.446935Z","shell.execute_reply":"2025-12-21T20:32:01.466083Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CopyMoveLightningModule(pl.LightningModule):\n    \"\"\"PyTorch Lightning module for copy-move detection\"\"\"\n\n    def __init__(\n        self,\n        backbone='resnet50', # Name of the backbone model\n        pretrained=False,\n        image_size=512, # Target image size\n        learnable_bg_suppression=True,\n        distance_threshold=3,\n        learning_rate=1e-4, # Learning rate for the optimizer\n        weight_decay=1e-4, # Weight decay for the optimizer\n        alpha=0.5, # Weight for Dice loss\n        beta=0.1, # Weight for complexity regularization\n        weights_path: Optional[str] = None # Path to local pretrained weights file\n    ):\n        super().__init__()\n        self.save_hyperparameters() # Saves all __init__ parameters as hyperparameters\n\n        self.model = CopyMoveHeatmapDetector(\n            backbone=backbone,\n            pretrained=pretrained,\n            output_size=(image_size, image_size),\n            distance_threshold=distance_threshold,\n            learnable_bg_suppression=learnable_bg_suppression,\n            # Pass the weights_path to the detector\n            weights_path=weights_path,\n        )\n\n        self.learning_rate = learning_rate\n        self.weight_decay = weight_decay\n        self.alpha = alpha\n        self.beta = beta\n\n        # Define metrics using torchmetrics\n        metrics = MetricCollection({\n            'precision': BinaryPrecision(),\n            'recall': BinaryRecall(),\n            'f1': BinaryF1Score(),\n            'iou': BinaryJaccardIndex(),\n            'accuracy': BinaryAccuracy()\n        })\n\n        self.train_metrics = metrics.clone(prefix='train/') # Metrics for training\n        self.val_metrics = metrics.clone(prefix='val/') # Metrics for validation\n\n    def forward(self, x):\n        return self.model(x) # Pass input through the detection model\n\n    def dice_loss(self, pred, target, smooth=1.0):\n        # Apply sigmoid to predictions for Dice loss since the model output is now logits\n        pred_sigmoid = torch.sigmoid(pred)\n        pred_flat = pred_sigmoid.contiguous().view(-1)\n        target_flat = target.contiguous().view(-1)\n        intersection = (pred_flat * target_flat).sum()\n        dice = (2. * intersection + smooth) / (pred_flat.sum() + target_flat.sum() + smooth)\n        # Return the Dice loss\n        return 1 - dice\n\n    def compute_loss(self, pred_heatmap, gt_mask, complexity_map):\n        # Use BCEWithLogitsLoss for stability with autocasting\n        bce_loss = F.binary_cross_entropy_with_logits(pred_heatmap, gt_mask.float())\n        dice_loss = self.dice_loss(pred_heatmap, gt_mask.float())\n\n        main_loss = self.alpha * dice_loss + (1 - self.alpha) * bce_loss # Combined main loss\n\n        if complexity_map is not None:\n            # Apply sigmoid to prediction heatmap for complexity regularization\n            pred_heatmap_sigmoid = torch.sigmoid(pred_heatmap)\n            # Complexity regularization loss\n            complexity_reg = F.mse_loss(complexity_map * gt_mask, gt_mask.float())\n            # Add complexity regularization to main loss\n            main_loss = main_loss + self.beta * complexity_reg\n\n        # Return total loss, Dice loss, and BCE loss\n        return main_loss, dice_loss, bce_loss\n\n    def training_step(self, batch, batch_idx):\n        images, masks, _ = batch # Unpack the batch\n        pred_heatmap, complexity_map = self(images) # Forward pass\n\n        loss, dice, bce = self.compute_loss(pred_heatmap, masks, complexity_map) # Compute loss\n\n        # Apply sigmoid for metric calculation and logging\n        pred_binary = (torch.sigmoid(pred_heatmap) > 0.5).long() # Binarize predictions\n        target_binary = masks.long() # Convert target mask to long\n        metrics = self.train_metrics(pred_binary, target_binary) # Compute training metrics\n\n        self.log('train/loss', loss, prog_bar=True, on_step=True) # Log training loss\n        self.log('train/dice_loss', dice, on_step=True) # Log training Dice loss\n        self.log('train/bce_loss', bce, on_step=True) # Log training BCE loss\n        self.log_dict(metrics, prog_bar=True, on_step=True) # Log training metrics\n        return loss # Return the loss\n\n    def validation_step(self, batch, batch_idx):\n        images, masks, _ = batch # Unpack the batch\n        pred_heatmap, complexity_map = self(images) # Forward pass\n\n        loss, dice, bce = self.compute_loss(pred_heatmap, masks, complexity_map) # Compute loss\n\n        # Apply sigmoid for metric calculation and logging\n        pred_binary = (torch.sigmoid(pred_heatmap) > 0.5).long() # Binarize predictions\n        target_binary = masks.long() # Convert target mask to long\n        metrics = self.val_metrics(pred_binary, target_binary) # Compute validation metrics\n\n        self.log('val/loss', loss, prog_bar=True, on_step=True) # Log validation loss\n        self.log('val/dice_loss', dice, on_step=True) # Log validation Dice loss\n        self.log('val/bce_loss', bce, on_step=True) # Log validation BCE loss\n        self.log_dict(metrics, prog_bar=True, on_step=True) # Log validation metrics\n        return loss # Return the loss\n\n    def on_train_epoch_end(self):\n        self.train_metrics.reset() # Reset training metrics at the end of each epoch\n\n    def on_validation_epoch_end(self):\n        self.val_metrics.reset() # Reset validation metrics at the end of each epoch\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.AdamW(\n            self.parameters(),\n            lr=self.learning_rate,\n            weight_decay=self.weight_decay\n        ) # AdamW optimizer\n\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n            optimizer,\n            T_max=self.trainer.max_epochs,\n            eta_min=1e-6\n        ) # Cosine Annealing learning rate scheduler\n\n        return {\n            'optimizer': optimizer,\n            'lr_scheduler': {\n                'scheduler': scheduler,\n                'interval': 'epoch'\n            }\n        } # Return optimizer and learning rate scheduler configuration","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:32:01.467586Z","iopub.execute_input":"2025-12-21T20:32:01.467780Z","iopub.status.idle":"2025-12-21T20:32:01.490347Z","shell.execute_reply.started":"2025-12-21T20:32:01.467765Z","shell.execute_reply":"2025-12-21T20:32:01.489559Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create model\nmodel = CopyMoveLightningModule(\n    backbone=BACKBONE,\n    image_size=IMAGE_SIZE,\n    learnable_bg_suppression=LEARN_BG_SUPPRESSION,\n    learning_rate=LEARNING_RATE,\n    weight_decay=WEIGHT_DECAY,\n    weights_path=f\"{PATH_MODELS}/{BACKBONE}_weights.pth\",\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:32:01.491300Z","iopub.execute_input":"2025-12-21T20:32:01.491555Z","iopub.status.idle":"2025-12-21T20:32:02.399498Z","shell.execute_reply.started":"2025-12-21T20:32:01.491531Z","shell.execute_reply":"2025-12-21T20:32:02.398691Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"code","source":"# Setup callbacks for the PyTorch Lightning trainer\ncheckpoint_callback = ModelCheckpoint(\n    # Filename format for saved checkpoints\n    filename='best-{epoch:03d}-{val/iou:.4f}',\n    # Metric to monitor for saving the best model\n    monitor='val/iou',\n    # Save checkpoint when the monitored metric is maximized\n    mode='max',\n    # Save the top 3 models based on the monitored metric\n    save_top_k=3,\n    # Save the last checkpoint\n    save_last=True\n)\n\n# Metric to monitor for early stopping\nearly_stop_callback = EarlyStopping(\n    # Number of epochs with no improvement after which training will be stopped\n    monitor='val/iou',\n    patience=10,\n    # Stop training when the monitored metric is maximized\n    mode='max',\n    # Print message when early stopping is triggered\n    verbose=True\n)\n\n# Log the learning rate at the end of each epoch\nlr_monitor = LearningRateMonitor(\n    logging_interval='epoch',\n)\n# Logger to save training logs in CSV format\nlogger = CSVLogger(\n    save_dir='logs',\n    name='copy_move',\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:32:02.400344Z","iopub.execute_input":"2025-12-21T20:32:02.400637Z","iopub.status.idle":"2025-12-21T20:32:02.424073Z","shell.execute_reply.started":"2025-12-21T20:32:02.400600Z","shell.execute_reply":"2025-12-21T20:32:02.423499Z"},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"Backbone: {BACKBONE} | Epochs: {NUM_EPOCHS} | Batch: {BATCH_SIZE}\")\n\n# Create trainer\ntrainer = pl.Trainer(\n    # Set the maximum number of training epochs\n    max_epochs=NUM_EPOCHS,\n    # Automatically select the accelerator (GPU, CPU, etc.)\n    accelerator='auto',\n    # Add defined callbacks\n    callbacks=[checkpoint_callback, early_stop_callback, lr_monitor],\n    # Use the configured logger for logging training progress\n    logger=logger,\n    # Log training metrics every 10 steps\n    log_every_n_steps=10,\n    # Use mixed precision training on GPU\n    precision='16-mixed' if DEVICE == 'cuda' else 32,\n    # Clip gradients to prevent exploding gradients\n    gradient_clip_val=1.0,\n    # Accumulate gradients over 6 batches before updating model weights\n    accumulate_grad_batches=6,\n)\n\n# Train\ntrainer.fit(model, data_module) # Start the training process\nprint(f\"✓ Training complete. Best model: {checkpoint_callback.best_model_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-21T20:32:02.424739Z","iopub.execute_input":"2025-12-21T20:32:02.424965Z"},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import seaborn as sns\nimport pandas as pd\nsns.set()\n\n# Read the metrics.csv using the trainer's logger directory\nmetrics = pd.read_csv(f\"{trainer.logger.log_dir}/metrics.csv\")\n\n# Remove the step column and set epoch as index\n# metrics.set_index(\"step\", inplace=True)\ndisplay(metrics.dropna(axis=1, how=\"all\").head())\n\n# Melt the DataFrame to long-form for plotting\nmetrics_melted = metrics.reset_index().melt(id_vars='epoch', var_name='metric', value_name='value')\n\n# Define metric groups\nmetric_groups = {\n    'Loss': [c for c in metrics.columns if \"loss\" in c],\n    'Metrics': [c for c in metrics.columns if any(m in c for m in [\"precision\", \"recall\", \"f1\", \"iou\", \"accuracy\"])],\n}\n\n\n# Plot metrics for each group in a separate chart\nfor title, metric_list in metric_groups.items():\n    # Filter melted DataFrame for the current group\n    group_metrics = metrics_melted[metrics_melted['metric'].isin(metric_list)]\n\n    plt.figure(figsize=(10, 5))\n    sns.lineplot(data=group_metrics, x='epoch', y='value', hue='metric')\n    plt.title(f'{title} over Epochs', fontsize=14, fontweight='bold')\n    plt.xlabel('Epoch', fontsize=12)\n    plt.ylabel(title, fontsize=12)\n    plt.grid(True, alpha=0.3)\n    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left') # Move legend outside and to the right\n    # plt.yscale('log')  # Set y-axis to logarithmic scale\nplt.show()","metadata":{"trusted":true,"_kg_hide-output":false,"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Inference","metadata":{}},{"cell_type":"code","source":"# Load best model\nbest_checkpoint = checkpoint_callback.best_model_path\nmodel = CopyMoveHeatmapDetector(\n    backbone=BACKBONE,\n    pretrained=False,\n    output_size=(IMAGE_SIZE, IMAGE_SIZE),\n    distance_threshold=3,\n    learnable_bg_suppression=LEARN_BG_SUPPRESSION\n)\n\ncheckpoint = torch.load(best_checkpoint, map_location=DEVICE)\nstate_dict = {k.replace('model.', ''): v for k, v in checkpoint['state_dict'].items()}\nmodel.load_state_dict(state_dict)\nmodel = model.to(DEVICE)\nmodel.eval()","metadata":{"trusted":true,"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def post_process_mask(mask, kernel_size=5, min_area_percentage=0.01):\n    \"\"\"\n    Applies morphological closing and filters small objects from a binary mask.\n\n    Args:\n        mask (np.ndarray): Binary mask (uint8).\n        kernel_size (int): Size of the kernel for morphological operations. If None, no morphological closing is applied.\n        min_area_percentage (float): Minimum area of connected components as a percentage of total image area. If None, no area filtering is applied.\n\n    Returns:\n        np.ndarray: Processed binary mask.\n    \"\"\"\n    processed_mask = mask.copy()\n\n    # Apply morphological closing if kernel_size is provided\n    if kernel_size is not None and kernel_size > 0:\n        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))\n        processed_mask = cv2.morphologyEx(processed_mask, cv2.MORPH_CLOSE, kernel)\n\n    # Filter small objects if min_area_percentage is provided\n    if min_area_percentage is not None and min_area_percentage > 0:\n        # Find connected components\n        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(processed_mask, 8, cv2.CV_32S)\n\n        # Calculate the minimum area in pixels based on the percentage\n        total_image_area = mask.shape[0] * mask.shape[1]\n        min_area_pixels = total_image_area * min_area_percentage\n\n        # Create a new mask with only large components\n        filtered_mask = np.zeros_like(mask, dtype=np.uint8)\n        for i in range(1, num_labels):  # Start from 1 to exclude the background\n            if stats[i, cv2.CC_STAT_AREA] >= min_area_pixels:\n                filtered_mask[labels == i] = 255 # Keep components larger than min_area\n        processed_mask = filtered_mask\n\n    # Ensure the output is binary (0 or 1)\n    return (processed_mask > 0).astype(np.uint8)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create output directories\noutput_dir = Path(OUTPUT_DIR)\noutput_dir.mkdir(exist_ok=True, parents=True)\nmasks_dir = output_dir / 'masks'\nheatmaps_dir = output_dir / 'heatmaps'\nmasks_dir.mkdir(exist_ok=True)\nheatmaps_dir.mkdir(exist_ok=True)\n\n# Prepare transform\ntransform = get_val_transforms(IMAGE_SIZE)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Validation Set","metadata":{}},{"cell_type":"code","source":"# Prepare validation DataLoader\nval_loader = data_module.val_dataloader()\n\n# Create output directories for validation results if they don't exist\nval_output_dir = Path(OUTPUT_DIR) / 'validation'\nval_masks_dir = val_output_dir / 'masks'\nval_heatmaps_dir = val_output_dir / 'heatmaps'\nval_output_dir.mkdir(exist_ok=True, parents=True)\nval_masks_dir.mkdir(exist_ok=True)\nval_heatmaps_dir.mkdir(exist_ok=True)\n\n# Prepare transform (using validation transforms)\ntransform = get_val_transforms(IMAGE_SIZE)\n\n# List to store results for threshold optimization\nthreshold_results = []\n\n# Define a range of thresholds to test\nthresholds_to_test = np.linspace(0, 1, 51) # Test 51 thresholds from 0 to 1\n\n# Initialize metrics for threshold optimization\nthreshold_metrics = MetricCollection({\n    'precision': BinaryPrecision(),\n    'recall': BinaryRecall(),\n    'f1': BinaryF1Score(),\n    'iou': BinaryJaccardIndex(),\n    'accuracy': BinaryAccuracy()\n}).to(DEVICE)\n\nprint(\"Optimizing threshold on Validation Set...\")\n\nfor threshold in tqdm(thresholds_to_test, desc=\"Testing Thresholds\"):\n    all_preds = []\n    all_targets = []\n\n    for images, masks, filenames in val_loader:\n        with torch.no_grad():\n            images = images.to(DEVICE)\n            pred_heatmap_tensor, _ = model(images)\n\n        # Apply the current threshold to the heatmap\n        binary_mask_tensor = (torch.sigmoid(pred_heatmap_tensor) > threshold).long()\n\n        all_preds.append(binary_mask_tensor.cpu())\n        all_targets.append(masks.long().cpu())\n\n    # Concatenate results from all batches\n    all_preds = torch.cat(all_preds, dim=0).view(-1)\n    all_targets = torch.cat(all_targets, dim=0).view(-1)\n\n    # Compute metrics for the current threshold\n    current_metrics = threshold_metrics(all_preds.to(DEVICE), all_targets.to(DEVICE))\n\n    threshold_results.append({\n        'threshold': threshold,\n        'precision': current_metrics['precision'].item(),\n        'recall': current_metrics['recall'].item(),\n        'f1': current_metrics['f1'].item(),\n        'iou': current_metrics['iou'].item(),\n        'accuracy': current_metrics['accuracy'].item()\n    })\n    threshold_metrics.reset() # Reset metrics for the next threshold\n\n# Convert results to DataFrame and find the best threshold\nthreshold_df = pd.DataFrame(threshold_results)\n# Find the threshold that maximizes IoU (or another metric like F1)\nbest_threshold_row = threshold_df.loc[threshold_df['iou'].idxmax()]\nbest_threshold = best_threshold_row['threshold']\n\nprint(f\"\\nBest Threshold found based on IoU: {best_threshold:.4f}\")\nprint(\"Metrics at best threshold:\")\nprint(best_threshold_row)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Plot metrics vs. threshold using pandas plotting\nplt.figure(figsize=(12, 6))\n\n# Select the metric columns and plot them directly from the DataFrame\nthreshold_df.plot(x='threshold', y=['precision', 'recall', 'f1', 'iou', 'accuracy'], ax=plt.gca())\n\nplt.title('Validation Metrics vs. Threshold', fontsize=14, fontweight='bold')\nplt.xlabel('Threshold', fontsize=12)\nplt.ylabel('Score', fontsize=12)\nplt.grid(True, alpha=0.3)\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Now, run inference again using the best threshold to save masks and heatmaps\nresults = []\ndisplay_samples = []\nnum_display_samples = 12 # Number of samples to display\n\nfor images, masks, filenames in tqdm(val_loader, desc=\"Predicting on Validation Set (Best Threshold)\"):\n    with torch.no_grad():\n        images = images.to(DEVICE)\n        pred_heatmap_tensor, _ = model(images)\n\n    for i in range(images.shape[0]):\n        img = images[i].cpu().numpy().transpose(1, 2, 0)\n        # Denormalize the image for display\n        mean = np.array([0.485, 0.456, 0.406])\n        std = np.array([0.229, 0.224, 0.225])\n        img = std * img + mean\n        img = np.clip(img, 0, 1) # Clip values to be within [0, 1]\n\n        gt_mask = masks[i].squeeze().cpu().numpy()\n        heatmap = pred_heatmap_tensor[i, 0].cpu().numpy()\n\n        # Resize heatmap to original image size for visualization and post-processing\n        original_size = img.shape[:2]\n        heatmap_resized = cv2.resize(heatmap, (original_size[1], original_size[0]))\n        # Apply the best threshold to get the binary mask (0/255 for post_processing)\n        binary_mask = (heatmap_resized > best_threshold).astype(np.uint8) * 255\n        # Apply post-processing\n        processed_mask = post_process_mask(binary_mask, kernel_size=None, min_area_percentage=None)\n\n        # Store sample for display\n        if len(display_samples) < num_display_samples:\n            display_samples.append({\n                'image': img,\n                'gt_mask': gt_mask,\n                'heatmap': heatmap_resized,\n                'predicted_mask': processed_mask,\n                'filename': filenames[i]\n            })\n\n        # Save heatmap and mask\n        filename_stem = Path(filenames[i]).stem\n        np.save(val_heatmaps_dir / f\"{filename_stem}.npy\", heatmap_resized)\n        np.save(val_masks_dir / f\"{filename_stem}.npy\", processed_mask)\n\n        results.append({\n            'image_name': filenames[i],\n            'has_forgery': processed_mask.sum() > 0,\n            'forgery_percentage': (processed_mask.sum() / (processed_mask.size * 255)) * 100\n        })\n\nresults_df = pd.DataFrame(results)\nresults_df.to_csv(val_output_dir / 'validation_predictions_summary.csv', index=False)\n\n# You can now use the best_threshold for inference on the test set.\n# Update the THRESHOLD constant or pass best_threshold to the test inference loop.\nTHRESHOLD = best_threshold\nprint(f\"Updated global THRESHOLD to: {THRESHOLD:.4f}\")\nprint(f\"\\n✓ Forged: {results_df['has_forgery'].sum()}, Authentic: {(~results_df['has_forgery']).sum()}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"✓ Validation Inference complete. Forged: {results_df['has_forgery'].sum()}, Authentic: {(~results_df['has_forgery']).sum()}\")\n\nfor sample in display_samples:\n    fig, axes = plt.subplots(1, 3, figsize=(18, 6)) # 3 columns: GT Mask, Image with Contours, Predicted Mask\n\n    # Display Ground Truth Mask\n    axes[0].imshow(sample['gt_mask'], cmap='gray')\n    axes[0].set_title(f\"Ground Truth Mask\\n{sample['filename']}\", fontsize=10)\n    axes[0].axis('off')\n\n    # Display Original Image with GT and Predicted Mask Contours\n    axes[1].imshow(sample['image'])\n    # Draw GT contours using contour\n    axes[1].contour(sample['gt_mask'], levels=[0.5], colors=['green'], linewidths=1) # Green for GT\n    # Draw Predicted contours using contour\n    axes[1].contour(sample['predicted_mask'], levels=[128], colors=['red'], linewidths=1) # Red for Predicted (assuming 0/255 mask)\n    axes[1].set_title(f\"Image with Contours (GT: Green, Pred: Red)\\n{sample['filename']}\", fontsize=10)\n    axes[1].axis('off')\n\n    # Display Predicted Binary Mask\n    axes[2].imshow(sample['predicted_mask'], cmap='gray')\n    axes[2].set_title(\"Predicted Binary Mask\", fontsize=10)\n    axes[2].axis('off')\n\n    plt.tight_layout()\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Test Set","metadata":{}},{"cell_type":"code","source":"# Predict on test set\ntest_img_dir = Path(TEST_IMAGES_DIR)\nprint(f\"Checking test image directory: {test_img_dir}\")\ntest_image_files = sorted(list(test_img_dir.glob('*.png')))\nprint(f\"Found {len(test_image_files)} test image files.\")\nresults = []\n\nprint(f\"Running Inference on {len(list(Path(TEST_IMAGES_DIR).glob('*.png')))} images...\")\n\nfor img_path in tqdm(test_image_files, desc=\"Predicting\"):\n    image = cv2.imread(str(img_path))\n    original_size = image.shape[:2]\n    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n    transformed = transform(image=image_rgb)\n    image_tensor = transformed['image'].unsqueeze(0).to(DEVICE)\n\n    with torch.no_grad():\n        heatmap_tensor, _ = model(image_tensor)\n\n    heatmap = heatmap_tensor[0, 0].cpu().numpy()\n    heatmap = cv2.resize(heatmap, (original_size[1], original_size[0]))\n    binary_mask = (heatmap > THRESHOLD).astype(np.uint8) * 255 # Convert to 0/255 for post_processing\n\n    # Apply post-processing\n    # Using 0.01% of the image area as the minimum area threshold\n    processed_mask = post_process_mask(binary_mask, min_area_percentage=0.0005)\n\n    np.save(heatmaps_dir / f\"{img_path.stem}.npy\", heatmap)\n    np.save(masks_dir / f\"{img_path.stem}.npy\", processed_mask) # Save processed mask\n\n    results.append({\n        'image_name': img_path.name,\n        'has_forgery': processed_mask.sum() > 0,\n        'forgery_percentage': (processed_mask.sum() / (processed_mask.size * 255)) * 100 # Adjust calculation for 0/255 mask\n    })\n\nresults_df = pd.DataFrame(results)\nresults_df.to_csv(output_dir / 'predictions_summary.csv', index=False)\n\nprint(f\"✓ Inference complete. Forged: {results_df['has_forgery'].sum()}, Authentic: {(~results_df['has_forgery']).sum()}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Assuming inference has been run and outputs are saved in OUTPUT_DIR\noutput_dir = Path(OUTPUT_DIR)\nmasks_dir = output_dir / 'masks'\nheatmaps_dir = output_dir / 'heatmaps'\ntest_img_dir = Path(TEST_IMAGES_DIR)\n\n# Get a list of test image files and corresponding saved mask/heatmap files\ntest_image_files = sorted(list(test_img_dir.glob('*.png')))\nsaved_mask_files = sorted(list(masks_dir.glob('*.npy')))\nsaved_heatmap_files = sorted(list(heatmaps_dir.glob('*.npy')))\n\n# Create a list of tuples (image_path, heatmap_path, mask_path) for files that exist\ndisplay_files = []\nfor img_path in test_image_files:\n    stem = img_path.stem\n    heatmap_path = heatmaps_dir / f\"{stem}.npy\"\n    mask_path = masks_dir / f\"{stem}.npy\"\n    if heatmap_path.exists() and mask_path.exists():\n        display_files.append((img_path, heatmap_path, mask_path))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Determine the number of samples to show\nnum_samples_to_show = min(5, len(display_files)) # Show up to 5 samples\n\n# Select random samples\nrandom_display_files = random.sample(display_files, num_samples_to_show)\n\n# Display predictions for selected samples\nfor img_path, heatmap_path, mask_path in random_display_files:\n    image = cv2.imread(str(img_path))\n    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n    # Load saved heatmap and mask\n    heatmap_resized = np.load(heatmap_path)\n    binary_mask = np.load(mask_path)\n\n    # Create a figure with 3 subplots: Original Image, Heatmap, Binary Mask\n    fig, axes = plt.subplots(1, 3, figsize=(18, 6)) # Adjust figsize as needed\n\n    # Display Original Image\n    axes[0].imshow(image_rgb)\n    axes[0].set_title(f\"Original Image\\n{img_path.name}\", fontsize=10)\n    axes[0].axis('off')\n\n    # Display Heatmap\n    heatmap_display = axes[1].imshow(heatmap_resized, cmap='viridis', vmin=0, vmax=1)\n    axes[1].set_title(\"Predicted Heatmap\", fontsize=10)\n    axes[1].axis('off')\n    fig.colorbar(heatmap_display, ax=axes[1]) # Add colorbar for heatmap\n\n    # Display Binary Mask\n    axes[2].imshow(binary_mask, cmap='gray') # Use grayscale cmap for binary mask\n    axes[2].set_title(\"Predicted Binary Mask\", fontsize=10)\n    axes[2].axis('off')\n\n    plt.tight_layout()\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Submission\n\nYou must create a row in the submission file for each image. For images that do not contain a copy-move forgery predict the string \"authentic\". For all others submit run length encoded masks as serialized using the rle_encode function in the metric. The file should contain a header and have the following format:\n\n```\ncase_id,annotation\n1,authentic\n2,\"[123, 4]\"\n```","metadata":{}},{"cell_type":"code","source":"def rle_encode(mask, fg_val=1):\n    \"\"\"Convert binary mask to RLE using the competition metric format\"\"\"\n    dots = np.where(mask.T.flatten() == fg_val)[0]\n    run_lengths = []\n    prev = -2\n    \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    \n    return run_lengths","metadata":{"trusted":true,"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Assuming inference has been run and outputs are saved in OUTPUT_DIR\noutput_dir = Path(OUTPUT_DIR)\nmasks_dir = output_dir / 'masks'\ntest_img_dir = Path(TEST_IMAGES_DIR) # Need test_img_dir to get all image names\n\n# Get a list of all test image filenames (without extension)\ntest_image_stems = sorted([img_path.stem for img_path in test_img_dir.glob('*.png')])\n\nsubmissions = []\n\nfor img_stem in tqdm(test_image_stems, desc=\"Creating submission\"):\n    mask_path = masks_dir / f\"{img_stem}.npy\"\n\n    if mask_path.exists():\n        mask = np.load(mask_path)\n        rle_annotation = rle_encode(mask)\n        # print(rle_annotation)\n        if rle_annotation: # If any pixels are marked as forgery\n            # Ensure the annotation is correctly formatted with quotes if it's an RLE string\n            submissions.append({'case_id': img_stem, 'annotation': f'\"{repr(rle_annotation)}\"'})\n        else: # For authentic images, the annotation should be 'authentic' without quotes\n            submissions.append({'case_id': img_stem, 'annotation': 'authentic'})\n    else:\n        # If no mask file exists (e.g., for an authentic image if they were in test set)\n        # Based on previous EDA, test images are all forged according to train_masks count.\n        # However, for robustness, we can assume no mask means authentic if needed.\n        # For this competition, test images should have corresponding masks.\n        print(f\"Warning: No mask found for test image {img_stem}. Assuming authentic.\")\n        submissions.append({'case_id': img_stem, 'annotation': 'authentic'})\n\nsubmission_df = pd.DataFrame(submissions)","metadata":{"trusted":true,"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from pprint import pprint\n\ndef write_submission_csv(data, filename):\n    \"\"\"Writes submission data to a CSV file using writelines by generating lines upfront.\"\"\"\n    lines = [\"case_id,annotation\\n\"]  # Header line\n    for row in data:\n        annotation = row['annotation']\n        lines.append(f\"{row['case_id']},{annotation}\\n\")\n    # pprint(lines)\n    with open(filename, 'w') as f:\n        f.writelines(lines)\n\n\n# Assuming 'submissions' list is already created from the previous inference step\n# (or you can re-run the inference part to generate it)\nwrite_submission_csv(submissions, 'submission.csv')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!head submission.csv","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}