{"cells":[{"cell_type":"markdown","id":"e2e641d5","metadata":{},"source":"# FANet vs MFAD Benchmark: DSB-2018\n\n**Purpose:** Head-to-head comparison for IEEE T-MI submission\n\n| Property | Value |\n|----------|-------|\n| FANet Original | Custom Encoder + SE + MixPool (hard binary gating) |\n| MFAD (Proposed) | ResNet-34 + Multi-scale Feedback Attention Decoder + Gradient Decoupling |\n| Loss | 0.5×BCE + 0.5×Dice |\n| Baseline Dice (Paper) | N/A |\n| Metrics | Dice, mIoU, Sensitivity, Specificity, FPR, Precision |\n"},{"cell_type":"code","execution_count":null,"id":"1467133d","metadata":{},"outputs":[],"source":"import os, sys, random, time, csv, json\nimport cv2\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom glob import glob\nfrom tqdm import tqdm\nfrom PIL import Image\nfrom collections import OrderedDict\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.models as models\n\ntry:\n    import albumentations as A\n    HAS_ALBUM = True\nexcept ImportError:\n    os.system(\"pip install -q albumentations\")\n    import albumentations as A\n    HAS_ALBUM = True\n\ntry:\n    import tifffile as tiff\nexcept ImportError:\n    os.system(\"pip install -q tifffile\")\n    import tifffile as tiff\n\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Device: {DEVICE}\")\nif DEVICE.type == 'cuda':\n    print(f\"GPU: {torch.cuda.get_device_name(0)}\")"},{"cell_type":"code","execution_count":null,"id":"7984c948","metadata":{},"outputs":[],"source":"# ================================================================\n# UniversalDataset — handles jpg/png/tif/gif, auto-binarises masks\n# ================================================================\nclass UniversalDataset(Dataset):\n    def __init__(self, images_path, masks_path, image_size, is_train=False):\n        self.images_path = list(images_path)\n        self.masks_path  = list(masks_path)\n        self.image_size  = image_size\n        self.is_train    = is_train\n        self.n_samples   = len(self.images_path)\n\n        # Paper augmentation suite\n        self.transform = A.Compose([\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.1,\n                               rotate_limit=45, p=0.6),\n            A.ElasticTransform(p=0.3, alpha=120, sigma=120*0.05,\n                               alpha_affine=120*0.03),\n            A.GridDistortion(p=0.3),\n            A.OpticalDistortion(p=0.3, distort_limit=0.05, shift_limit=0.05),\n            A.ToGray(p=0.2),\n            A.RandomBrightnessContrast(p=0.3),\n            A.CoarseDropout(max_holes=8, max_height=32, max_width=32, p=0.2),\n        ])\n\n    def _read_image(self, path):\n        if path.lower().endswith('.gif'):\n            return np.array(Image.open(path).convert('RGB'))\n        img = cv2.imread(path, cv2.IMREAD_UNCHANGED)\n        if img is None:\n            return np.array(Image.open(path).convert('RGB'))\n        if len(img.shape) == 2:\n            return np.stack([img]*3, axis=-1)\n        if img.shape[2] == 4:\n            return cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)\n        return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n    def _read_mask(self, path):\n        if path.lower().endswith('.gif'):\n            return np.array(Image.open(path).convert('L'))\n        m = cv2.imread(path, cv2.IMREAD_UNCHANGED)\n        if m is None:\n            return np.array(Image.open(path).convert('L'))\n        if len(m.shape) == 3:\n            return cv2.cvtColor(m, cv2.COLOR_BGR2GRAY)\n        return m\n\n    def __getitem__(self, index):\n        image = self._read_image(self.images_path[index])\n        mask  = self._read_mask(self.masks_path[index])\n\n        # Resize\n        image = cv2.resize(image, self.image_size, interpolation=cv2.INTER_LINEAR)\n        mask  = cv2.resize(mask,  self.image_size, interpolation=cv2.INTER_NEAREST)\n\n        # Augmentation (train only)\n        if self.is_train:\n            aug = self.transform(image=image, mask=mask)\n            image, mask = aug['image'], aug['mask']\n\n        # Normalise image -> [0,1] float32\n        image = image.astype(np.float32)\n        mx = image.max()\n        if mx > 0:\n            image = image / (65535.0 if mx > 255 else 255.0)\n        image = np.transpose(image, (2, 0, 1))  # HWC -> CHW\n\n        # Binarise mask -> {0, 1}\n        mask = mask.astype(np.float32)\n        uv = np.unique(mask)\n        if len(uv) > 2:\n            mask = (mask > (mask.max() + mask.min()) / 2.0).astype(np.float32)\n        elif mask.max() > 0:\n            mask = (mask == mask.max()).astype(np.float32)\n        mask = np.expand_dims(mask, axis=0)\n\n        return image, mask\n\n    def __len__(self):\n        return self.n_samples\n\n\n# ================================================================\n# EM Volume Dataset (for Electron Microscopy 3D TIFF)\n# ================================================================\nclass EMVolumeDataset(Dataset):\n    def __init__(self, volume, labels, indices, image_size, is_train=False):\n        self.volume  = volume\n        self.labels  = labels\n        self.indices = indices\n        self.image_size = image_size\n        self.is_train = is_train\n        self.transform = A.Compose([\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.1,\n                               rotate_limit=45, p=0.5),\n            A.ElasticTransform(p=0.3, alpha=120, sigma=120*0.05,\n                               alpha_affine=120*0.03),\n            A.GridDistortion(p=0.2),\n        ])\n\n    def __getitem__(self, idx):\n        i = self.indices[idx]\n        img = self.volume[i]\n        msk = self.labels[i]\n\n        if len(img.shape) == 2:\n            img = np.stack([img]*3, axis=-1)\n\n        img = cv2.resize(img, self.image_size, interpolation=cv2.INTER_LINEAR)\n        msk = cv2.resize(msk, self.image_size, interpolation=cv2.INTER_NEAREST)\n\n        if self.is_train:\n            aug = self.transform(image=img, mask=msk)\n            img, msk = aug['image'], aug['mask']\n\n        img = img.astype(np.float32)\n        mx = img.max()\n        if mx > 0:\n            img = img / (65535.0 if mx > 255 else 255.0)\n        img = np.transpose(img, (2, 0, 1))\n\n        msk = msk.astype(np.float32)\n        if msk.max() > 0:\n            msk = (msk / msk.max()).astype(np.float32)\n        msk = np.expand_dims(msk, 0)\n\n        return img, msk\n\n    def __len__(self):\n        return len(self.indices)"},{"cell_type":"code","execution_count":null,"id":"12121f33","metadata":{},"outputs":[],"source":"# ================================================================\n# Data Science Bowl 2018 (DSB-2018) — Nuclei Segmentation\n# ~670 images, ~80:20 split\n# Image size: 256x256\n# Masks are provided as separate images per nucleus -> need merging\n# ================================================================\nimport zipfile\nimport shutil\n\nDATASET_NAME = \"DSB-2018\"\nIMAGE_SIZE = (256, 256)\nEPOCHS = 200\nBATCH_SIZE = 16\nLR = 1e-4\n\n_zips = glob(\"/kaggle/input/**/stage1_train.zip\", recursive=True)\nZIP_PATH = _zips[0] if _zips else \"/kaggle/input/data-science-bowl-2018/stage1_train.zip\"\nORIG_DIR = \"/kaggle/working/stage1_train\"\nMERGED_MASK_DIR = \"/kaggle/working/dsb2018_masks\"\n\nif os.path.exists(ZIP_PATH) and not os.path.exists(ORIG_DIR):\n    print(f\"Extracting {ZIP_PATH}...\")\n    os.makedirs(ORIG_DIR, exist_ok=True)\n    with zipfile.ZipFile(ZIP_PATH, 'r') as zip_ref:\n        zip_ref.extractall(ORIG_DIR)\nelif not os.path.exists(ZIP_PATH) and not os.path.exists(ORIG_DIR):\n    print(\"WARNING: data-science-bowl-2018 dataset is not added to Kaggle Input!\")\n\nos.makedirs(MERGED_MASK_DIR, exist_ok=True)\n\nall_imgs = []\nall_msks = []\n\nprint(\"Merging individual nucleus masks...\")\nif os.path.exists(ORIG_DIR):\n    folder_names = sorted(os.listdir(ORIG_DIR))\n    for folder in tqdm(folder_names, desc=\"Merging Masks\"):\n        folder_path = os.path.join(ORIG_DIR, folder)\n        if not os.path.isdir(folder_path): continue\n\n        img_file = os.path.join(folder_path, \"images\", folder + \".png\")\n        mask_files = glob(os.path.join(folder_path, \"masks\", \"*.png\"))\n\n        if not os.path.exists(img_file) or len(mask_files) == 0: continue\n\n        merged_path = os.path.join(MERGED_MASK_DIR, folder + \".png\")\n        \n        if not os.path.exists(merged_path):\n            merged = None\n            for mf in mask_files:\n                m = np.array(Image.open(mf).convert('L'))\n                if merged is None:\n                    merged = m\n                else:\n                    merged = np.maximum(merged, m)\n            Image.fromarray(merged).save(merged_path)\n            \n        all_imgs.append(img_file)\n        all_msks.append(merged_path)\n\nimages = all_imgs\nmasks = all_msks\n\nassert len(images) > 0, \"No images found!\"\n\nrandom.seed(42)\ncombined = list(zip(images, masks))\nrandom.shuffle(combined)\nimages, masks = zip(*combined)\n\ntotal = len(images)\nsplit = max(1, int(total * 0.8))\ntrain_x, train_y = list(images[:split]), list(masks[:split])\nvalid_x, valid_y = list(images[split:]), list(masks[split:])\n\ntrain_dataset = UniversalDataset(train_x, train_y, IMAGE_SIZE, is_train=True)\nvalid_dataset = UniversalDataset(valid_x, valid_y, IMAGE_SIZE, is_train=False)\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True,\n                          num_workers=2, pin_memory=True)\nvalid_loader = DataLoader(valid_dataset, batch_size=1, shuffle=False,\n                          num_workers=2, pin_memory=True)\n\nprint(f\"DSB-2018: Train={len(train_dataset)}, Val={len(valid_dataset)}\")"},{"cell_type":"markdown","id":"55af60e3","metadata":{},"source":"## Model Definitions\nBoth architectures are defined inline for reproducibility."},{"cell_type":"code","execution_count":null,"id":"32f32ced","metadata":{},"outputs":[],"source":"# ================================================================\n# FANet Original — Exact reproduction from the paper's codebase\n# Custom Encoder/Decoder + SE + MixPool with hard binary gating\n# Input: [image, mask] as a list\n# ================================================================\nclass SELayer(nn.Module):\n    '''Squeeze-and-Excitation channel attention.'''\n    def __init__(self, channel, reduction=16):\n        super().__init__()\n        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n        self.fc = nn.Sequential(\n            nn.Linear(channel, channel // reduction, bias=False),\n            nn.ReLU(inplace=True),\n            nn.Linear(channel // reduction, channel, bias=False),\n            nn.Sigmoid()\n        )\n\n    def forward(self, x):\n        b, c, _, _ = x.size()\n        y = self.avg_pool(x).view(b, c)\n        y = self.fc(y).view(b, c, 1, 1)\n        return x * y.expand_as(x)\n\n\nclass ResidualBlock(nn.Module):\n    '''3x3->3x3 Residual block with SE attention.'''\n    def __init__(self, in_c, out_c):\n        super().__init__()\n        self.conv1 = nn.Conv2d(in_c, out_c, kernel_size=3, padding=1)\n        self.bn1 = nn.BatchNorm2d(out_c)\n        self.conv2 = nn.Conv2d(out_c, out_c, kernel_size=3, padding=1)\n        self.bn2 = nn.BatchNorm2d(out_c)\n        self.conv3 = nn.Conv2d(in_c, out_c, kernel_size=1, padding=0)\n        self.bn3 = nn.BatchNorm2d(out_c)\n        self.se = SELayer(out_c, out_c)\n        self.relu = nn.ReLU(inplace=True)\n\n    def forward(self, x):\n        x1 = self.relu(self.bn1(self.conv1(x)))\n        x2 = self.bn2(self.conv2(x1))\n        x3 = self.se(self.bn3(self.conv3(x)))\n        return self.relu(x2 + x3)\n\n\nclass MixPool_Original(nn.Module):\n    '''MixPool with HARD BINARY gating (original FANet behaviour).'''\n    def __init__(self, in_c, out_c):\n        super().__init__()\n        self.fmask = nn.Sequential(\n            nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_c),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(out_c, 1, kernel_size=1, padding=0),\n            nn.Sigmoid()\n        )\n        self.conv1 = nn.Sequential(\n            nn.Conv2d(in_c, out_c // 2, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_c // 2),\n            nn.ReLU(inplace=True)\n        )\n        self.conv2 = nn.Sequential(\n            nn.Conv2d(in_c, out_c // 2, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_c // 2),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x, m):\n        fmask = self.fmask(x)\n        stride_h = m.shape[2] // x.shape[2]\n        stride_w = m.shape[3] // x.shape[3]\n        m = nn.MaxPool2d((stride_h, stride_w))(m)\n        m_fg = m[:, 0:1]\n\n        # HARD BINARY: gradient killed by (> 0.5)\n        fmask_g = (fmask > 0.5).float()\n        keep = torch.maximum(fmask_g, m_fg)\n\n        x1 = self.conv1(x * keep)\n        x2 = self.conv2(x)\n        return torch.cat([x1, x2], dim=1)\n\n\nclass EncoderBlock_Orig(nn.Module):\n    def __init__(self, in_c, out_c):\n        super().__init__()\n        self.r1 = ResidualBlock(in_c, out_c)\n        self.r2 = ResidualBlock(out_c, out_c)\n        self.p1 = MixPool_Original(out_c, out_c)\n        self.pool = nn.MaxPool2d((2, 2))\n\n    def forward(self, inputs, masks):\n        x = self.r1(inputs)\n        x = self.r2(x)\n        p = self.p1(x, masks)\n        o = self.pool(p)\n        return o, x\n\n\nclass DecoderBlock_Orig(nn.Module):\n    def __init__(self, in_c, out_c):\n        super().__init__()\n        self.upsample = nn.ConvTranspose2d(in_c, in_c, kernel_size=4,\n                                            stride=2, padding=1)\n        self.r1 = ResidualBlock(in_c + in_c, out_c)\n        self.r2 = ResidualBlock(out_c, out_c)\n        self.p1 = MixPool_Original(out_c, out_c)\n\n    def forward(self, inputs, skip, masks):\n        x = self.upsample(inputs)\n        x = torch.cat([x, skip], dim=1)\n        x = self.r1(x)\n        x = self.r2(x)\n        return self.p1(x, masks)\n\n\nclass FANet_Original(nn.Module):\n    '''FANet with hard binary MixPool gating — exact paper reproduction.'''\n    def __init__(self):\n        super().__init__()\n        self.e1 = EncoderBlock_Orig(3, 32)\n        self.e2 = EncoderBlock_Orig(32, 64)\n        self.e3 = EncoderBlock_Orig(64, 128)\n        self.e4 = EncoderBlock_Orig(128, 256)\n        self.d1 = DecoderBlock_Orig(256, 128)\n        self.d2 = DecoderBlock_Orig(128, 64)\n        self.d3 = DecoderBlock_Orig(64, 32)\n        self.d4 = DecoderBlock_Orig(32, 16)\n        self.output = nn.Conv2d(16 + 1, 1, kernel_size=1, padding=0)\n\n    def forward(self, x, m_fg=None):\n        if m_fg is None:\n            m_fg = torch.zeros(x.size(0), 1, x.size(2), x.size(3),\n                               device=x.device)\n        masks = m_fg\n        p1, s1 = self.e1(x, masks)\n        p2, s2 = self.e2(p1, masks)\n        p3, s3 = self.e3(p2, masks)\n        p4, s4 = self.e4(p3, masks)\n        d1 = self.d1(p4, s4, masks)\n        d2 = self.d2(d1, s3, masks)\n        d3 = self.d3(d2, s2, masks)\n        d4 = self.d4(d3, s1, masks)\n        d5 = torch.cat([d4, masks[:, 0:1]], dim=1)\n        return self.output(d5)"},{"cell_type":"code","execution_count":null,"id":"8308ac62","metadata":{},"outputs":[],"source":"# ================================================================\n# FANet-MFAD (Proposed) — ResNet-34 + Multi-scale Feedback Attention\n# Decoder with Gradient Decoupling (.detach())\n# ================================================================\nclass ConvBlock(nn.Module):\n    def __init__(self, in_c, out_c):\n        super().__init__()\n        self.conv1 = nn.Conv2d(in_c, out_c, 3, padding=1)\n        self.bn1   = nn.BatchNorm2d(out_c)\n        self.conv2 = nn.Conv2d(out_c, out_c, 3, padding=1)\n        self.bn2   = nn.BatchNorm2d(out_c)\n        self.relu  = nn.ReLU(inplace=True)\n\n    def forward(self, x):\n        return self.relu(self.bn2(self.conv2(self.relu(self.bn1(self.conv1(x))))))\n\n\nclass FANet_MFAD(nn.Module):\n    '''Multi-scale Feedback Attention Decoder with Gradient Decoupling.\n    \n    Key innovations:\n      1. ResNet-34 pre-trained encoder (transfer learning)\n      2. 4-channel input (RGB + detached prev_mask) for spatial guidance\n      3. Learned gate on prev_mask with .detach() — Feedback Firewall\n      4. Multi-scale attention: feat * (1 + att(mask)) at every decoder level\n    '''\n    def __init__(self):\n        super().__init__()\n        resnet = models.resnet34(weights=models.ResNet34_Weights.IMAGENET1K_V1)\n\n        # Encoder stem: accept 4-ch input (RGB + prev_mask)\n        self.e1_conv = nn.Conv2d(4, 64, 7, stride=2, padding=3, bias=False)\n        with torch.no_grad():\n            self.e1_conv.weight[:, :3] = resnet.conv1.weight\n            self.e1_conv.weight[:, 3:4] = 0.0\n        self.e1_bn   = resnet.bn1\n        self.e1_relu = resnet.relu\n        self.pool    = resnet.maxpool\n\n        self.e2 = resnet.layer1   # 64\n        self.e3 = resnet.layer2   # 128\n        self.e4 = resnet.layer3   # 256\n        self.b  = resnet.layer4   # 512\n\n        # Decoder with skip connections\n        self.up1 = nn.ConvTranspose2d(512, 256, 2, stride=2)\n        self.d1  = ConvBlock(512, 256)\n        self.att1 = nn.Sequential(nn.Conv2d(1, 256, 1),\n                                  nn.BatchNorm2d(256), nn.Sigmoid())\n\n        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)\n        self.d2  = ConvBlock(256, 128)\n        self.att2 = nn.Sequential(nn.Conv2d(1, 128, 1),\n                                  nn.BatchNorm2d(128), nn.Sigmoid())\n\n        self.up3 = nn.ConvTranspose2d(128, 64, 2, stride=2)\n        self.d3  = ConvBlock(128, 64)\n        self.att3 = nn.Sequential(nn.Conv2d(1, 64, 1),\n                                  nn.BatchNorm2d(64), nn.Sigmoid())\n\n        self.up4 = nn.ConvTranspose2d(64, 64, 2, stride=2)\n        self.d4  = ConvBlock(128, 64)\n        self.att4 = nn.Sequential(nn.Conv2d(1, 64, 1),\n                                  nn.BatchNorm2d(64), nn.Sigmoid())\n\n        self.up5 = nn.ConvTranspose2d(64, 32, 2, stride=2)\n        self.d5  = ConvBlock(32, 32)\n        self.att5 = nn.Sequential(nn.Conv2d(1, 32, 1),\n                                  nn.BatchNorm2d(32), nn.Sigmoid())\n\n        self.out = nn.Conv2d(32, 1, 1)\n        self.learned_gate = nn.Conv2d(1, 1, 1)\n\n    def _mfad(self, feat, prev_mask, att_layer):\n        '''Multi-scale Feedback Attention: feat * (1 + att(mask))'''\n        m = F.interpolate(prev_mask, size=feat.shape[2:], mode='nearest')\n        return feat * (1 + att_layer(m))\n\n    def forward(self, x, prev_mask=None):\n        if prev_mask is None:\n            prev_mask = torch.zeros(x.size(0), 1, x.size(2), x.size(3),\n                                    device=x.device)\n        else:\n            # GRADIENT DECOUPLING: .detach() severs the Feedback Trap\n            prev_mask = torch.sigmoid(self.learned_gate(prev_mask.detach()))\n\n        # Encoder\n        s1 = self.e1_relu(self.e1_bn(self.e1_conv(\n            torch.cat([x, prev_mask], dim=1))))\n        s2 = self.e2(self.pool(s1))\n        s3 = self.e3(s2)\n        s4 = self.e4(s3)\n        b  = self.b(s4)\n\n        # Decoder with multi-scale feedback attention\n        d1 = self.d1(torch.cat([self.up1(b), s4], 1))\n        d1 = self._mfad(d1, prev_mask, self.att1)\n\n        d2 = self.d2(torch.cat([self.up2(d1), s3], 1))\n        d2 = self._mfad(d2, prev_mask, self.att2)\n\n        d3 = self.d3(torch.cat([self.up3(d2), s2], 1))\n        d3 = self._mfad(d3, prev_mask, self.att3)\n\n        d4 = self.d4(torch.cat([self.up4(d3), s1], 1))\n        d4 = self._mfad(d4, prev_mask, self.att4)\n\n        d5 = self.d5(self.up5(d4))\n        d5 = self._mfad(d5, prev_mask, self.att5)\n\n        return self.out(d5)"},{"cell_type":"code","execution_count":null,"id":"a02e07ef","metadata":{},"outputs":[],"source":"# ================================================================\n# Loss: 0.5*BCE + 0.5*Dice (matches original FANet paper exactly)\n# ================================================================\nclass DiceBCELoss(nn.Module):\n    @torch.autocast(device_type='cuda', enabled=False)  # BCE is unsafe under autocast -> run loss in fp32\n    def forward(self, inputs, targets, smooth=1):\n        inputs, targets = inputs.float(), targets.float()\n        probs = torch.sigmoid(inputs).view(-1)\n        tgt   = targets.view(-1)\n        bce   = F.binary_cross_entropy(probs, tgt, reduction='mean')\n        inter = (probs * tgt).sum()\n        dice  = 1 - (2. * inter + smooth) / (probs.sum() + tgt.sum() + smooth)\n        return 0.5 * bce + 0.5 * dice\n\n\n# ================================================================\n# Metrics — per-image, computed on binarised predictions\n# ================================================================\ndef compute_metrics(y_true_np, y_pred_np):\n    '''Compute Dice, mIoU, Sensitivity, Specificity, FPR.\n    \n    Args:\n        y_true_np: binary GT mask [H, W] or [1, H, W]\n        y_pred_np: binary pred mask [H, W] or [1, H, W]\n    Returns:\n        dict with keys: dice, miou, sensitivity, specificity, fpr, precision\n    '''\n    yt = y_true_np.flatten().astype(np.float32)\n    yp = y_pred_np.flatten().astype(np.float32)\n    \n    eps = 1e-7\n    tp = np.sum((yt == 1) & (yp == 1))\n    fp = np.sum((yt == 0) & (yp == 1))\n    fn = np.sum((yt == 1) & (yp == 0))\n    tn = np.sum((yt == 0) & (yp == 0))\n\n    dice = (2 * tp) / (2 * tp + fp + fn + eps)\n    iou  = tp / (tp + fp + fn + eps)\n    sensitivity = tp / (tp + fn + eps)          # = Recall\n    specificity = tn / (tn + fp + eps)\n    fpr  = fp / (fp + tn + eps)\n    precision = tp / (tp + fp + eps)\n    \n    return {\n        'dice': dice,\n        'miou': iou,\n        'sensitivity': sensitivity,\n        'specificity': specificity,\n        'fpr': fpr,\n        'precision': precision,\n    }"},{"cell_type":"code","execution_count":null,"id":"a9f35b0a","metadata":{},"outputs":[],"source":"# ================================================================\n# Training function — works for both FANet_Original and FANet_MFAD\n# ================================================================\ndef train_model(model, train_loader, valid_loader, model_name,\n                epochs=EPOCHS, lr=LR, device=DEVICE):\n    #Train a model and return best checkpoint path + training history.#\n    criterion = DiceBCELoss()\n    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n        optimizer, mode='max', factor=0.5, patience=10, min_lr=1e-6)\n\n    # Mixed precision for speed on T4/P100\n    scaler = torch.amp.GradScaler('cuda', enabled=(device.type == 'cuda'))\n    use_amp = (device.type == 'cuda')\n\n    ckpt_path = f\"{model_name}_{DATASET_NAME.replace(' ', '_')}_best.pth\"\n    best_dice = 0.0\n    history = {'train_loss': [], 'val_dice': []}\n\n    for epoch in range(epochs):\n        # ---- TRAIN ----\n        model.train()\n        train_loss = 0.0\n        for x, y in tqdm(train_loader,\n                         desc=f'[{model_name}] E{epoch+1}/{epochs} Train',\n                         leave=False):\n            x, y = x.to(device), y.to(device)\n            optimizer.zero_grad()\n            with torch.amp.autocast('cuda', enabled=use_amp):\n                out = model(x)                       # iter 1 (no feedback)\n                out = model(x, torch.sigmoid(out))   # iter 2 (with feedback)\n                loss = criterion(out, y)\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n            train_loss += loss.item()\n\n        train_loss /= len(train_loader)\n\n        # ---- VALID ----\n        model.eval()\n        val_dice_sum = 0.0\n        n_val = 0\n        with torch.no_grad():\n            for x, y in valid_loader:\n                x, y = x.to(device), y.to(device)\n                with torch.amp.autocast('cuda', enabled=use_amp):\n                    out = model(x)\n                    out = model(x, torch.sigmoid(out))\n                pred = (torch.sigmoid(out).cpu().numpy() > 0.5).astype(np.float32)\n                gt   = y.cpu().numpy()\n                # Per-image dice\n                for b_i in range(pred.shape[0]):\n                    m = compute_metrics(gt[b_i], pred[b_i])\n                    val_dice_sum += m['dice']\n                    n_val += 1\n\n        val_dice = val_dice_sum / max(n_val, 1)\n        lr_now = optimizer.param_groups[0]['lr']\n        scheduler.step(val_dice)\n\n        history['train_loss'].append(train_loss)\n        history['val_dice'].append(val_dice)\n\n        print(f\"  [{model_name}] Epoch {epoch+1:03d} | \"\n              f\"LR {lr_now:.1e} | Loss {train_loss:.4f} | \"\n              f\"Val Dice {val_dice:.4f}\")\n\n        if val_dice > best_dice:\n            best_dice = val_dice\n            torch.save(model.state_dict(), ckpt_path)\n\n    print(f\"  [{model_name}] Best Val Dice = {best_dice:.4f}\")\n    return ckpt_path, history"},{"cell_type":"code","execution_count":null,"id":"f9495fda","metadata":{},"outputs":[],"source":"# ================================================================\n# Full Evaluation — loads best checkpoint, computes all metrics\n# ================================================================\ndef evaluate_model(model, ckpt_path, valid_loader, model_name, device=DEVICE):\n    '''Load best checkpoint and compute per-image metrics on val set.'''\n    model.load_state_dict(torch.load(ckpt_path, map_location=device,\n                                     weights_only=True))\n    model.eval()\n\n    all_metrics = []\n    all_preds   = []\n    all_gts     = []\n    all_images  = []\n\n    with torch.no_grad():\n        for x, y in tqdm(valid_loader,\n                         desc=f'[{model_name}] Evaluating', leave=False):\n            x, y = x.to(device), y.to(device)\n            out = model(x)\n            out = model(x, torch.sigmoid(out))  # 2-iter feedback\n            pred = (torch.sigmoid(out).cpu().numpy() > 0.5).astype(np.float32)\n            gt   = y.cpu().numpy()\n            imgs = x.cpu().numpy()\n\n            for b_i in range(pred.shape[0]):\n                m = compute_metrics(gt[b_i], pred[b_i])\n                all_metrics.append(m)\n                all_preds.append(pred[b_i])\n                all_gts.append(gt[b_i])\n                all_images.append(imgs[b_i])\n\n    # Aggregate\n    agg = {}\n    for key in all_metrics[0]:\n        vals = [m[key] for m in all_metrics]\n        agg[key] = {'mean': np.mean(vals), 'std': np.std(vals)}\n\n    return agg, all_metrics, all_images, all_gts, all_preds\n\n\ndef print_results(name, agg):\n    '''Pretty-print aggregated results.'''\n    print(f\"\\n{'='*60}\")\n    print(f\"  {name} — Final Results on {DATASET_NAME}\")\n    print(f\"{'='*60}\")\n    for k, v in agg.items():\n        print(f\"  {k:>15s}: {v['mean']:.4f} ± {v['std']:.4f}\")\n    print(f\"{'='*60}\\n\")\n\n\n# ================================================================\n# Qualitative Comparison (side-by-side visualisation)\n# ================================================================\ndef plot_comparison(images_orig, gts_orig, preds_orig,\n                    images_mfad, gts_mfad, preds_mfad,\n                    n_samples=5):\n    #Plot: Image | GT | FANet Mask | MFAD Mask.#\n    n = min(n_samples, len(images_orig), len(images_mfad))\n    fig, axes = plt.subplots(n, 4, figsize=(16, 4 * n))\n    if n == 1:\n        axes = axes[np.newaxis, :]\n\n    for i in range(n):\n        img = np.transpose(images_orig[i], (1, 2, 0))\n        img = np.clip(img, 0, 1)\n        gt  = gts_orig[i][0]\n        p_orig = preds_orig[i][0]\n        p_mfad = preds_mfad[i][0]\n\n        axes[i, 0].imshow(img)\n        axes[i, 0].set_title('Input Image' if i == 0 else '')\n        axes[i, 0].axis('off')\n\n        axes[i, 1].imshow(gt, cmap='gray')\n        axes[i, 1].set_title('Ground Truth' if i == 0 else '')\n        axes[i, 1].axis('off')\n\n        axes[i, 2].imshow(p_orig, cmap='gray')\n        axes[i, 2].set_title('FANet Original' if i == 0 else '')\n        axes[i, 2].axis('off')\n\n        axes[i, 3].imshow(p_mfad, cmap='gray')\n        axes[i, 3].set_title('MFAD (Ours)' if i == 0 else '')\n        axes[i, 3].axis('off')\n\n    plt.suptitle(f'Qualitative Comparison — {DATASET_NAME}', fontsize=16, y=1.01)\n    plt.tight_layout()\n    plt.savefig(f'comparison_{DATASET_NAME.replace(\" \", \"_\")}.png',\n                dpi=150, bbox_inches='tight')\n    plt.show()\n\n\n# ================================================================\n# Save results to CSV\n# ================================================================\ndef save_benchmark_csv(agg_orig, agg_mfad, dataset_name):\n    '''Save benchmark results to CSV.'''\n    csv_path = f'benchmark_results_{dataset_name.replace(\" \", \"_\")}.csv'\n    with open(csv_path, 'w', newline='') as f:\n        writer = csv.writer(f)\n        writer.writerow(['Model', 'Metric', 'Mean', 'Std'])\n        for model_name, agg in [('FANet_Original', agg_orig),\n                                  ('FANet_MFAD', agg_mfad)]:\n            for k, v in agg.items():\n                writer.writerow([model_name, k,\n                                f\"{v['mean']:.6f}\", f\"{v['std']:.6f}\"])\n    print(f\"Results saved to {csv_path}\")\n    return csv_path"},{"cell_type":"markdown","id":"8597811c","metadata":{},"source":"## Run Benchmark: DSB-2018\nTraining both models on the same data split, then evaluating."},{"cell_type":"code","execution_count":null,"id":"7e9d0770","metadata":{},"outputs":[],"source":"# ================================================================\n# RUN BENCHMARK: Train + Evaluate + Compare + Save\n# ================================================================\nprint(f\"\\n{'#'*70}\")\nprint(f\"#  BENCHMARK: {DATASET_NAME}\")\nprint(f\"#  Image Size: {IMAGE_SIZE} | Epochs: {EPOCHS} | LR: {LR}\")\nprint(f\"#  Train: {len(train_loader.dataset)} | Val: {len(valid_loader.dataset)}\")\nprint(f\"{'#'*70}\\n\")\n\n# ---- 1. Instantiate both models ----\nmodel_orig = FANet_Original().to(DEVICE)\nmodel_mfad = FANet_MFAD().to(DEVICE)\n\norig_params = sum(p.numel() for p in model_orig.parameters())\nmfad_params = sum(p.numel() for p in model_mfad.parameters())\nprint(f\"FANet Original: {orig_params/1e6:.2f}M params\")\nprint(f\"FANet MFAD:     {mfad_params/1e6:.2f}M params\\n\")\n\n# ---- 2. Train FANet Original ----\nprint(\"=\" * 60)\nprint(\"  TRAINING: FANet Original\")\nprint(\"=\" * 60)\nckpt_orig, hist_orig = train_model(model_orig, train_loader, valid_loader,\n                                    \"FANet_Orig\", epochs=EPOCHS, lr=LR)\n\n# ---- 3. Train FANet MFAD ----\nprint(\"\\n\" + \"=\" * 60)\nprint(\"  TRAINING: FANet MFAD (Proposed)\")\nprint(\"=\" * 60)\nckpt_mfad, hist_mfad = train_model(model_mfad, train_loader, valid_loader,\n                                    \"FANet_MFAD\", epochs=EPOCHS, lr=LR)\n\n# ---- 4. Evaluate both models ----\nprint(\"\\n\" + \"=\" * 60)\nprint(\"  EVALUATION\")\nprint(\"=\" * 60)\n\n# Re-create fresh models for loading\nmodel_orig_eval = FANet_Original().to(DEVICE)\nmodel_mfad_eval = FANet_MFAD().to(DEVICE)\n\nagg_orig, metrics_orig, imgs_orig, gts_orig, preds_orig = \\\n    evaluate_model(model_orig_eval, ckpt_orig, valid_loader, \"FANet_Orig\")\nagg_mfad, metrics_mfad, imgs_mfad, gts_mfad, preds_mfad = \\\n    evaluate_model(model_mfad_eval, ckpt_mfad, valid_loader, \"FANet_MFAD\")\n\nprint_results(\"FANet Original\", agg_orig)\nprint_results(\"FANet MFAD (Proposed)\", agg_mfad)\n\n# ---- 5. Save CSV ----\ncsv_path = save_benchmark_csv(agg_orig, agg_mfad, DATASET_NAME)\n\n# ---- 6. Qualitative comparison ----\nplot_comparison(imgs_orig, gts_orig, preds_orig,\n                imgs_mfad, gts_mfad, preds_mfad, n_samples=5)\n\n# ---- 7. Training curves ----\nfig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\nax1.plot(hist_orig['train_loss'], label='FANet Original', alpha=0.8)\nax1.plot(hist_mfad['train_loss'], label='MFAD (Ours)', alpha=0.8)\nax1.set_xlabel('Epoch'); ax1.set_ylabel('Train Loss')\nax1.set_title('Training Loss'); ax1.legend(); ax1.grid(True, alpha=0.3)\n\nax2.plot(hist_orig['val_dice'], label='FANet Original', alpha=0.8)\nax2.plot(hist_mfad['val_dice'], label='MFAD (Ours)', alpha=0.8)\nax2.set_xlabel('Epoch'); ax2.set_ylabel('Val Dice')\nax2.set_title('Validation Dice Score'); ax2.legend(); ax2.grid(True, alpha=0.3)\n\nplt.suptitle(f'Training Curves — {DATASET_NAME}', fontsize=14)\nplt.tight_layout()\nplt.savefig(f'training_curves_{DATASET_NAME.replace(\" \", \"_\")}.png',\n            dpi=150, bbox_inches='tight')\nplt.show()\n\n# ---- 8. Summary ----\nprint(f\"\\n{'='*70}\")\nprint(f\"  BENCHMARK SUMMARY — {DATASET_NAME}\")\nprint(f\"{'='*70}\")\nprint(f\"  {'Metric':<15s} | {'FANet Orig':>12s} | {'MFAD (Ours)':>12s} | {'Delta':>10s}\")\nprint(f\"  {'-'*15}-+-{'-'*12}-+-{'-'*12}-+-{'-'*10}\")\nfor k in ['dice', 'miou', 'sensitivity', 'specificity', 'fpr', 'precision']:\n    vo = agg_orig[k]['mean']\n    vm = agg_mfad[k]['mean']\n    delta = vm - vo\n    sign = '+' if delta >= 0 else ''\n    print(f\"  {k:<15s} | {vo:>12.4f} | {vm:>12.4f} | {sign}{delta:>9.4f}\")\nprint(f\"{'='*70}\")"}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","file_extension":".py","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":5}