import torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\nimport cv2\nimport pandas as pd\nimport os\nfrom torchvision import transforms\nfrom typing import Tuple, List, Dict
\nDATA_ROOT = './'\nVALIDATION_IMAGES_DIR = os.path.join(DATA_ROOT, 'train_images') # Uses the same image dir as validation images are taken from training data\nVALIDATION_MASKS_CSV = os.path.join(DATA_ROOT, 'val_masks.csv') # File created by the training script\nMODEL_WEIGHTS_PATH = 'forgery_detection_model.pth' # Path to your best trained model
\nIMG_HEIGHT = 256 \nIMG_WIDTH = 256\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nBATCH_SIZE = 16
\ndef rle_decode(mask_rle: str, shape: Tuple[int, int]) -> np.ndarray:\n \"\"\"Decodes a run-length encoded mask to a binary array.\"\"\"\n if mask_rle == 'authentic' or not mask_rle:\n return np.zeros(shape, dtype=np.uint8)\n s = mask_rle.split()\n starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n starts -= 1\n ends = starts + lengths\n img = np.zeros(shape[0] * shape[1], dtype=np.uint8)\n for lo, hi in zip(starts, ends):\n img[lo:hi] = 1\n return img.reshape(shape)
\ndef calculate_iou(preds: np.ndarray, targets: np.ndarray) -> float:\n \"\"\"Calculates Mean Intersection over Union (IoU) across flattened arrays.\"\"\"\n intersection = np.sum(preds * targets)\n union = np.sum(preds) + np.sum(targets) - intersection\n iou = (intersection + 1e-6) / (union + 1e-6)\n return iou
\nclass NoiseResidual(nn.Module):\n def init(self, kernel_size=5):\n super(NoiseResidual, self).init()\n self.kernel_size = kernel_size\n self.pad = kernel_size // 2
\ndef forward(, ):\n channels = .shape[]\n pool = nn.AvgPool2d(kernel_size=.kernel_size, stride=, padding=.pad).to(.device)\n smooth_x = .clone()\n i in range(channels):\n smooth_x[:, i:i+, :, :] = pool([:, i:i+, :, :])\n residual = - smooth_x\n torch.cat([, residual], dim=)\ndef conv_block(in_channels, out_channels):\n return nn.Sequential(\n nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n nn.ReLU(inplace=True),\n nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),\n nn.ReLU(inplace=True)\n )
\nclass ForgeryDetectionUNet(nn.Module):\n def init(self, in_channels=6, out_channels=1):\n super(ForgeryDetectionUNet, self).init()\n self.noise_preprocessor = NoiseResidual(kernel_size=5)\n self.enc1 = conv_block(in_channels, 64)\n self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n self.enc2 = conv_block(64, 128)\n self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)\n self.enc3 = conv_block(128, 256)\n self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)\n self.bottleneck = conv_block(256, 512)\n self.upconv3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)\n self.dec3 = conv_block(512, 256)\n self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)\n self.dec2 = conv_block(256, 128)\n self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)\n self.dec1 = conv_block(128, 64)\n self.out_conv = nn.Conv2d(64, out_channels, kernel_size=1)
\n ():\n x = .noise_preprocessor(x)\n e1 = .enc1(x); p1 = .pool1(e1)\n e2 = .enc2(p1); p2 = .pool2(e2)\n e3 = .enc3(p2); p3 = .pool3(e3)\n b = .bottleneck(p3)\n u3 = .upconv3(b); u3 = torch.cat((u3, e3), dim=); d3 = .dec3(u3)\n u2 = .upconv2(d3); u2 = torch.cat((u2, e2), dim=); d2 = .dec2(u2)\n u1 = .upconv1(d2); u1 = torch.cat((u1, e1), dim=); d1 = .dec1(u1)\n out = .out_conv(d1)\n torch.sigmoid(out)\nclass ValidationForgeryDataset(Dataset):\n def init(self, df: pd.DataFrame, img_dir: str):\n self.df = df\n self.img_dir = img_dir\n self.case_ids = df['case_id'].unique().tolist()\n self.mask_map: Dict[str, List[str]] = self._create_mask_map(df)
\n () -> :\n mask_map = {}\n case_id, group df.groupby():\n mask_map[case_id] = group[].tolist()\n mask_map\n\n ():\n (.case_ids)\n\n () -> [torch.Tensor, torch.Tensor]:\n case_id = .case_ids[idx]\n img_path = os.path.join(.img_dir, )\n\n image = cv2.imread(img_path)\n image :\n FileNotFoundError()\n\n image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n masks_rle = .mask_map.get(case_id, [])\n combined_mask = np.zeros((image.shape[], image.shape[]), dtype=np.uint8)\n\n rle masks_rle:\n rle != :\n mask = rle_decode(rle, (image.shape[], image.shape[]))\n combined_mask = np.maximum(combined_mask, mask) \n\n image = cv2.resize(image, (IMG_WIDTH, IMG_HEIGHT))\n mask = cv2.resize(combined_mask, (IMG_WIDTH, IMG_HEIGHT), interpolation=cv2.INTER_NEAREST)\n\n image_tensor = transforms.ToTensor()(image.astype(np.float32) / )\n mask_tensor = torch.from_numpy(mask).unsqueeze().()\n\n image_tensor, mask_tensor\ndef optimize_threshold(model_path: str, val_df_path: str, val_img_dir: str):\n print(\"--- Starting Threshold Optimization ---\")
\n\nmodel = ForgeryDetectionUNet(=6, =1).to(DEVICE)\ntry:\n model.load_state_dict(torch.load(model_path, =DEVICE))\nexcept FileNotFoundError:\n (f)\n return\n\nmodel.eval()\n\ntry:\n val_df = pd.read_csv(val_df_path)\nexcept FileNotFoundError:\n (f)\n return\n\nval_dataset = ValidationForgeryDataset(val_df, val_img_dir)\nval_dataloader = DataLoader(val_dataset, =BATCH_SIZE, =)\n\n\nall_preds = []\nall_targets = []\n\n(f)\n\nwith torch.no_grad():\n images, masks val_dataloader:\n images = images.(DEVICE)\n outputs = model(images)\n\n all_preds.append(outputs.cpu().numpy())\n all_targets.append(masks.cpu().numpy())\n\nraw_preds = np.concatenate(all_preds, =0).flatten()\nraw_targets = np.concatenate(all_targets, =0).flatten()\n\n\nthresholds = np.arange(0.1, 0.95, 0.05)\nbest_iou = 0.0\nbest_threshold = 0.5\n\n()\n( * 35)\n\n T thresholds:\n binary_preds = (raw_preds > T).astype(np.uint8)\n current_iou = calculate_iou(binary_preds, raw_targets)\n\n (f)\n\n current_iou > best_iou:\n best_iou = current_iou\n best_threshold = T\n\n\n( * 35)\n(f)\n(f)\n(f)\n( * 35)\n()\nreturn best_threshold\nif name == 'main':\n optimize_threshold(\n model_path=MODEL_WEIGHTS_PATH, \n val_df_path=VALIDATION_MASKS_CSV,\n val_img_dir=VALIDATION_IMAGES_DIR\n )
", "messages": [ { "id": 3314414, "postDate": "2025-11-10T07:11:56.300Z", "content": "import torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\nimport cv2\nimport pandas as pd\nimport os\nfrom torchvision import transforms\nfrom typing import Tuple, List, Dict
\nDATA_ROOT = './'\nVALIDATION_IMAGES_DIR = os.path.join(DATA_ROOT, 'train_images') # Uses the same image dir as validation images are taken from training data\nVALIDATION_MASKS_CSV = os.path.join(DATA_ROOT, 'val_masks.csv') # File created by the training script\nMODEL_WEIGHTS_PATH = 'forgery_detection_model.pth' # Path to your best trained model
\nIMG_HEIGHT = 256 \nIMG_WIDTH = 256\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nBATCH_SIZE = 16
\ndef rle_decode(mask_rle: str, shape: Tuple[int, int]) -> np.ndarray:\n \"\"\"Decodes a run-length encoded mask to a binary array.\"\"\"\n if mask_rle == 'authentic' or not mask_rle:\n return np.zeros(shape, dtype=np.uint8)\n s = mask_rle.split()\n starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n starts -= 1\n ends = starts + lengths\n img = np.zeros(shape[0] * shape[1], dtype=np.uint8)\n for lo, hi in zip(starts, ends):\n img[lo:hi] = 1\n return img.reshape(shape)
\ndef calculate_iou(preds: np.ndarray, targets: np.ndarray) -> float:\n \"\"\"Calculates Mean Intersection over Union (IoU) across flattened arrays.\"\"\"\n intersection = np.sum(preds * targets)\n union = np.sum(preds) + np.sum(targets) - intersection\n iou = (intersection + 1e-6) / (union + 1e-6)\n return iou
\nclass NoiseResidual(nn.Module):\n def init(self, kernel_size=5):\n super(NoiseResidual, self).init()\n self.kernel_size = kernel_size\n self.pad = kernel_size // 2
\ndef forward(, ):\n channels = .shape[]\n pool = nn.AvgPool2d(kernel_size=.kernel_size, stride=, padding=.pad).to(.device)\n smooth_x = .clone()\n i in range(channels):\n smooth_x[:, i:i+, :, :] = pool([:, i:i+, :, :])\n residual = - smooth_x\n torch.cat([, residual], dim=)\ndef conv_block(in_channels, out_channels):\n return nn.Sequential(\n nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n nn.ReLU(inplace=True),\n nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),\n nn.ReLU(inplace=True)\n )
\nclass ForgeryDetectionUNet(nn.Module):\n def init(self, in_channels=6, out_channels=1):\n super(ForgeryDetectionUNet, self).init()\n self.noise_preprocessor = NoiseResidual(kernel_size=5)\n self.enc1 = conv_block(in_channels, 64)\n self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n self.enc2 = conv_block(64, 128)\n self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)\n self.enc3 = conv_block(128, 256)\n self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)\n self.bottleneck = conv_block(256, 512)\n self.upconv3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)\n self.dec3 = conv_block(512, 256)\n self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)\n self.dec2 = conv_block(256, 128)\n self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)\n self.dec1 = conv_block(128, 64)\n self.out_conv = nn.Conv2d(64, out_channels, kernel_size=1)
\n ():\n x = .noise_preprocessor(x)\n e1 = .enc1(x); p1 = .pool1(e1)\n e2 = .enc2(p1); p2 = .pool2(e2)\n e3 = .enc3(p2); p3 = .pool3(e3)\n b = .bottleneck(p3)\n u3 = .upconv3(b); u3 = torch.cat((u3, e3), dim=); d3 = .dec3(u3)\n u2 = .upconv2(d3); u2 = torch.cat((u2, e2), dim=); d2 = .dec2(u2)\n u1 = .upconv1(d2); u1 = torch.cat((u1, e1), dim=); d1 = .dec1(u1)\n out = .out_conv(d1)\n torch.sigmoid(out)\nclass ValidationForgeryDataset(Dataset):\n def init(self, df: pd.DataFrame, img_dir: str):\n self.df = df\n self.img_dir = img_dir\n self.case_ids = df['case_id'].unique().tolist()\n self.mask_map: Dict[str, List[str]] = self._create_mask_map(df)
\n () -> :\n mask_map = {}\n case_id, group df.groupby():\n mask_map[case_id] = group[].tolist()\n mask_map\n\n ():\n (.case_ids)\n\n () -> [torch.Tensor, torch.Tensor]:\n case_id = .case_ids[idx]\n img_path = os.path.join(.img_dir, )\n\n image = cv2.imread(img_path)\n image :\n FileNotFoundError()\n\n image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n masks_rle = .mask_map.get(case_id, [])\n combined_mask = np.zeros((image.shape[], image.shape[]), dtype=np.uint8)\n\n rle masks_rle:\n rle != :\n mask = rle_decode(rle, (image.shape[], image.shape[]))\n combined_mask = np.maximum(combined_mask, mask) \n\n image = cv2.resize(image, (IMG_WIDTH, IMG_HEIGHT))\n mask = cv2.resize(combined_mask, (IMG_WIDTH, IMG_HEIGHT), interpolation=cv2.INTER_NEAREST)\n\n image_tensor = transforms.ToTensor()(image.astype(np.float32) / )\n mask_tensor = torch.from_numpy(mask).unsqueeze().()\n\n image_tensor, mask_tensor\ndef optimize_threshold(model_path: str, val_df_path: str, val_img_dir: str):\n print(\"--- Starting Threshold Optimization ---\")
\n\nmodel = ForgeryDetectionUNet(=6, =1).to(DEVICE)\ntry:\n model.load_state_dict(torch.load(model_path, =DEVICE))\nexcept FileNotFoundError:\n (f)\n return\n\nmodel.eval()\n\ntry:\n val_df = pd.read_csv(val_df_path)\nexcept FileNotFoundError:\n (f)\n return\n\nval_dataset = ValidationForgeryDataset(val_df, val_img_dir)\nval_dataloader = DataLoader(val_dataset, =BATCH_SIZE, =)\n\n\nall_preds = []\nall_targets = []\n\n(f)\n\nwith torch.no_grad():\n images, masks val_dataloader:\n images = images.(DEVICE)\n outputs = model(images)\n\n all_preds.append(outputs.cpu().numpy())\n all_targets.append(masks.cpu().numpy())\n\nraw_preds = np.concatenate(all_preds, =0).flatten()\nraw_targets = np.concatenate(all_targets, =0).flatten()\n\n\nthresholds = np.arange(0.1, 0.95, 0.05)\nbest_iou = 0.0\nbest_threshold = 0.5\n\n()\n( * 35)\n\n T thresholds:\n binary_preds = (raw_preds > T).astype(np.uint8)\n current_iou = calculate_iou(binary_preds, raw_targets)\n\n (f)\n\n current_iou > best_iou:\n best_iou = current_iou\n best_threshold = T\n\n\n( * 35)\n(f)\n(f)\n(f)\n( * 35)\n()\nreturn best_threshold\nif name == 'main':\n optimize_threshold(\n model_path=MODEL_WEIGHTS_PATH, \n val_df_path=VALIDATION_MASKS_CSV,\n val_img_dir=VALIDATION_IMAGES_DIR\n )
", "rawMarkdown": "import torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\nimport cv2\nimport pandas as pd\nimport os\nfrom torchvision import transforms\nfrom typing import Tuple, List, Dict\n\n# --- Configuration and Constants ---\nDATA_ROOT = './'\nVALIDATION_IMAGES_DIR = os.path.join(DATA_ROOT, 'train_images') # Uses the same image dir as validation images are taken from training data\nVALIDATION_MASKS_CSV = os.path.join(DATA_ROOT, 'val_masks.csv') # File created by the training script\nMODEL_WEIGHTS_PATH = 'forgery_detection_model.pth' # Path to your best trained model\n\nIMG_HEIGHT = 256 \nIMG_WIDTH = 256\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nBATCH_SIZE = 16\n\n# --- Utility Functions (RLE and IoU Metric) ---\n\ndef rle_decode(mask_rle: str, shape: Tuple[int, int]) -> np.ndarray:\n \"\"\"Decodes a run-length encoded mask to a binary array.\"\"\"\n if mask_rle == 'authentic' or not mask_rle:\n return np.zeros(shape, dtype=np.uint8)\n s = mask_rle.split()\n starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n starts -= 1\n ends = starts + lengths\n img = np.zeros(shape[0] * shape[1], dtype=np.uint8)\n for lo, hi in zip(starts, ends):\n img[lo:hi] = 1\n return img.reshape(shape)\n\ndef calculate_iou(preds: np.ndarray, targets: np.ndarray) -> float:\n \"\"\"Calculates Mean Intersection over Union (IoU) across flattened arrays.\"\"\"\n intersection = np.sum(preds * targets)\n union = np.sum(preds) + np.sum(targets) - intersection\n iou = (intersection + 1e-6) / (union + 1e-6)\n return iou\n\n# --- Model Architecture (MUST match training script) ---\n\nclass NoiseResidual(nn.Module):\n def __init__(self, kernel_size=5):\n super(NoiseResidual, self).__init__()\n self.kernel_size = kernel_size\n self.pad = kernel_size // 2\n\n def forward(self, x):\n channels = x.shape[1]\n pool = nn.AvgPool2d(kernel_size=self.kernel_size, stride=1, padding=self.pad).to(x.device)\n smooth_x = x.clone()\n for i in range(channels):\n smooth_x[:, i:i+1, :, :] = pool(x[:, i:i+1, :, :])\n residual = x - smooth_x\n return torch.cat([x, residual], dim=1)\n\ndef conv_block(in_channels, out_channels):\n return nn.Sequential(\n nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n nn.ReLU(inplace=True),\n nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),\n nn.ReLU(inplace=True)\n )\n\nclass ForgeryDetectionUNet(nn.Module):\n def __init__(self, in_channels=6, out_channels=1):\n super(ForgeryDetectionUNet, self).__init__()\n self.noise_preprocessor = NoiseResidual(kernel_size=5)\n self.enc1 = conv_block(in_channels, 64)\n self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n self.enc2 = conv_block(64, 128)\n self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)\n self.enc3 = conv_block(128, 256)\n self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)\n self.bottleneck = conv_block(256, 512)\n self.upconv3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)\n self.dec3 = conv_block(512, 256)\n self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)\n self.dec2 = conv_block(256, 128)\n self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)\n self.dec1 = conv_block(128, 64)\n self.out_conv = nn.Conv2d(64, out_channels, kernel_size=1)\n \n def forward(self, x):\n x = self.noise_preprocessor(x)\n e1 = self.enc1(x); p1 = self.pool1(e1)\n e2 = self.enc2(p1); p2 = self.pool2(e2)\n e3 = self.enc3(p2); p3 = self.pool3(e3)\n b = self.bottleneck(p3)\n u3 = self.upconv3(b); u3 = torch.cat((u3, e3), dim=1); d3 = self.dec3(u3)\n u2 = self.upconv2(d3); u2 = torch.cat((u2, e2), dim=1); d2 = self.dec2(u2)\n u1 = self.upconv1(d2); u1 = torch.cat((u1, e1), dim=1); d1 = self.dec1(u1)\n out = self.out_conv(d1)\n return torch.sigmoid(out)\n\n# --- Validation Data Loader (Must match ForgeryDataset structure) ---\n\nclass ValidationForgeryDataset(Dataset):\n def __init__(self, df: pd.DataFrame, img_dir: str):\n self.df = df\n self.img_dir = img_dir\n self.case_ids = df['case_id'].unique().tolist()\n self.mask_map: Dict[str, List[str]] = self._create_mask_map(df)\n \n def _create_mask_map(self, df: pd.DataFrame) -> dict:\n mask_map = {}\n for case_id, group in df.groupby('case_id'):\n mask_map[case_id] = group['annotation'].tolist()\n return mask_map\n\n def __len__(self):\n return len(self.case_ids)\n\n def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:\n case_id = self.case_ids[idx]\n img_path = os.path.join(self.img_dir, f\"{case_id}.png\")\n \n image = cv2.imread(img_path)\n if image is None:\n raise FileNotFoundError(f\"Image not found: {img_path}\")\n \n image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n \n masks_rle = self.mask_map.get(case_id, [])\n combined_mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)\n \n for rle in masks_rle:\n if rle != 'authentic':\n mask = rle_decode(rle, (image.shape[0], image.shape[1]))\n combined_mask = np.maximum(combined_mask, mask) \n\n image = cv2.resize(image, (IMG_WIDTH, IMG_HEIGHT))\n mask = cv2.resize(combined_mask, (IMG_WIDTH, IMG_HEIGHT), interpolation=cv2.INTER_NEAREST)\n \n image_tensor = transforms.ToTensor()(image.astype(np.float32) / 255.0)\n mask_tensor = torch.from_numpy(mask).unsqueeze(0).float()\n \n return image_tensor, mask_tensor\n\n# --- Main Optimization Function ---\n\ndef optimize_threshold(model_path: str, val_df_path: str, val_img_dir: str):\n print(\"--- Starting Threshold Optimization ---\")\n \n # 1. Load Model and Data\n model = ForgeryDetectionUNet(in_channels=6, out_channels=1).to(DEVICE)\n try:\n model.load_state_dict(torch.load(model_path, map_location=DEVICE))\n except FileNotFoundError:\n print(f\"ERROR: Model weights not found at {model_path}. Please run forgery_training_template.py first.\")\n return\n\n model.eval()\n \n try:\n val_df = pd.read_csv(val_df_path)\n except FileNotFoundError:\n print(f\"ERROR: Validation mask CSV not found at {val_df_path}. Run training first.\")\n return\n\n val_dataset = ValidationForgeryDataset(val_df, val_img_dir)\n val_dataloader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)\n \n # 2. Collect all predictions and ground truth masks\n all_preds = []\n all_targets = []\n \n print(f\"Generating predictions for {len(val_dataset)} validation images...\")\n\n with torch.no_grad():\n for images, masks in val_dataloader:\n images = images.to(DEVICE)\n outputs = model(images)\n \n all_preds.append(outputs.cpu().numpy())\n all_targets.append(masks.cpu().numpy())\n\n raw_preds = np.concatenate(all_preds, axis=0).flatten()\n raw_targets = np.concatenate(all_targets, axis=0).flatten()\n \n # 3. Perform Threshold Sweep (0.10 to 0.90 in steps of 0.05)\n thresholds = np.arange(0.1, 0.95, 0.05)\n best_iou = 0.0\n best_threshold = 0.5\n \n print(\"\\nStarting threshold sweep...\")\n print(\"-\" * 35)\n\n for T in thresholds:\n binary_preds = (raw_preds > T).astype(np.uint8)\n current_iou = calculate_iou(binary_preds, raw_targets)\n \n print(f\"Threshold {T:.2f}: IoU = {current_iou:.6f}\")\n \n if current_iou > best_iou:\n best_iou = current_iou\n best_threshold = T\n\n # 4. Report Results\n print(\"-\" * 35)\n print(f\"Optimization Complete!\")\n print(f\"🎉 Best IoU found: {best_iou:.6f}\")\n print(f\"🔑 Optimal Threshold: {best_threshold:.2f}\")\n print(\"-\" * 35)\n print(\"ACTION REQUIRED: Update OPTIMAL_THRESHOLD in generate_submission.py with this value!\")\n return best_threshold\n\n\nif __name__ == '__main__':\n optimize_threshold(\n model_path=MODEL_WEIGHTS_PATH, \n val_df_path=VALIDATION_MASKS_CSV,\n val_img_dir=VALIDATION_IMAGES_DIR\n )", "votes": 2 } ], "comments": [], "raw_markdown_by_id": { "3314414": "import torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\nimport cv2\nimport pandas as pd\nimport os\nfrom torchvision import transforms\nfrom typing import Tuple, List, Dict\n\n# --- Configuration and Constants ---\nDATA_ROOT = './'\nVALIDATION_IMAGES_DIR = os.path.join(DATA_ROOT, 'train_images') # Uses the same image dir as validation images are taken from training data\nVALIDATION_MASKS_CSV = os.path.join(DATA_ROOT, 'val_masks.csv') # File created by the training script\nMODEL_WEIGHTS_PATH = 'forgery_detection_model.pth' # Path to your best trained model\n\nIMG_HEIGHT = 256 \nIMG_WIDTH = 256\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nBATCH_SIZE = 16\n\n# --- Utility Functions (RLE and IoU Metric) ---\n\ndef rle_decode(mask_rle: str, shape: Tuple[int, int]) -> np.ndarray:\n \"\"\"Decodes a run-length encoded mask to a binary array.\"\"\"\n if mask_rle == 'authentic' or not mask_rle:\n return np.zeros(shape, dtype=np.uint8)\n s = mask_rle.split()\n starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n starts -= 1\n ends = starts + lengths\n img = np.zeros(shape[0] * shape[1], dtype=np.uint8)\n for lo, hi in zip(starts, ends):\n img[lo:hi] = 1\n return img.reshape(shape)\n\ndef calculate_iou(preds: np.ndarray, targets: np.ndarray) -> float:\n \"\"\"Calculates Mean Intersection over Union (IoU) across flattened arrays.\"\"\"\n intersection = np.sum(preds * targets)\n union = np.sum(preds) + np.sum(targets) - intersection\n iou = (intersection + 1e-6) / (union + 1e-6)\n return iou\n\n# --- Model Architecture (MUST match training script) ---\n\nclass NoiseResidual(nn.Module):\n def __init__(self, kernel_size=5):\n super(NoiseResidual, self).__init__()\n self.kernel_size = kernel_size\n self.pad = kernel_size // 2\n\n def forward(self, x):\n channels = x.shape[1]\n pool = nn.AvgPool2d(kernel_size=self.kernel_size, stride=1, padding=self.pad).to(x.device)\n smooth_x = x.clone()\n for i in range(channels):\n smooth_x[:, i:i+1, :, :] = pool(x[:, i:i+1, :, :])\n residual = x - smooth_x\n return torch.cat([x, residual], dim=1)\n\ndef conv_block(in_channels, out_channels):\n return nn.Sequential(\n nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n nn.ReLU(inplace=True),\n nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),\n nn.ReLU(inplace=True)\n )\n\nclass ForgeryDetectionUNet(nn.Module):\n def __init__(self, in_channels=6, out_channels=1):\n super(ForgeryDetectionUNet, self).__init__()\n self.noise_preprocessor = NoiseResidual(kernel_size=5)\n self.enc1 = conv_block(in_channels, 64)\n self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n self.enc2 = conv_block(64, 128)\n self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)\n self.enc3 = conv_block(128, 256)\n self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)\n self.bottleneck = conv_block(256, 512)\n self.upconv3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)\n self.dec3 = conv_block(512, 256)\n self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)\n self.dec2 = conv_block(256, 128)\n self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)\n self.dec1 = conv_block(128, 64)\n self.out_conv = nn.Conv2d(64, out_channels, kernel_size=1)\n \n def forward(self, x):\n x = self.noise_preprocessor(x)\n e1 = self.enc1(x); p1 = self.pool1(e1)\n e2 = self.enc2(p1); p2 = self.pool2(e2)\n e3 = self.enc3(p2); p3 = self.pool3(e3)\n b = self.bottleneck(p3)\n u3 = self.upconv3(b); u3 = torch.cat((u3, e3), dim=1); d3 = self.dec3(u3)\n u2 = self.upconv2(d3); u2 = torch.cat((u2, e2), dim=1); d2 = self.dec2(u2)\n u1 = self.upconv1(d2); u1 = torch.cat((u1, e1), dim=1); d1 = self.dec1(u1)\n out = self.out_conv(d1)\n return torch.sigmoid(out)\n\n# --- Validation Data Loader (Must match ForgeryDataset structure) ---\n\nclass ValidationForgeryDataset(Dataset):\n def __init__(self, df: pd.DataFrame, img_dir: str):\n self.df = df\n self.img_dir = img_dir\n self.case_ids = df['case_id'].unique().tolist()\n self.mask_map: Dict[str, List[str]] = self._create_mask_map(df)\n \n def _create_mask_map(self, df: pd.DataFrame) -> dict:\n mask_map = {}\n for case_id, group in df.groupby('case_id'):\n mask_map[case_id] = group['annotation'].tolist()\n return mask_map\n\n def __len__(self):\n return len(self.case_ids)\n\n def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:\n case_id = self.case_ids[idx]\n img_path = os.path.join(self.img_dir, f\"{case_id}.png\")\n \n image = cv2.imread(img_path)\n if image is None:\n raise FileNotFoundError(f\"Image not found: {img_path}\")\n \n image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n \n masks_rle = self.mask_map.get(case_id, [])\n combined_mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)\n \n for rle in masks_rle:\n if rle != 'authentic':\n mask = rle_decode(rle, (image.shape[0], image.shape[1]))\n combined_mask = np.maximum(combined_mask, mask) \n\n image = cv2.resize(image, (IMG_WIDTH, IMG_HEIGHT))\n mask = cv2.resize(combined_mask, (IMG_WIDTH, IMG_HEIGHT), interpolation=cv2.INTER_NEAREST)\n \n image_tensor = transforms.ToTensor()(image.astype(np.float32) / 255.0)\n mask_tensor = torch.from_numpy(mask).unsqueeze(0).float()\n \n return image_tensor, mask_tensor\n\n# --- Main Optimization Function ---\n\ndef optimize_threshold(model_path: str, val_df_path: str, val_img_dir: str):\n print(\"--- Starting Threshold Optimization ---\")\n \n # 1. Load Model and Data\n model = ForgeryDetectionUNet(in_channels=6, out_channels=1).to(DEVICE)\n try:\n model.load_state_dict(torch.load(model_path, map_location=DEVICE))\n except FileNotFoundError:\n print(f\"ERROR: Model weights not found at {model_path}. Please run forgery_training_template.py first.\")\n return\n\n model.eval()\n \n try:\n val_df = pd.read_csv(val_df_path)\n except FileNotFoundError:\n print(f\"ERROR: Validation mask CSV not found at {val_df_path}. Run training first.\")\n return\n\n val_dataset = ValidationForgeryDataset(val_df, val_img_dir)\n val_dataloader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)\n \n # 2. Collect all predictions and ground truth masks\n all_preds = []\n all_targets = []\n \n print(f\"Generating predictions for {len(val_dataset)} validation images...\")\n\n with torch.no_grad():\n for images, masks in val_dataloader:\n images = images.to(DEVICE)\n outputs = model(images)\n \n all_preds.append(outputs.cpu().numpy())\n all_targets.append(masks.cpu().numpy())\n\n raw_preds = np.concatenate(all_preds, axis=0).flatten()\n raw_targets = np.concatenate(all_targets, axis=0).flatten()\n \n # 3. Perform Threshold Sweep (0.10 to 0.90 in steps of 0.05)\n thresholds = np.arange(0.1, 0.95, 0.05)\n best_iou = 0.0\n best_threshold = 0.5\n \n print(\"\\nStarting threshold sweep...\")\n print(\"-\" * 35)\n\n for T in thresholds:\n binary_preds = (raw_preds > T).astype(np.uint8)\n current_iou = calculate_iou(binary_preds, raw_targets)\n \n print(f\"Threshold {T:.2f}: IoU = {current_iou:.6f}\")\n \n if current_iou > best_iou:\n best_iou = current_iou\n best_threshold = T\n\n # 4. Report Results\n print(\"-\" * 35)\n print(f\"Optimization Complete!\")\n print(f\"🎉 Best IoU found: {best_iou:.6f}\")\n print(f\"🔑 Optimal Threshold: {best_threshold:.2f}\")\n print(\"-\" * 35)\n print(\"ACTION REQUIRED: Update OPTIMAL_THRESHOLD in generate_submission.py with this value!\")\n return best_threshold\n\n\nif __name__ == '__main__':\n optimize_threshold(\n model_path=MODEL_WEIGHTS_PATH, \n val_df_path=VALIDATION_MASKS_CSV,\n val_img_dir=VALIDATION_IMAGES_DIR\n )" } }