{"cells": [{"cell_type": "markdown", "id": "5998a1df", "metadata": {}, "source": "# Weakly-supervised crack segmentation\n\nSegmenting cracks pixel by pixel while training on one bit per image: does this image contain a\ncrack or not.\n\nDataset: [Crack Segmentation Dataset](https://www.kaggle.com/datasets/lakshaymiddha/crack-segmentation-dataset), only the `train/` directory is used.\n\nJakub Laskowski (160287), Jakub G\u00f3rniak (160326)\n\n## What we are allowed to look at\n\n| Stage | What it sees |\n|---|---|\n| Label creation | Each mask is opened once and turned into one bit: is there any crack pixel in this image. The mask is then thrown away. |\n| Training (all stages) | Images and that one bit. |\n| Evaluation | The real masks, on a held out split, after everything is frozen. |\n\nNo threshold and no hyperparameter is picked by looking at a mask.\n\n## Why this is not a crack detector\n\nNothing here knows what a crack looks like. There are no edge detectors, no ridge filters, no\nmorphological thinning and no colour rules. We only assume that\n\n1. a positive image has at least one positive pixel and a negative image has none,\n2. the thing we are looking for covers a small part of the image,\n3. its boundaries line up with edges in the image.\n\nThe same three things are true for \"find the tumour cells given only healthy/sick slide labels\", so\nchanging `DATA_ROOT` is enough to move the notebook to that problem.\n\n## The three stages\n\n```\n       image + 1 bit                    coarse evidence                 full-resolution mask\n    +----------------+              +--------------------+            +--------------------+\n    |  1. MIL        |   score map  |  2. CAM -> pseudo  |  pseudo-   |  3. U-Net          |\n    |  classifier    |------------->|     mask           |----------->|  distillation      |\n    |  (ResNet-34)   |   56x56      |  TTA + guided      |  masks     |  (ResNet-34 enc.)  |\n    +----------------+              |  filter + dual thr |            +--------------------+\n                                    +--------------------+\n```\n\n1. MIL classifier. A fully convolutional ResNet-34 outputs a per pixel logit map and top-k pooling\n   turns it into one logit for the whole image. Training only needs the bit, but the map in the\n   middle is already a class activation map.\n2. CAM to pseudo-mask. Multi-scale and flip TTA, a guided filter to sharpen the edges, then two\n   thresholds: sure foreground, sure background, and an ignore band in between that we drop from the\n   next loss.\n3. Distillation. A U-Net is trained on the pseudo-masks. It averages out the CAM noise and gives\n   full resolution masks. Then one round of self-training where the U-Net relabels the data, always\n   forcing negative images to be completely background."}, {"cell_type": "code", "execution_count": null, "id": "0dfa39d8", "metadata": {}, "outputs": [], "source": "import math\nimport os\nimport random\nimport re\nimport time\nimport warnings\nfrom pathlib import Path\n\nimport albumentations as A\nimport cv2\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nfrom albumentations.pytorch import ToTensorV2\nfrom PIL import Image\nfrom sklearn.metrics import average_precision_score, roc_auc_score\nfrom sklearn.model_selection import train_test_split\nfrom torch.utils.data import DataLoader, Dataset\nfrom tqdm.auto import tqdm\n\nwarnings.filterwarnings(\"ignore\")\n\n# --------------------------------------------------------------------------- config\ndef find_data_root():\n    \"\"\"Kaggle does not always mount the dataset where we expect, so go and look for it.\"\"\"\n    if \"DATA_ROOT\" in os.environ:\n        return Path(os.environ[\"DATA_ROOT\"])\n    mounted = Path(\"/kaggle/input\")\n    for depth in range(1, 7):\n        pattern = \"/\".join([\"*\"] * depth) + \"/train/images\"\n        for hit in sorted(mounted.glob(pattern)):\n            if hit.is_dir():\n                return hit.parent.parent\n    return mounted / \"crack-segmentation-dataset\" / \"crack_segmentation_dataset\"\n\n\nDATA_ROOT = find_data_root()\nTRAIN_DIR = DATA_ROOT / \"train\"          # the only directory we are allowed to train on\nWORK_DIR = Path(os.environ.get(\"WORK_DIR\", \"/kaggle/working\"))\n\n# these are the values behind the reported results, the env overrides are there so we could\n# run the whole notebook quickly as a smoke test\nIMG_SIZE = int(os.environ.get(\"IMG_SIZE\", 448))          # classifier/segmenter input resolution\nBATCH_SIZE = int(os.environ.get(\"BATCH_SIZE\", 8))\nCLS_EPOCHS = int(os.environ.get(\"CLS_EPOCHS\", 8))        # stage 1 - MIL classifier\nSEG_EPOCHS = int(os.environ.get(\"SEG_EPOCHS\", 10))       # stage 3 - U-Net, per self-training round\nSELF_TRAINING_ROUNDS = int(os.environ.get(\"SELF_TRAINING_ROUNDS\", 2))\nSEED = 42\n\nTOPK_RATIO = 0.02       # \"the object covers ~2% of the image\" prior, used by the MIL pooling\nIGNORE_INDEX = 255      # pixels the pseudo-labeller is unsure about\n\nFG_THRESHOLD = 0.45     # dual threshold on the normalised CAM\nBG_THRESHOLD = 0.20\nAREA_RATIO = TOPK_RATIO          # the foreground threshold keeps roughly this much of the image\nRELABEL_AREA_CAP = 2 * TOPK_RATIO  # self-training is not allowed to grow the mask past this\nGUIDED_RADIUS = 11      # the 8 we used at 320 px, scaled up with IMG_SIZE\nTTA_SCALES = (0.5, 0.75, 1.0, 1.5)   # 1.5 x 448 already beats the 2.0 x 320 we used before\n\n# both stages stop once the validation number stops moving, which is what keeps the run inside\n# the session limit. the first run took 3.6 h and the classifier had already peaked by epoch 3\nCLS_PATIENCE = int(os.environ.get(\"CLS_PATIENCE\", 2))\nSEG_PATIENCE = int(os.environ.get(\"SEG_PATIENCE\", 3))\nMIN_EPOCHS = int(os.environ.get(\"MIN_EPOCHS\", 3))\nPUZZLE_PROB = 0.5       # the puzzle branch costs a second forward pass, so only pay it half the time\nW_CLDICE = 0.3          # topology term for the segmenter, see stage 3\n\n# pseudo-masks are an intermediate, not a result. on kaggle everything under /kaggle/working is\n# committed as the kernel output, and 20k pngs in there make the real artefacts impossible to pull\nSCRATCH_DIR = Path(os.environ.get(\"SCRATCH_DIR\", \"/kaggle/temp\" if Path(\"/kaggle\").is_dir() else WORK_DIR))\n\nIMAGENET_MEAN = (0.485, 0.456, 0.406)\nIMAGENET_STD = (0.229, 0.224, 0.225)\n\n\ndef set_seed(seed: int = SEED) -> None:\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n\n\ndef sample_n(frame, n, seed=SEED):\n    \"\"\"Sample at most n rows so the figures still work on a small subset.\"\"\"\n    return frame.sample(min(n, len(frame)), random_state=seed)\n\n\nset_seed()\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ntorch.backends.cudnn.benchmark = True    # the input size never changes, so let cudnn pick its kernels once\nRUN_STARTED = time.time()\n\n\ndef stage_done(name):\n    \"\"\"Print how far into the session we are, the notebook has to finish inside 12 h.\"\"\"\n    print(f\"[{(time.time() - RUN_STARTED) / 60:6.1f} min] {name}\")\n\n\nprint(f\"device: {DEVICE} | torch {torch.__version__}\")\nprint(f\"data:   {TRAIN_DIR} (exists: {TRAIN_DIR.exists()})\")\nif not TRAIN_DIR.exists():\n    for depth in range(1, 6):\n        dirs = sorted(str(p) for p in Path(\"/kaggle/input\").glob(\"/\".join([\"*\"] * depth)) if p.is_dir())\n        print(f\"mounted, depth {depth}:\", dirs[:10])\n"}, {"cell_type": "markdown", "id": "5719e3d0", "metadata": {}, "source": "## 1. Data preparation, throwing the masks away\n\nThe cell below is the only place where a training mask is opened. Each one becomes\n`label = int(mask.any())` and we cache the result in a CSV with just `image_path, label`. Everything\nafter that reads the CSV, so pixel information cannot leak into the training by accident.\n\nThe split is stratified on the label and on the source sub-dataset (the filename prefix, `CFD_`,\n`DeepCrack_`, `Volker_` and so on). Those sub-datasets look very different, so a plain random split\nwould give us a validation set that is too easy."}, {"cell_type": "code", "execution_count": null, "id": "32101c75", "metadata": {}, "outputs": [], "source": "LABEL_CSV = WORK_DIR / \"image_labels.csv\"\nWORK_DIR.mkdir(parents=True, exist_ok=True)\n\n\ndef source_group(stem: str) -> str:\n    \"\"\"Filename prefix, 'CFD_006' -> 'cfd'.\"\"\"\n    match = re.match(r\"^([A-Za-z]+)\", stem)\n    return match.group(1).lower() if match else \"unknown\"\n\n\ndef build_label_table() -> pd.DataFrame:\n    \"\"\"Open every mask once, keep one bit per image and forget the pixels.\"\"\"\n    if LABEL_CSV.exists():\n        return pd.read_csv(LABEL_CSV)\n\n    image_paths = sorted((TRAIN_DIR / \"images\").glob(\"*\"))\n    mask_dir = TRAIN_DIR / \"masks\"\n    rows = []\n    for path in tqdm(image_paths, desc=\"deriving image-level labels\"):\n        mask_path = next(mask_dir.glob(path.stem + \".*\"))\n        mask = np.asarray(Image.open(mask_path).convert(\"L\"))\n        rows.append(\n            {\n                \"image_path\": str(path),\n                \"stem\": path.stem,\n                \"group\": source_group(path.stem),\n                \"label\": int((mask > 127).any()),   # the one bit we keep\n            }\n        )\n    table = pd.DataFrame(rows)\n    table.to_csv(LABEL_CSV, index=False)\n    return table\n\n\ntable = build_label_table()\nassert len(table), f\"no images under {TRAIN_DIR / 'images'}, check the dataset is attached\"\n\n# stratify on (label, source) so every split sees every imaging condition. a group needs about\n# eight images to still leave two in the holdout, smaller ones join the biggest group that\n# has the same label\nlabels = table.label.astype(str)\nstrata = labels + \"|\" + table.group\ncounts = strata.value_counts()\nbiggest = {lab: counts[counts.index.str.startswith(lab + \"|\")].index[0] for lab in labels.unique()}\nstrata = strata.mask(strata.map(counts) < 8, labels.map(biggest))\n\ntrain_idx, holdout_idx = train_test_split(np.arange(len(table)), test_size=0.25, random_state=SEED, stratify=strata)\nval_idx, test_idx = train_test_split(holdout_idx, test_size=0.6, random_state=SEED, stratify=strata.iloc[holdout_idx])\n\ntable[\"split\"] = \"train\"\ntable.loc[val_idx, \"split\"] = \"val\"\ntable.loc[test_idx, \"split\"] = \"test\"\n\ntrain_df = table[table.split == \"train\"].reset_index(drop=True)\nval_df = table[table.split == \"val\"].reset_index(drop=True)\ntest_df = table[table.split == \"test\"].reset_index(drop=True)\n\nprint(f\"total images: {len(table)}\")\nprint(table.groupby(\"split\").agg(n=(\"label\", \"size\"), positives=(\"label\", \"sum\")))\n"}, {"cell_type": "code", "execution_count": null, "id": "ec55286e", "metadata": {}, "outputs": [], "source": "# EDA on what we are allowed to see: the images and their one bit.\ncounts = table.groupby([\"group\", \"label\"]).size().unstack(fill_value=0)\ncounts = counts.loc[counts.sum(axis=1).sort_values(ascending=False).index]\n\nfig = plt.figure(figsize=(14, 4.5))\ngs = fig.add_gridspec(1, 2, width_ratios=[1.3, 1])\n\nax = fig.add_subplot(gs[0])\ncounts.plot(kind=\"bar\", stacked=True, ax=ax, color=[\"#9ecae1\", \"#e6550d\"], width=0.8)\nax.set_title(\"Images per source sub-dataset\")\nax.set_xlabel(\"\")\nax.set_ylabel(\"images\")\nax.legend([\"no crack\", \"crack\"])\nax.tick_params(axis=\"x\", rotation=45)\n\nax = fig.add_subplot(gs[1])\noverall = table.label.value_counts().sort_index()\nax.bar([\"no crack\", \"crack\"], overall.values, color=[\"#9ecae1\", \"#e6550d\"])\nfor i, v in enumerate(overall.values):\n    ax.text(i, v, f\"{v}\\n({v / len(table):.0%})\", ha=\"center\", va=\"bottom\")\nax.set_title(\"Overall class balance\")\nax.set_ylim(0, overall.max() * 1.2)\nplt.tight_layout()\nplt.show()\n\nsizes = {Image.open(p).size for p in sample_n(table, 50).image_path}\nprint(f\"image sizes in a random sample: {sizes}\")\n\nsample = pd.concat([sample_n(table[table.label == 1], 6), sample_n(table[table.label == 0], 6)])\nfig, axes = plt.subplots(2, 6, figsize=(15, 5.2))\nfor ax, (_, row) in zip(axes.ravel(), sample.iterrows(), strict=False):\n    ax.imshow(Image.open(row.image_path))\n    ax.set_title(f\"{'crack' if row.label else 'no crack'} \u00b7 {row.group}\", fontsize=9)\nfor ax in axes.ravel():\n    ax.axis(\"off\")\nfig.suptitle(\"What the model gets to see during training\", y=1.0)\nplt.tight_layout()\nplt.show()\n"}, {"cell_type": "markdown", "id": "e432031c", "metadata": {}, "source": "### Preprocessing and augmentation\n\nRandom resized crop to 448x448 (we started at 320 and the extra resolution matters, a crack is only\na few pixels wide), ImageNet normalisation, and augmentations that do not change the label and do\nnot depend on the orientation: flips and 90 degree rotations (a crack upside down is still a crack)\nand small affine warps.\n\nThe photometric half of the list is heavier than the task needs on purpose. 87% of the images\ncontain a crack and roughly 96% of the crack free ones come from a single source sub-dataset, so\n\"is this image from `noncrack`?\" is very nearly the right answer to \"is this image crack free?\".\nGrayscale, gamma, strong brightness/contrast and JPEG artefacts are there to make the source harder\nto recognise than the crack. For the same reason the training loader draws from a\n`WeightedRandomSampler` so both classes turn up equally often instead of 87/13.\n\n`CrackDataset` returns the image and the label and never stores a mask path, so the rule is enforced\nby the code and not by us remembering it."}, {"cell_type": "code", "execution_count": null, "id": "cb695a52", "metadata": {}, "outputs": [], "source": "# almost every crack free image comes from one source sub-dataset, so the classifier could get a\n# good score by recognising the sub-dataset instead of the crack. the photometric augmentations are\n# heavier than they need to be for the task on purpose, they are there to take that shortcut away\ntrain_tf = A.Compose([\n    A.RandomResizedCrop(size=(IMG_SIZE, IMG_SIZE), scale=(0.5, 1.0)),\n    A.HorizontalFlip(p=0.5),\n    A.VerticalFlip(p=0.5),\n    A.RandomRotate90(p=0.5),\n    A.Affine(scale=(0.8, 1.25), rotate=(-20, 20), p=0.5),\n    A.ToGray(p=0.3),\n    A.RandomBrightnessContrast(0.35, 0.35, p=0.8),\n    A.RandomGamma(gamma_limit=(70, 140), p=0.5),\n    A.HueSaturationValue(8, 20, 10, p=0.3),\n    A.ImageCompression(quality_range=(40, 95), p=0.3),\n    A.OneOf([A.GaussianBlur(blur_limit=(3, 5)), A.MotionBlur(blur_limit=5)], p=0.2),\n    A.GaussNoise(p=0.3),\n    A.Normalize(IMAGENET_MEAN, IMAGENET_STD),\n    ToTensorV2(),\n])\n\neval_tf = A.Compose([\n    A.Resize(IMG_SIZE, IMG_SIZE),\n    A.Normalize(IMAGENET_MEAN, IMAGENET_STD),\n    ToTensorV2(),\n])\n\n\ndef denormalise(tensor):\n    \"\"\"CHW normalised tensor -> HWC image in [0, 1], for plotting.\"\"\"\n    array = tensor.detach().cpu().numpy().transpose(1, 2, 0)\n    return np.clip(array * np.array(IMAGENET_STD) + np.array(IMAGENET_MEAN), 0, 1)\n\n\nclass CrackDataset(Dataset):\n    \"\"\"Images and one binary label, there is no mask anywhere in here.\"\"\"\n\n    def __init__(self, frame, transform, with_raw=False):\n        self.paths = frame.image_path.tolist()\n        self.labels = frame.label.astype(np.float32).tolist()\n        self.stems = frame.stem.tolist()\n        self.transform = transform\n        self.with_raw = with_raw\n\n    def __len__(self):\n        return len(self.paths)\n\n    def __getitem__(self, idx):\n        raw = np.asarray(Image.open(self.paths[idx]).convert(\"RGB\"))\n        item = {\"image\": self.transform(image=raw)[\"image\"], \"label\": self.labels[idx], \"stem\": self.stems[idx]}\n        if self.with_raw:\n            item[\"raw\"] = cv2.resize(raw, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_LINEAR)\n        return item\n\n\ndef make_loader(frame, transform, shuffle=False, batch_size=BATCH_SIZE, with_raw=False, sampler=None):\n    return DataLoader(\n        CrackDataset(frame, transform, with_raw),\n        batch_size=batch_size,\n        shuffle=shuffle and sampler is None,\n        sampler=sampler,\n        num_workers=2,\n        pin_memory=torch.cuda.is_available(),\n        drop_last=shuffle,\n    )\n\n\ndef balanced_sampler(frame):\n    \"\"\"Only 13% of the images are crack free, so draw both classes equally often.\"\"\"\n    counts = frame.label.value_counts()\n    weights = frame.label.map(1.0 / counts).to_numpy()\n    return torch.utils.data.WeightedRandomSampler(weights, num_samples=len(frame), replacement=True)\n\n\ntrain_loader = make_loader(train_df, train_tf, shuffle=True, sampler=balanced_sampler(train_df))\nval_loader = make_loader(val_df, eval_tf)\nprint(f\"train batches: {len(train_loader)} | val batches: {len(val_loader)}\")\n"}, {"cell_type": "markdown", "id": "0d9301ed", "metadata": {}, "source": "## 2. Stage 1, a classifier that also gives us a map\n\n### Picking the architecture\n\nThe image is a bag and every position of the feature map is an instance. The label says either \"at\nleast one instance is positive\" or \"all of them are negative\", which is multiple instance learning.\nIt is the same setup as a slide that is positive if it contains at least one tumour cell.\n\nSo we build the network as a segmenter and only collapse it to one number at the end:\n\n```\nimage -> ResNet-34 (ImageNet, dilated) -> 3x3 conv + BN + ReLU -> 1x1 conv -> score map S (B,1,40,40)\n                                                                                  |\n                                                                            top-k pooling\n                                                                                  v\n                                                                            image logit (B,1)\n```\n\n| Choice | Alternative | Why |\n|---|---|---|\n| Dilated ResNet-34 (no stride in `layer3`/`layer4`, dilation 2 and 4, output stride 8) | plain ResNet-50 at stride 32 | Cracks are a few pixels wide. Stride 32 gives a 14x14 map which is useless. Dilation gives 56x56 with the same parameters and the same receptive field, and ResNet-34 keeps the 4x compute increase manageable. `replace_stride_with_dilation` only works on `Bottleneck` nets so we had to do it by hand. |\n| Top-k pooling, k = 2% of the positions | global average pooling | GAP averages 3136 positions, so a 40 pixel crack is well under 2% of the logit and the gradient drowns. Top-k is a soft maximum and its k says \"the object is small\". Global max pooling is the other extreme, correct in principle but a single position gives a very noisy gradient. |\n| No bias on the 1x1 conv | with a bias | Without it the image logit stays a pooled linear projection of the features, which is what makes S an actual CAM and not an approximation of one. |\n\nThere is also an auxiliary max pooling head on the same score map. It says that on a negative image\nno position is allowed to fire, which is how the negative images give us real supervision for the\nbackground."}, {"cell_type": "code", "execution_count": null, "id": "5f00688b", "metadata": {}, "outputs": [], "source": "def dilate_stage(stage, dilation):\n    \"\"\"Turn a stride-2 ResNet stage into a stride-1 dilated stage (keeps the receptive field).\"\"\"\n    for module in stage.modules():\n        if isinstance(module, nn.Conv2d):\n            if module.stride == (2, 2):\n                module.stride = (1, 1)\n            if module.kernel_size == (3, 3):\n                module.dilation = (dilation, dilation)\n                module.padding = (dilation, dilation)\n\n\nclass MILClassifier(nn.Module):\n    \"\"\"Fully-convolutional ResNet-34 + 1x1 classifier + top-k MIL pooling.\"\"\"\n\n    def __init__(self, pretrained=True, topk_ratio=TOPK_RATIO):\n        super().__init__()\n        weights = torchvision.models.ResNet34_Weights.IMAGENET1K_V1 if pretrained else None\n        backbone = torchvision.models.resnet34(weights=weights)\n        # torchvision only does replace_stride_with_dilation for Bottleneck nets,\n        # so we do it by hand. Output stride goes from 32 to 8.\n        dilate_stage(backbone.layer3, 2)\n        dilate_stage(backbone.layer4, 4)\n\n        self.encoder = nn.Sequential(*list(backbone.children())[:-2])   # -> (B, 512, H/8, W/8)\n        self.neck = nn.Sequential(\n            nn.Conv2d(512, 256, 3, padding=1, bias=False),\n            nn.BatchNorm2d(256),\n            nn.ReLU(inplace=True),\n            nn.Dropout2d(0.1),\n        )\n        self.classifier = nn.Conv2d(256, 1, 1, bias=False)   # no bias, otherwise it is not a real CAM\n        self.topk_ratio = topk_ratio\n\n    def score_map(self, x):\n        \"\"\"Per-pixel logits at input_size / 8.\"\"\"\n        return self.classifier(self.neck(self.encoder(x)))\n\n    def topk_pool(self, score):\n        flat = score.flatten(2)                                   # (B, 1, H*W)\n        k = max(1, round(self.topk_ratio * flat.shape[-1]))\n        return flat.topk(k, dim=-1).values.mean(-1)               # (B, 1)\n\n    @staticmethod\n    def max_pool(score):\n        return score.flatten(2).max(-1).values                    # (B, 1)\n\n    def forward(self, x):\n        score = self.score_map(x)\n        return {\"logit\": self.topk_pool(score), \"max_logit\": self.max_pool(score), \"score_map\": score}\n\n\ndef normalise_map(maps, eps=1e-6):\n    \"\"\"Per-image min-max scaling of an activation map to [0, 1].\"\"\"\n    flat = maps.flatten(2)\n    lo = flat.min(-1, keepdim=True).values\n    hi = flat.max(-1, keepdim=True).values\n    return ((flat - lo) / (hi - lo + eps)).view_as(maps).clamp(0, 1)\n\n\nmodel = MILClassifier().to(DEVICE)\nwith torch.no_grad():\n    out = model(torch.randn(2, 3, IMG_SIZE, IMG_SIZE, device=DEVICE))\nprint(f\"score map: {tuple(out['score_map'].shape)}   image logit: {tuple(out['logit'].shape)}\")\nprint(f\"trainable parameters: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M\")\nprint(f\"top-k pooling uses k = {max(1, round(TOPK_RATIO * out['score_map'][0, 0].numel()))} of {out['score_map'][0, 0].numel()} positions\")\n"}, {"cell_type": "markdown", "id": "ee73ef8c", "metadata": {}, "source": "### The loss: one bit plus self-consistency\n\n$$\\mathcal{L} = \\underbrace{\\mathrm{BCE}(\\mathrm{topk}(S), y)}_{\\text{bag}} + 0.4\\,\\underbrace{\\mathrm{BCE}(\\max(S), y)}_{\\text{strict MIL}} + \\underbrace{\\mathrm{BCE}(\\mathrm{topk}(\\tilde S), y) + \\lambda_p \\lVert S - \\tilde S \\rVert_1}_{\\text{puzzle consistency}} + \\lambda_s\\,\\overline{\\sigma(S)}\\big|_{y=1}$$\n\nThe usual problem with CAM based segmentation is that a classifier only needs the most obvious part\nof the object to get the label right, so the map lights up a small blob and stops there. Two extra\nterms deal with that and neither of them needs more than the bit we already have.\n\n* Puzzle consistency (Puzzle-CAM). Cut the image into 2x2 tiles, classify each tile on its own and\n  stitch the four maps back together into $\\tilde S$. A tile cannot use evidence that lives in\n  another tile, so the network has to respond everywhere the evidence is, and $\\tilde S$ has to\n  agree with the whole image map $S$. This was the single most useful term we tried.\n* Foreground area prior ($\\lambda_s = 0.02$, positives only), a small push towards a smaller\n  activated area. This is assumption 2, set it to 0 if the target is large.\n\nThe regularisers are warmed up linearly over the first two epochs. Before that $S$ is basically\nnoise and forcing it to be consistent with itself only slows the training down.\n\nThe puzzle term needs a second forward pass over the tiles and that is most of the cost of a step,\nso we only apply it on a random half of the steps. That halves its strength as well as its price,\nwhich is a fair trade here: the classifier reaches validation AUROC 1.0 either way and the budget\nis better spent on the segmenter. Training stops as soon as validation AUROC has been flat for two\nepochs, which on the full dataset saves about a third of stage 1 for no loss at all."}, {"cell_type": "code", "execution_count": null, "id": "2570dd82", "metadata": {}, "outputs": [], "source": "W_MAX, W_PUZZLE, W_SPARSITY = 0.4, 0.5, 0.02\nREG_WARMUP_EPOCHS = 2\n\n\ndef tile(x, n=2):\n    \"\"\"(B,C,H,W) -> (B*n*n, C, H/n, W/n), row-major tile order.\"\"\"\n    b, c, h, w = x.shape\n    x = x.reshape(b, c, n, h // n, n, w // n).permute(0, 2, 4, 1, 3, 5)\n    return x.reshape(b * n * n, c, h // n, w // n)\n\n\ndef untile(x, batch, n=2):\n    \"\"\"Inverse of `tile`.\"\"\"\n    _, c, th, tw = x.shape\n    x = x.reshape(batch, n, n, c, th, tw).permute(0, 3, 1, 4, 2, 5)\n    return x.reshape(batch, c, n * th, n * tw)\n\n\ndef mil_loss(model, images, labels, reg_scale=1.0):\n    out = model(images)\n    terms = {\n        \"bag\": F.binary_cross_entropy_with_logits(out[\"logit\"].squeeze(1), labels),\n        \"strict\": W_MAX * F.binary_cross_entropy_with_logits(out[\"max_logit\"].squeeze(1), labels),\n    }\n\n    # the puzzle branch needs a second forward pass over the tiles, which is most of the cost of a\n    # step. sampling it halves that for half the regularisation, and stage 1 reaches AUROC 1.0 either way\n    if reg_scale > 0 and W_PUZZLE > 0 and random.random() < PUZZLE_PROB:\n        tiled = untile(model.score_map(tile(images)), images.shape[0])\n        terms[\"puzzle_cls\"] = F.binary_cross_entropy_with_logits(model.topk_pool(tiled).squeeze(1), labels)\n        terms[\"puzzle_l1\"] = reg_scale * W_PUZZLE * F.l1_loss(out[\"score_map\"], tiled)\n\n    if reg_scale > 0 and W_SPARSITY > 0 and (labels > 0.5).any():\n        terms[\"sparsity\"] = reg_scale * W_SPARSITY * torch.sigmoid(out[\"score_map\"][labels > 0.5]).mean()\n\n    total = torch.stack(list(terms.values())).sum()\n    return total, {k: float(v.detach()) for k, v in terms.items()} | {\"total\": float(total.detach())}\n"}, {"cell_type": "code", "execution_count": null, "id": "b9d92b35", "metadata": {}, "outputs": [], "source": "CLS_CKPT = WORK_DIR / \"mil_classifier.pth\"\n\n\ndef cosine_with_warmup(optimizer, total_steps, warmup=200):\n    def lr_lambda(step):\n        if step < warmup:\n            return (step + 1) / warmup\n        progress = (step - warmup) / max(total_steps - warmup, 1)\n        return 0.5 * (1 + math.cos(math.pi * min(progress, 1.0)))\n    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n\n\n@torch.no_grad()\ndef classifier_scores(model, loader):\n    \"\"\"Image level probabilities and labels, no masks involved.\"\"\"\n    model.eval()\n    scores, labels = [], []\n    for batch in loader:\n        logit = model(batch[\"image\"].to(DEVICE))[\"logit\"].squeeze(1)\n        scores.append(torch.sigmoid(logit).float().cpu().numpy())\n        labels.append(batch[\"label\"].numpy())\n    return np.concatenate(scores), np.concatenate(labels)\n\n\ndef train_classifier(model, epochs=CLS_EPOCHS, patience=CLS_PATIENCE, min_epochs=MIN_EPOCHS):\n    head_params = [p for n, p in model.named_parameters() if not n.startswith(\"encoder\")]\n    optimizer = torch.optim.AdamW(\n        [{\"params\": model.encoder.parameters(), \"lr\": 3e-5}, {\"params\": head_params, \"lr\": 3e-4}],\n        weight_decay=1e-4,\n    )\n    scheduler = cosine_with_warmup(optimizer, epochs * len(train_loader))\n    scaler = torch.amp.GradScaler(\"cuda\", enabled=DEVICE.type == \"cuda\")\n    history, best_auc, stale = [], -1.0, 0\n\n    for epoch in range(epochs):\n        model.train()\n        reg_scale = min(1.0, epoch / REG_WARMUP_EPOCHS) if REG_WARMUP_EPOCHS else 1.0\n        running, started = {}, time.time()\n\n        for batch in tqdm(train_loader, desc=f\"epoch {epoch + 1}/{epochs}\", leave=False):\n            images = batch[\"image\"].to(DEVICE, non_blocking=True)\n            labels = batch[\"label\"].to(DEVICE, non_blocking=True)\n            with torch.autocast(DEVICE.type, enabled=DEVICE.type == \"cuda\"):\n                loss, report = mil_loss(model, images, labels, reg_scale)\n            optimizer.zero_grad(set_to_none=True)\n            scaler.scale(loss).backward()\n            scaler.unscale_(optimizer)\n            nn.utils.clip_grad_norm_(model.parameters(), 5.0)\n            scaler.step(optimizer)\n            scaler.update()\n            scheduler.step()\n            for k, v in report.items():\n                running[k] = running.get(k, 0.0) + v\n\n        train_stats = {k: v / len(train_loader) for k, v in running.items()}\n        scores, labels = classifier_scores(model, val_loader)\n        val_auc = roc_auc_score(labels, scores)\n        val_acc = ((scores >= 0.5) == labels).mean()\n        history.append({\"epoch\": epoch + 1, **train_stats, \"val_auroc\": val_auc, \"val_acc\": val_acc})\n        print(f\"epoch {epoch + 1:2d}  loss {train_stats['total']:.4f}  val AUROC {val_auc:.4f}  val acc {val_acc:.4f}  ({time.time() - started:.0f}s)\")\n\n        if val_auc > best_auc + 1e-4:               # selection uses image-level labels only\n            best_auc, stale = val_auc, 0\n            torch.save(model.state_dict(), CLS_CKPT)\n        else:\n            stale += 1\n\n        if epoch + 1 >= min_epochs and stale >= patience:\n            print(f\"val AUROC has not improved for {patience} epochs, stopping\")\n            break\n    return pd.DataFrame(history)\n\n\nset_seed()\ncls_history = train_classifier(model)\nmodel.load_state_dict(torch.load(CLS_CKPT, map_location=DEVICE))\nprint(f\"best val AUROC: {cls_history.val_auroc.max():.4f}\")\nstage_done(\"stage 1, MIL classifier\")\n"}, {"cell_type": "code", "execution_count": null, "id": "dd0fe84c", "metadata": {}, "outputs": [], "source": "# the gate decides whether an image gets a mask at all. an empty prediction scores dice 1 on an\n# image with no crack and dice 0 on one with a crack, so the threshold is worth choosing rather than\n# leaving at 0.5. we pick it on validation using the image level labels, which is all we are allowed\nval_scores, val_labels = classifier_scores(model, val_loader)\ncandidates = np.linspace(0.05, 0.95, 19)\ngate_f1 = []\nfor threshold in candidates:\n    predicted = val_scores >= threshold\n    hits = float((predicted & (val_labels > 0.5)).sum())\n    gate_f1.append(2 * hits / max(float(predicted.sum()) + float((val_labels > 0.5).sum()), 1e-9))\nGATE_THRESHOLD = float(candidates[int(np.argmax(gate_f1))])\nprint(f\"gate threshold picked on val: {GATE_THRESHOLD:.2f} (val F1 {max(gate_f1):.4f})\")\n\ntest_loader = make_loader(test_df, eval_tf)\ntest_scores, test_labels = classifier_scores(model, test_loader)\ntest_pred = (test_scores >= GATE_THRESHOLD).astype(int)\n\ntp = int(((test_pred == 1) & (test_labels == 1)).sum())\nfp = int(((test_pred == 1) & (test_labels == 0)).sum())\nfn = int(((test_pred == 0) & (test_labels == 1)).sum())\ntn = int(((test_pred == 0) & (test_labels == 0)).sum())\nprecision, recall = tp / max(tp + fp, 1), tp / max(tp + fn, 1)\n\ncls_metrics = {\n    \"accuracy\": (tp + tn) / len(test_labels),\n    \"precision\": precision,\n    \"recall\": recall,\n    \"f1\": 2 * precision * recall / max(precision + recall, 1e-9),\n    \"auroc\": roc_auc_score(test_labels, test_scores),\n    \"average_precision\": average_precision_score(test_labels, test_scores),\n}\n\nfig, axes = plt.subplots(1, 3, figsize=(15, 4))\nloss_cols = [c for c in cls_history.columns if c not in {\"epoch\", \"val_auroc\", \"val_acc\", \"total\"}]\nfor col in loss_cols:\n    axes[0].plot(cls_history.epoch, cls_history[col], label=col, alpha=0.8)\naxes[0].plot(cls_history.epoch, cls_history.total, \"k--\", lw=2, label=\"total\")\naxes[0].set_title(\"Training loss terms\")\naxes[0].set_xlabel(\"epoch\")\naxes[0].legend(fontsize=8)\naxes[0].grid(alpha=0.3)\n\naxes[1].plot(cls_history.epoch, cls_history.val_auroc, \"o-\", label=\"val AUROC\")\naxes[1].plot(cls_history.epoch, cls_history.val_acc, \"s-\", label=\"val accuracy\")\naxes[1].set_title(\"Image-level validation\")\naxes[1].set_xlabel(\"epoch\")\naxes[1].set_ylim(0.5, 1.02)\naxes[1].legend()\naxes[1].grid(alpha=0.3)\n\naxes[2].imshow([[tn, fp], [fn, tp]], cmap=\"Blues\")\nfor (i, j), v in np.ndenumerate([[tn, fp], [fn, tp]]):\n    axes[2].text(j, i, str(v), ha=\"center\", va=\"center\", fontsize=14)\naxes[2].set_xticks([0, 1], [\"pred no crack\", \"pred crack\"])\naxes[2].set_yticks([0, 1], [\"no crack\", \"crack\"])\naxes[2].set_title(\"Test confusion matrix\")\nplt.tight_layout()\nplt.show()\n\nprint(\"image-level test metrics:\")\nfor k, v in cls_metrics.items():\n    print(f\"  {k:18s} {v:.4f}\")\n\n# 87% of the images contain a crack, so accuracy alone looks good no matter what. average precision\n# and the per source breakdown are the numbers that actually say whether the classifier learned the\n# crack or just learned which sub-dataset the image came from\nper_source = test_df.assign(correct=(test_pred == test_labels.astype(int)))\nsummary = per_source.groupby(\"group\").agg(\n    n=(\"label\", \"size\"), positives=(\"label\", \"sum\"), accuracy=(\"correct\", \"mean\"))\nprint(\"\\nper source accuracy:\")\nprint(summary.sort_values(\"n\", ascending=False).round(4).to_string())\n\n# most of the sources are 100% cracked and nearly every crack-free image comes from one of them,\n# so \"which sub-dataset is this\" already answers \"is there a crack\" almost perfectly. the readout\n# below throws away the sources that only ever show one class, which is the part of the image-level\n# score that the shortcut cannot explain. it is a small sample and we report it as such\nmixed = summary[(summary.positives > 0) & (summary.positives < summary.n)].index\ninside = test_df.group.isin(mixed).to_numpy()\nif inside.sum() and 0 < test_labels[inside].sum() < inside.sum():\n    positives_inside = int(test_labels[inside].sum())\n    print(f\"\\nsources containing both classes ({list(mixed)}):\")\n    print(f\"  n {int(inside.sum())} ({positives_inside} cracked / {int(inside.sum()) - positives_inside} not)\")\n    print(f\"  auroc              {roc_auc_score(test_labels[inside], test_scores[inside]):.4f}\")\n    print(f\"  average_precision  {average_precision_score(test_labels[inside], test_scores[inside]):.4f}\")\nelse:\n    print(\"\\nno source contains both classes, the image-level score cannot be separated from the shortcut\")\nstage_done(\"stage 1 evaluation\")\n"}, {"cell_type": "markdown", "id": "b1b7a2c8", "metadata": {}, "source": "## 3. Stage 2, from a 56x56 score map to pixels\n\nThree generic steps, in this order:\n\n1. Multi-scale and flip TTA. We recompute the score map at scales 0.5 / 0.75 / 1.0 / 1.5 with\n   horizontal and vertical flips and average the 12 passes. The small scales give context and the\n   1.5 scale gives resolution (at 672 px input the map is effectively 84x84). The scales used to go\n   up to 2.0, but that was there to make up for a 320 px base, and 1.5 x 448 is already finer than\n   2.0 x 320 for less than half the compute. Averaging also cancels\n   part of the orientation bias that a stack of convolutions always has.\n2. Guided filter. The CAM is smooth and limited by the stride, but the image itself already knows\n   where the edges are, so a guided filter (He et al.) copies that structure onto the probability\n   map. This is assumption 3 and it is not task specific, it is the same filter you would use to\n   snap a tumour heat map to the cell boundaries.\n3. Two thresholds. Instead of deciding everywhere, pixels above `fg` become foreground, pixels below\n   `bg` become background and everything in between becomes `ignore` and is dropped from the stage 3\n   loss. Letting the pseudo-labeller say \"I don't know\" is what stops the U-Net from copying the\n   CAM's mistakes. The foreground threshold is a per image quantile that keeps `TOPK_RATIO` of the\n   pixels, the same 2% object size prior the MIL pooling already uses. Otsu was the obvious choice and it was the wrong one: it splits the histogram into two\n   classes of similar size, while a crack is around 2% of the image, so it cut too low and the\n   pseudo-masks came out much thicker than the crack. On a synthetic crack (a 1-3 px line on a\n   noisy background, so we know the answer) Otsu gets 0.10 IoU against 0.25-0.33 for the quantile.\n   Like Otsu, the quantile is computed from the CAM alone, never from a mask.\n\nOn top of that there is the image level gate: if the classifier says there is no crack, the whole\npseudo-mask is background. That is the only exact supervision we have and we apply it at every\nstage, self-training included."}, {"cell_type": "code", "execution_count": null, "id": "a55f150c", "metadata": {}, "outputs": [], "source": "@torch.no_grad()\ndef multi_scale_cam(model, images, scales=TTA_SCALES, flips=True, size=(IMG_SIZE, IMG_SIZE)):\n    \"\"\"Scale/flip-averaged score map, min-max normalised per image to [0, 1].\"\"\"\n    model.eval()\n    accumulator = torch.zeros(images.shape[0], 1, *size, device=images.device)\n    variants = [(), (-1,), (-2,)] if flips else [()]\n\n    for scale in scales:\n        scaled = images if scale == 1.0 else F.interpolate(\n            images, scale_factor=scale, mode=\"bilinear\", align_corners=False)\n        for dims in variants:\n            batch = torch.flip(scaled, dims) if dims else scaled\n            score = model.score_map(batch)\n            if dims:\n                score = torch.flip(score, dims)\n            accumulator += F.interpolate(F.relu(score), size=size, mode=\"bilinear\", align_corners=False)\n\n    return normalise_map(accumulator / (len(scales) * len(variants)))\n\n\ndef guided_filter(guide_rgb, source, radius=GUIDED_RADIUS, eps=1e-3):\n    \"\"\"Edge preserving filter, pulls source onto the edges of guide_rgb.\"\"\"\n    guide = cv2.cvtColor(guide_rgb, cv2.COLOR_RGB2GRAY).astype(np.float32) / 255.0\n    source = source.astype(np.float32)\n    ksize = (2 * radius + 1, 2 * radius + 1)\n\n    def box(x):\n        return cv2.boxFilter(x, -1, ksize, normalize=True, borderType=cv2.BORDER_REFLECT)\n\n    mean_g, mean_s = box(guide), box(source)\n    var_g = box(guide * guide) - mean_g * mean_g\n    cov_gs = box(guide * source) - mean_g * mean_s\n    a = cov_gs / (var_g + eps)\n    b = mean_s - a * mean_g\n    return np.clip(box(a) * guide + box(b), 0, 1)\n\n\ndef rescale(x, eps=1e-6):\n    lo, hi = float(x.min()), float(x.max())\n    return (x - lo) / (hi - lo + eps)\n\n\ndef area_threshold(probability, ratio=AREA_RATIO):\n    \"\"\"Per image threshold that keeps about `ratio` of the pixels, using the object size prior.\"\"\"\n    return float(np.quantile(probability, 1.0 - ratio))\n\n\ndef probability_to_pseudo_label(probability, is_positive):\n    \"\"\"0 = background, 1 = foreground, IGNORE_INDEX = unknown.\"\"\"\n    if not is_positive:\n        return np.zeros(probability.shape, np.uint8)          # no crack in the image, so no crack pixels\n    # otsu was the obvious choice here and it was the wrong one: it splits the histogram into two\n    # classes of similar size, while a crack covers about 2% of the image, so it always cut too low\n    # and the pseudo-mask came out several times thicker than the crack. the quantile uses the same\n    # size prior as the MIL pooling and it bounds the area by construction\n    fg = area_threshold(probability)\n    bg = min(BG_THRESHOLD, 0.75 * fg)\n    label = np.full(probability.shape, IGNORE_INDEX, np.uint8)\n    label[probability < bg] = 0\n    label[probability >= fg] = 1\n    return label\n\n\ndef refined_probability(cam_slice, raw_image):\n    \"\"\"CAM -> refined, renormalised probability map at IMG_SIZE.\"\"\"\n    return rescale(guided_filter(raw_image, cam_slice))\n"}, {"cell_type": "code", "execution_count": null, "id": "f4e03014", "metadata": {}, "outputs": [], "source": "# the whole point of the project is that no mask is ever used for training, and that is easy to\n# break by accident. every read goes through here and gets counted, and the counter for the\n# training split is asserted to be zero right before the results are computed\nTRAIN_STEMS = set(train_df.stem)\nGT_READS = {\"train\": 0, \"held out\": 0}\n\n\ndef load_gt_mask(stem, size=None):\n    \"\"\"Ground truth, only for the figures and the scoring, never for training.\"\"\"\n    GT_READS[\"train\" if stem in TRAIN_STEMS else \"held out\"] += 1\n    path = next((TRAIN_DIR / \"masks\").glob(stem + \".*\"))\n    mask = (np.asarray(Image.open(path).convert(\"L\")) > 127).astype(np.uint8)\n    return cv2.resize(mask, size, interpolation=cv2.INTER_NEAREST) if size else mask\n\n\n# walk a few positive test images through every step of stage 2\ndemo_df = sample_n(test_df[test_df.label == 1], 4).reset_index(drop=True)\ndemo_batch = next(iter(make_loader(demo_df, eval_tf, batch_size=len(demo_df), with_raw=True)))\ndemo_images = demo_batch[\"image\"].to(DEVICE)\n\nsingle_cam = multi_scale_cam(model, demo_images, scales=(1.0,), flips=False)[:, 0].cpu().numpy()\ntta_cam = multi_scale_cam(model, demo_images)[:, 0].cpu().numpy()\n\ntitles = [\"image\", \"CAM (single scale)\", \"CAM + TTA\", \"+ guided filter\", \"pseudo-mask\", \"ground truth\"]\nfig, axes = plt.subplots(len(demo_df), 6, figsize=(17, 2.9 * len(demo_df)), squeeze=False)\nfor i in range(len(demo_df)):\n    raw = demo_batch[\"raw\"][i].numpy()\n    refined = refined_probability(tta_cam[i], raw)\n    pseudo = probability_to_pseudo_label(refined, is_positive=True)\n    truth = load_gt_mask(demo_batch[\"stem\"][i], (IMG_SIZE, IMG_SIZE))\n\n    for j, panel in enumerate([raw, single_cam[i], tta_cam[i], refined, pseudo, truth]):\n        ax = axes[i][j]\n        if j == 0:\n            ax.imshow(panel)\n        elif j == 4:\n            ax.imshow(np.where(panel == IGNORE_INDEX, 0.5, panel), cmap=\"gray\", vmin=0, vmax=1)\n        elif j == 5:\n            ax.imshow(panel, cmap=\"gray\", vmin=0, vmax=1)\n        else:\n            ax.imshow(panel, cmap=\"inferno\", vmin=0, vmax=1)\n        if i == 0:\n            ax.set_title(titles[j], fontsize=11)\n        ax.axis(\"off\")\nfig.text(0.63, 0.005, \"pseudo-mask: black = background, grey = ignore, white = foreground\", ha=\"center\", fontsize=9)\nplt.tight_layout()\nplt.show()\n"}, {"cell_type": "code", "execution_count": null, "id": "78a464df", "metadata": {}, "outputs": [], "source": "PSEUDO_DIR = SCRATCH_DIR / \"pseudo\"\n\n\n@torch.no_grad()\ndef generate_pseudo_masks(model, frame, out_dir):\n    \"\"\"Write one uint8 PNG per image: {0 background, 1 foreground, 255 ignore}.\"\"\"\n    out_dir.mkdir(parents=True, exist_ok=True)\n    loader = make_loader(frame, eval_tf, with_raw=True)\n    composition = {\"foreground\": 0.0, \"background\": 0.0, \"ignore\": 0.0}\n\n    for batch in tqdm(loader, desc=f\"pseudo-labelling -> {out_dir.name}\", leave=False):\n        cams = multi_scale_cam(model, batch[\"image\"].to(DEVICE))[:, 0].cpu().numpy()\n        for i, stem in enumerate(batch[\"stem\"]):\n            refined = refined_probability(cams[i], batch[\"raw\"][i].numpy())\n            # train/val images come with their image-level bit, so the gate is exact here\n            label = probability_to_pseudo_label(refined, bool(batch[\"label\"][i] > 0.5))\n            Image.fromarray(label).save(out_dir / f\"{stem}.png\")\n            composition[\"foreground\"] += float((label == 1).mean())\n            composition[\"background\"] += float((label == 0).mean())\n            composition[\"ignore\"] += float((label == IGNORE_INDEX).mean())\n\n    return {k: v / len(frame) for k, v in composition.items()}\n\n\nset_seed()\npseudo_stats = {split: generate_pseudo_masks(model, frame, PSEUDO_DIR / split)\n                for split, frame in [(\"train\", train_df), (\"val\", val_df)]}\n\nprint(\"average pseudo-mask composition\")\nprint(pd.DataFrame(pseudo_stats).T.round(4))\nstage_done(\"stage 2, pseudo-labelling\")\n"}, {"cell_type": "markdown", "id": "3dfd805b", "metadata": {}, "source": "## 4. Stage 3, distilling the pseudo-masks into a U-Net\n\nA CAM is a low resolution detector, a U-Net with skip connections is what actually draws the\noutline. Training the U-Net on the confident parts of the pseudo-masks gives us three things:\n\n* Resolution. The skip connections carry the stride 2 and stride 4 detail that the classifier's\n  stride 8 trunk had thrown away, so the output is really full resolution.\n* Denoising. The CAM errors are mostly independent between images, so a network that has to explain\n  thousands of them with one set of weights fits the part that is consistent and averages the rest\n  away. The student ends up better than the teacher.\n* Room to disagree. Around a tenth of the pixels sit in the ignore band and have no target at all,\n  so the network is not punished for correcting the teacher where the teacher was unsure.\n\nLoss: partial BCE, a Tversky term with `alpha = 0.3, beta = 0.7`, and a clDice term, all three only\non the pixels that are not ignored.\n\n* Tversky charges more for a false positive than for a miss. The first version used `pos_weight = 4`\n  and plain Dice, which does the opposite: it buys recall with precision, and the predictions were\n  already too thick, so it made the main failure mode worse.\n* Tversky controls how wide the crack comes out but says nothing about whether it stays in one\n  piece, and in the first full run the U-Net scored worse on boundary F1 than the map it had learned\n  from. clDice compares centrelines instead of areas: it skeletonises both masks with a\n  differentiable erode/dilate and asks whether each skeleton lies inside the other mask, so a crack\n  broken into fragments loses even when the overlap is fine. On its own it would not notice a mask\n  that is too thick, which is exactly the half Tversky already covers, so the two are complementary.\n\nSelf-training: after round 1 the U-Net relabels the training set itself, `p >= 0.7` is foreground,\n`p <= 0.3` is background, the rest is ignore, and negative images are forced to background. The\nforeground threshold is also capped by area so one round cannot mark more than about 4% of an image,\notherwise every round hands the next one a slightly thicker crack. Round 2 trains on those.\n\nPicking the model we ship needs care here. Every round is trained against its own pseudo-masks, so\nthe validation losses of round 1 and round 2 are computed against different targets and comparing\nthem is meaningless. In the first full run round 2 looked worse purely for that reason and still\noverwrote the checkpoint. Now each round is additionally scored against the fixed stage-2 masks,\nwhich gives every round one number on the same scale, and the best of those is the model that gets\nused. A later round is allowed to be worse and if it is, we keep the earlier one. All of this is\nstill pseudo-masks, so no real mask is involved."}, {"cell_type": "code", "execution_count": null, "id": "1d4b97f4", "metadata": {}, "outputs": [], "source": "class DecoderBlock(nn.Module):\n    def __init__(self, in_ch, skip_ch, out_ch):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.Conv2d(in_ch + skip_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),\n            nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),\n        )\n\n    def forward(self, x, skip=None):\n        x = F.interpolate(x, scale_factor=2.0, mode=\"nearest\")\n        if skip is not None:\n            x = torch.cat([x, skip], dim=1)\n        return self.block(x)\n\n\nclass UNet(nn.Module):\n    \"\"\"Classic U-Net with an ImageNet-pretrained ResNet-34 encoder.\"\"\"\n\n    def __init__(self, pretrained=True):\n        super().__init__()\n        weights = torchvision.models.ResNet34_Weights.IMAGENET1K_V1 if pretrained else None\n        backbone = torchvision.models.resnet34(weights=weights)\n        self.stem = nn.Sequential(backbone.conv1, backbone.bn1, backbone.relu)   # /2,  64\n        self.pool, self.layer1 = backbone.maxpool, backbone.layer1               # /4,  64\n        self.layer2, self.layer3, self.layer4 = backbone.layer2, backbone.layer3, backbone.layer4\n        self.dec4 = DecoderBlock(512, 256, 256)\n        self.dec3 = DecoderBlock(256, 128, 128)\n        self.dec2 = DecoderBlock(128, 64, 64)\n        self.dec1 = DecoderBlock(64, 64, 32)\n        self.dec0 = DecoderBlock(32, 0, 16)\n        self.head = nn.Conv2d(16, 1, 3, padding=1)\n\n    def forward(self, x):\n        f0 = self.stem(x)\n        f1 = self.layer1(self.pool(f0))\n        f2 = self.layer2(f1)\n        f3 = self.layer3(f2)\n        f4 = self.layer4(f3)\n        d = self.dec4(f4, f3)\n        d = self.dec3(d, f2)\n        d = self.dec2(d, f1)\n        d = self.dec1(d, f0)\n        return self.head(self.dec0(d))\n\n\nclass PseudoMaskDataset(Dataset):\n    \"\"\"Image plus the mask we generated ourselves.\"\"\"\n\n    def __init__(self, frame, pseudo_dir, transform):\n        self.paths = frame.image_path.tolist()\n        self.stems = frame.stem.tolist()\n        self.pseudo_dir = Path(pseudo_dir)\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.paths)\n\n    def __getitem__(self, idx):\n        image = np.asarray(Image.open(self.paths[idx]).convert(\"RGB\"))\n        pseudo = np.asarray(Image.open(self.pseudo_dir / f\"{self.stems[idx]}.png\"))\n        # pseudo-masks live at IMG_SIZE, the images at their native size\n        image = cv2.resize(image, (pseudo.shape[1], pseudo.shape[0]), interpolation=cv2.INTER_LINEAR)\n        augmented = self.transform(image=image, mask=pseudo)\n        return {\"image\": augmented[\"image\"], \"target\": augmented[\"mask\"].long()}\n\n\n# only flips and rotations, an affine warp would have to invent border pixels and we\n# would not know whether to mark them as ignore\nseg_train_tf = A.Compose([\n    A.Resize(IMG_SIZE, IMG_SIZE),\n    A.HorizontalFlip(p=0.5),\n    A.VerticalFlip(p=0.5),\n    A.RandomRotate90(p=0.5),\n    A.RandomBrightnessContrast(0.25, 0.25, p=0.5),\n    A.HueSaturationValue(8, 20, 10, p=0.3),\n    A.GaussNoise(p=0.2),\n    A.Normalize(IMAGENET_MEAN, IMAGENET_STD),\n    ToTensorV2(),\n])\nseg_eval_tf = A.Compose([A.Resize(IMG_SIZE, IMG_SIZE), A.Normalize(IMAGENET_MEAN, IMAGENET_STD), ToTensorV2()])\n\n\ndef soft_erode(x):\n    return -F.max_pool2d(-x, 3, stride=1, padding=1)\n\n\ndef soft_open(x):\n    return F.max_pool2d(soft_erode(x), 3, stride=1, padding=1)\n\n\ndef soft_skeleton(x, iters=5):\n    \"\"\"Differentiable morphological thinning, the soft skeleton from the clDice paper.\"\"\"\n    opened = soft_open(x)\n    skeleton = F.relu(x - opened)\n    for _ in range(iters):\n        x = soft_erode(x)\n        delta = F.relu(x - soft_open(x))\n        skeleton = skeleton + F.relu(delta - skeleton * delta)\n    return skeleton\n\n\ndef cl_dice(probability, reference, eps=1.0):\n    \"\"\"Dice between each mask and the other one's centreline, so a broken crack is punished.\"\"\"\n    pred_skeleton, true_skeleton = soft_skeleton(probability), soft_skeleton(reference)\n    dims = (1, 2, 3)\n    precision = ((pred_skeleton * reference).sum(dims) + eps) / (pred_skeleton.sum(dims) + eps)\n    sensitivity = ((true_skeleton * probability).sum(dims) + eps) / (true_skeleton.sum(dims) + eps)\n    return (2 * precision * sensitivity / (precision + sensitivity)).mean()\n\n\ndef pseudo_loss(logits, target, pos_weight=1.0, alpha=0.3, beta=0.7):\n    \"\"\"Partial BCE + Tversky + clDice, evaluated only where the pseudo-labeller was confident.\"\"\"\n    logits = logits.squeeze(1)\n    valid = target != IGNORE_INDEX\n    if valid.sum() == 0:\n        return logits.sum() * 0.0\n    weight = torch.tensor(pos_weight, device=logits.device)\n    ce = F.binary_cross_entropy_with_logits(logits[valid], target[valid].float(), pos_weight=weight)\n\n    # dice is symmetric and pos_weight 4 was actively buying recall with precision, which is exactly\n    # the wrong trade here - the predictions were already too thick. tversky with beta > alpha\n    # charges more for a false positive than for a miss and pulls the crack back to its real width\n    probability = torch.sigmoid(logits) * valid\n    reference = (target == 1).float()\n    true_pos = (probability * reference).sum((1, 2))\n    false_neg = ((1 - probability) * reference).sum((1, 2))\n    false_pos = (probability * (1 - reference)).sum((1, 2))\n    tversky = 1 - ((true_pos + 1) / (true_pos + alpha * false_neg + beta * false_pos + 1)).mean()\n\n    # tversky controls how wide the crack comes out but says nothing about whether it stays in one\n    # piece, and the u-net was scoring worse on boundary f1 than the map it learned from. cldice\n    # compares centrelines, so a mask that is split into fragments loses even if the overlap is fine\n    topology = 1 - cl_dice(probability.unsqueeze(1), reference.unsqueeze(1))\n    return ce + tversky + W_CLDICE * topology\n\n\nprint(f\"U-Net parameters: {sum(p.numel() for p in UNet(pretrained=False).parameters()) / 1e6:.1f}M\")\n"}, {"cell_type": "code", "execution_count": null, "id": "8ecfe766", "metadata": {}, "outputs": [], "source": "SEG_CKPT = WORK_DIR / \"unet.pth\"\nCONFIDENCE = 0.7\n\n\ndef make_seg_loader(frame, pseudo_dir, transform, shuffle):\n    return DataLoader(\n        PseudoMaskDataset(frame, pseudo_dir, transform),\n        batch_size=BATCH_SIZE, shuffle=shuffle, num_workers=2,\n        pin_memory=torch.cuda.is_available(), drop_last=shuffle,\n    )\n\n\ndef run_seg_epoch(unet, loader, optimizer=None, scheduler=None, scaler=None):\n    training = optimizer is not None\n    unet.train(training)\n    total = 0.0\n    for batch in tqdm(loader, desc=\"train\" if training else \"val\", leave=False):\n        images = batch[\"image\"].to(DEVICE, non_blocking=True)\n        targets = batch[\"target\"].to(DEVICE, non_blocking=True)\n        with torch.set_grad_enabled(training), torch.autocast(DEVICE.type, enabled=DEVICE.type == \"cuda\"):\n            loss = pseudo_loss(unet(images), targets)\n        if training:\n            optimizer.zero_grad(set_to_none=True)\n            scaler.scale(loss).backward()\n            scaler.unscale_(optimizer)\n            nn.utils.clip_grad_norm_(unet.parameters(), 5.0)\n            scaler.step(optimizer)\n            scaler.update()\n            scheduler.step()\n        total += float(loss.detach())\n    return total / len(loader)\n\n\n@torch.no_grad()\ndef relabel_with_unet(unet, frame, out_dir):\n    \"\"\"Self-training, the U-Net relabels the data. The image level gate still applies.\"\"\"\n    out_dir.mkdir(parents=True, exist_ok=True)\n    unet.eval()\n    for batch in tqdm(make_loader(frame, eval_tf), desc=f\"re-labelling -> {out_dir.name}\", leave=False):\n        probability = torch.sigmoid(unet(batch[\"image\"].to(DEVICE))).squeeze(1).float().cpu().numpy()\n        for i, stem in enumerate(batch[\"stem\"]):\n            if batch[\"label\"][i] < 0.5:\n                label = np.zeros(probability[i].shape, np.uint8)     # negative image, always fully background\n            else:\n                # cap the area as well, otherwise every round hands the next one a slightly\n                # thicker crack and the thickening compounds\n                high = max(CONFIDENCE, area_threshold(probability[i], RELABEL_AREA_CAP))\n                label = np.full(probability[i].shape, IGNORE_INDEX, np.uint8)\n                label[probability[i] <= 1 - CONFIDENCE] = 0\n                label[probability[i] >= high] = 1\n            Image.fromarray(label).save(out_dir / f\"{stem}.png\")\n\n\nset_seed()\nunet = UNet().to(DEVICE)\nseg_history, pseudo_dir = [], PSEUDO_DIR\nROUND_CKPT = WORK_DIR / \"unet_round.pth\"\n\n# each round trains against its own pseudo-masks, so the per-round validation losses are computed\n# against different targets and cannot be compared. this loader keeps the stage-2 masks fixed and\n# gives every round one number on the same scale, which is how we pick the model we ship\nreference_loader = make_seg_loader(val_df, PSEUDO_DIR / \"val\", seg_eval_tf, shuffle=False)\nbest_reference = float(\"inf\")\n\nfor round_id in range(1, SELF_TRAINING_ROUNDS + 1):\n    train_seg_loader = make_seg_loader(train_df, pseudo_dir / \"train\", seg_train_tf, shuffle=True)\n    val_seg_loader = make_seg_loader(val_df, pseudo_dir / \"val\", seg_eval_tf, shuffle=False)\n\n    encoder_params = [p for n, p in unet.named_parameters() if n.startswith((\"stem\", \"layer\"))]\n    decoder_params = [p for n, p in unet.named_parameters() if not n.startswith((\"stem\", \"layer\"))]\n    optimizer = torch.optim.AdamW(\n        [{\"params\": encoder_params, \"lr\": 3e-5}, {\"params\": decoder_params, \"lr\": 3e-4}], weight_decay=1e-4)\n    scheduler = cosine_with_warmup(optimizer, SEG_EPOCHS * len(train_seg_loader), warmup=100)\n    scaler = torch.amp.GradScaler(\"cuda\", enabled=DEVICE.type == \"cuda\")\n    best_val, stale = float(\"inf\"), 0\n\n    for epoch in range(1, SEG_EPOCHS + 1):\n        started = time.time()\n        train_loss = run_seg_epoch(unet, train_seg_loader, optimizer, scheduler, scaler)\n        val_loss = run_seg_epoch(unet, val_seg_loader)\n        seg_history.append({\"round\": round_id, \"epoch\": epoch, \"train_loss\": train_loss, \"val_loss\": val_loss})\n        print(f\"round {round_id} epoch {epoch:2d}  train {train_loss:.4f}  val {val_loss:.4f}  ({time.time() - started:.0f}s)\")\n        if val_loss < best_val - 1e-4:              # selection on pseudo-masks, not ground truth\n            best_val, stale = val_loss, 0\n            torch.save(unet.state_dict(), ROUND_CKPT)\n        else:\n            stale += 1\n\n        if epoch >= MIN_EPOCHS and stale >= SEG_PATIENCE:\n            print(f\"round {round_id} has not improved for {SEG_PATIENCE} epochs, stopping\")\n            break\n\n    unet.load_state_dict(torch.load(ROUND_CKPT, map_location=DEVICE))\n    reference_loss = run_seg_epoch(unet, reference_loader)\n    print(f\"round {round_id} scores {reference_loss:.4f} against the stage-2 masks\")\n    if reference_loss >= best_reference:\n        # the round did not beat the best one, and the next round can only be built on this one,\n        # so there is nothing to gain by paying for it\n        print(f\"round {round_id} did not beat {best_reference:.4f}, stopping self-training\")\n        break\n    best_reference = reference_loss\n    torch.save(unet.state_dict(), SEG_CKPT)\n\n    if round_id < SELF_TRAINING_ROUNDS:\n        pseudo_dir = SCRATCH_DIR / f\"pseudo_round{round_id + 1}\"\n        for split, frame in [(\"train\", train_df), (\"val\", val_df)]:\n            relabel_with_unet(unet, frame, pseudo_dir / split)\n\n# a later round is allowed to be worse than an earlier one, and in that case we keep the earlier one\nunet.load_state_dict(torch.load(SEG_CKPT, map_location=DEVICE))\nseg_history = pd.DataFrame(seg_history)\nstage_done(\"stage 3, u-net self-training\")\n"}, {"cell_type": "markdown", "id": "04b71a71", "metadata": {}, "source": "## 5. Results\n\nEverything above is frozen now. The test split was never seen by any training loop, any threshold or\nany model selection decision, so we can finally open its masks.\n\nEvery variant gives a probability map at 448x448 which we upsample to the mask's own resolution\nbefore thresholding, so all of them are scored on the same pixels. The gate at test time uses the\nclassifier's prediction and not the true bit, because a deployed model would not have the bit\neither.\n\nThe CAM variants binarise with a per image quantile and the U-Net with a flat 0.5. That is two\ndifferent rules, and we spent a run finding out that making them one rule is a bad idea. Section 6\nhas the details.\n\n| Metric | What it tells us |\n|---|---|\n| IoU / Dice (over the whole dataset) | Overall pixel agreement, pooled over all images. |\n| Mean image IoU | The same but averaged per image over the cracked ones, so a few huge cracks do not dominate. |\n| Precision / recall | Whether we over or under segment. |\n| Boundary F1 (+-2 px) | Region metrics are dominated by area, which is harsh for something a few pixels wide. This asks whether the contour is roughly in the right place. |\n\nThe four variants below separate the contribution of each idea: the raw CAM, the TTA, the guided\nfilter, and the U-Net distillation with self-training."}, {"cell_type": "code", "execution_count": null, "id": "dd4326fa", "metadata": {}, "outputs": [], "source": "def boundary_f1(prediction, target, tolerance=2):\n    \"\"\"F1 between predicted and true contours, allowing a few pixels of slack.\"\"\"\n    if prediction.sum() == 0 and target.sum() == 0:\n        return 1.0\n    if prediction.sum() == 0 or target.sum() == 0:\n        return 0.0\n    kernel = np.ones((3, 3), np.uint8)\n    slack = np.ones((2 * tolerance + 1, 2 * tolerance + 1), np.uint8)\n    pred_edge = cv2.morphologyEx(prediction, cv2.MORPH_GRADIENT, kernel)\n    true_edge = cv2.morphologyEx(target, cv2.MORPH_GRADIENT, kernel)\n    precision = (pred_edge * cv2.dilate(true_edge, slack)).sum() / max(pred_edge.sum(), 1)\n    recall = (true_edge * cv2.dilate(pred_edge, slack)).sum() / max(true_edge.sum(), 1)\n    return float(2 * precision * recall / max(precision + recall, 1e-9))\n\n\n@torch.no_grad()\ndef unet_probability(unet, images):\n    \"\"\"U-Net probability averaged over the four flips, the same trick the CAM already gets.\"\"\"\n    unet.eval()\n    total = torch.zeros(images.shape[0], images.shape[2], images.shape[3], device=images.device)\n    for dims in [(), (-1,), (-2,), (-1, -2)]:\n        flipped = torch.flip(images, dims) if dims else images\n        logits = unet(flipped)\n        if dims:\n            logits = torch.flip(logits, dims)\n        total += torch.sigmoid(logits).squeeze(1).float()\n    return total / 4\n\n\n# we tried picking this on validation instead, as the threshold whose average predicted area comes\n# out at TOPK_RATIO, so that the U-Net and the CAM would share one size prior. it made things worse:\n# it landed at 0.99 and cost about 0.03 IoU. the pseudo-masks are cut at the 2% quantile on purpose,\n# and part of what the student gains over the teacher is drawing the crack wider and more complete\n# than the teacher dared, so putting the same 2% back at the decision boundary throws that away.\n# the prior belongs in stage 2. 0.5 is the boundary the BCE term was trained around and it needs\n# nothing but the model, so it stays. the sweep in section 6 is what caught this.\nUNET_THRESHOLD = 0.5\n\n\n@torch.no_grad()\ndef predict_probability(variant, batch):\n    \"\"\"Probability map at IMG_SIZE + the image-level gate from the classifier.\"\"\"\n    images = batch[\"image\"].to(DEVICE)\n    gate = (torch.sigmoid(model(images)[\"logit\"].squeeze(1)).cpu().numpy() >= GATE_THRESHOLD)\n\n    if variant == \"cam\":\n        probability = multi_scale_cam(model, images, scales=(1.0,), flips=False)[:, 0].cpu().numpy()\n    elif variant == \"cam_tta\":\n        probability = multi_scale_cam(model, images)[:, 0].cpu().numpy()\n    elif variant == \"cam_tta_gf\":\n        cams = multi_scale_cam(model, images)[:, 0].cpu().numpy()\n        probability = np.stack([refined_probability(cams[i], batch[\"raw\"][i].numpy()) for i in range(len(cams))])\n    elif variant == \"unet\":\n        probability = unet_probability(unet, images).cpu().numpy()\n    else:\n        raise ValueError(variant)\n    return probability, gate\n\n\ndef binarise(probability, variant):\n    \"\"\"Same size prior for both: a per-image quantile for the CAM, a calibrated one for the U-Net.\"\"\"\n    if variant == \"unet\":\n        return (probability >= UNET_THRESHOLD).astype(np.uint8)\n    return (probability >= area_threshold(probability)).astype(np.uint8)\n\n\ndef evaluate(variant, frame=None, tag=None):\n    frame = test_df if frame is None else frame\n    loader = make_loader(frame, eval_tf, with_raw=True)\n    tp = fp = fn = tn = 0\n    image_iou, boundary = [], []\n\n    for batch in tqdm(loader, desc=f\"evaluating {variant}\", leave=False):\n        probability, gate = predict_probability(variant, batch)\n        for i, stem in enumerate(batch[\"stem\"]):\n            truth = load_gt_mask(stem)\n            prob = cv2.resize(probability[i], (truth.shape[1], truth.shape[0]), interpolation=cv2.INTER_LINEAR)\n            prediction = binarise(prob, variant) if gate[i] else np.zeros_like(truth)\n\n            p, t = prediction.astype(bool), truth.astype(bool)\n            i_tp, i_fp, i_fn = int((p & t).sum()), int((p & ~t).sum()), int((~p & t).sum())\n            tp, fp, fn, tn = tp + i_tp, fp + i_fp, fn + i_fn, tn + int((~p & ~t).sum())\n            if t.any():\n                image_iou.append(i_tp / max(i_tp + i_fp + i_fn, 1))\n                boundary.append(boundary_f1(prediction, truth))\n\n    precision, recall = tp / max(tp + fp, 1), tp / max(tp + fn, 1)\n    return {\n        \"variant\": tag or variant,\n        \"IoU\": tp / max(tp + fp + fn, 1),\n        \"Dice\": 2 * tp / max(2 * tp + fp + fn, 1),\n        \"precision\": precision,\n        \"recall\": recall,\n        \"mean image IoU\": float(np.mean(image_iou)) if image_iou else 0.0,\n        \"boundary F1\": float(np.mean(boundary)) if boundary else 0.0,\n        \"pixel acc\": (tp + tn) / max(tp + tn + fp + fn, 1),\n    }\n\n\nset_seed()\nassert GT_READS[\"train\"] == 0, f\"a training mask was opened {GT_READS['train']} times\"\nprint(f\"supervision audit: {GT_READS['train']} training masks opened, \"\n      f\"{GT_READS['held out']} held-out masks read so far, and only for figures\")\n\nresults = pd.DataFrame([\n    evaluate(\"cam\", tag=\"1. CAM (single scale)\"),\n    evaluate(\"cam_tta\", tag=\"2. + multi-scale/flip TTA\"),\n    evaluate(\"cam_tta_gf\", tag=\"3. + guided-filter refinement\"),\n    evaluate(\"unet\", tag=f\"4. + U-Net distillation ({SELF_TRAINING_ROUNDS} rounds)\"),\n]).set_index(\"variant\")\n\ndisplay(results.round(4))\nstage_done(\"evaluation\")\n"}, {"cell_type": "code", "execution_count": null, "id": "88b4167e", "metadata": {}, "outputs": [], "source": "seg_history = pd.DataFrame(seg_history)\nfig, axes = plt.subplots(1, 3, figsize=(16, 4.2))\n\nx = np.arange(len(results))\nfor offset, metric, colour in [(-0.2, \"IoU\", \"#3182bd\"), (0.0, \"Dice\", \"#e6550d\"), (0.2, \"boundary F1\", \"#31a354\")]:\n    axes[0].bar(x + offset, results[metric], width=0.2, label=metric, color=colour)\naxes[0].set_xticks(x, [f\"{i + 1}\" for i in range(len(results))])\naxes[0].set_xlabel(\"ablation step\")\naxes[0].set_title(\"Ablation on the test split\")\naxes[0].legend()\naxes[0].grid(axis=\"y\", alpha=0.3)\n\naxes[1].plot(results.recall, results.precision, \"o-\", color=\"#756bb1\")\nfor name, row in results.iterrows():\n    axes[1].annotate(name.split(\".\")[0], (row.recall, row.precision), textcoords=\"offset points\", xytext=(6, 4))\naxes[1].set_xlabel(\"recall\")\naxes[1].set_ylabel(\"precision\")\naxes[1].set_title(\"Precision / recall trade-off\")\naxes[1].grid(alpha=0.3)\n\nfor round_id, group in seg_history.groupby(\"round\"):\n    steps = np.arange(len(group)) + (round_id - 1) * SEG_EPOCHS\n    axes[2].plot(steps, group.train_loss, label=f\"round {round_id} train\")\n    axes[2].plot(steps, group.val_loss, \"--\", label=f\"round {round_id} val\")\naxes[2].set_xlabel(\"epoch\")\naxes[2].set_title(\"U-Net loss on pseudo-masks\")\naxes[2].legend(fontsize=8)\naxes[2].grid(alpha=0.3)\n\nplt.tight_layout()\nplt.show()\n\nprint(results[[\"IoU\", \"Dice\", \"boundary F1\"]].round(4).to_string())\n"}, {"cell_type": "code", "execution_count": null, "id": "42d3870d", "metadata": {}, "outputs": [], "source": "@torch.no_grad()\ndef show_predictions(frame, title):\n    frame = frame.reset_index(drop=True)\n    batch = next(iter(make_loader(frame, eval_tf, batch_size=len(frame), with_raw=True)))\n    images = batch[\"image\"].to(DEVICE)\n    cam = multi_scale_cam(model, images)[:, 0].cpu().numpy()\n    unet_prob = torch.sigmoid(unet(images)).squeeze(1).float().cpu().numpy()\n\n    columns = [\"image\", \"CAM (stage 2)\", \"U-Net probability\", \"U-Net mask\", \"ground truth\", \"error map\"]\n    fig, axes = plt.subplots(len(frame), 6, figsize=(17, 2.85 * len(frame)), squeeze=False)\n    for i, stem in enumerate(batch[\"stem\"]):\n        truth = load_gt_mask(stem, (IMG_SIZE, IMG_SIZE))\n        prediction = (unet_prob[i] >= UNET_THRESHOLD).astype(np.uint8)\n        error = np.stack([prediction & (1 - truth), prediction & truth, truth & (1 - prediction)], -1).astype(float)\n\n        for j, panel in enumerate([batch[\"raw\"][i].numpy(), cam[i], unet_prob[i], prediction, truth, error]):\n            ax = axes[i][j]\n            if j in (1, 2):\n                ax.imshow(panel, cmap=\"inferno\", vmin=0, vmax=1)\n            elif j in (3, 4):\n                ax.imshow(panel, cmap=\"gray\", vmin=0, vmax=1)\n            else:\n                ax.imshow(panel)\n            if i == 0:\n                ax.set_title(columns[j], fontsize=11)\n            ax.axis(\"off\")\n    fig.suptitle(f\"{title}   \u00b7   error map: green = correct, red = false positive, blue = missed\", y=1.005)\n    plt.tight_layout()\n    plt.show()\n\n\n# everything below runs on validation. the test split is scored exactly once, by the table above,\n# and looking through test images to find the interesting ones would quietly undo that\npositives = val_df[val_df.label == 1]\nshow_predictions(sample_n(positives, 5), \"Typical predictions (validation)\")\n\n# failure analysis: rank a sample of cracked images by how badly the U-Net disagrees with the truth\nscan = sample_n(positives, 200)\nscores = []\nfor _, row in tqdm(scan.iterrows(), total=len(scan), desc=\"scanning for failures\"):\n    image = eval_tf(image=np.asarray(Image.open(row.image_path).convert(\"RGB\")))[\"image\"][None].to(DEVICE)\n    with torch.no_grad():\n        prediction = (torch.sigmoid(unet(image))[0, 0].float().cpu().numpy() >= UNET_THRESHOLD).astype(np.uint8)\n    truth = load_gt_mask(row.stem, (IMG_SIZE, IMG_SIZE))\n    scores.append((int((prediction & truth).sum()) / max(int((prediction | truth).sum()), 1), row.stem))\n\nworst = [stem for _, stem in sorted(scores)[:5]]\nshow_predictions(val_df[val_df.stem.isin(worst)], f\"Worst {len(worst)} of {len(scan)} sampled cracked images\")\n"}, {"cell_type": "code", "execution_count": null, "id": "789b88c3", "metadata": {}, "outputs": [], "source": "# The sweep below is an oracle curve: picking its best point would be choosing a hyperparameter\n# with masks. It runs on validation and it is only here to show how much the mask-free\n# calibrated threshold gives up against the best one could have done knowing the answer.\nthresholds = np.linspace(0.05, 0.95, 19)\nsweep_tp = np.zeros_like(thresholds)\nsweep_fp = np.zeros_like(thresholds)\nsweep_fn = np.zeros_like(thresholds)\nper_source = {}\n\nsample_df = val_df.sample(min(400, len(val_df)), random_state=SEED).reset_index(drop=True)\nfor batch in tqdm(make_loader(sample_df, eval_tf, with_raw=True), desc=\"threshold sweep\", leave=False):\n    probability, gate = predict_probability(\"unet\", batch)\n    for i, stem in enumerate(batch[\"stem\"]):\n        truth = load_gt_mask(stem, (IMG_SIZE, IMG_SIZE)).astype(bool)\n        prob = probability[i] * gate[i]\n        for k, threshold in enumerate(thresholds):\n            prediction = prob >= threshold\n            hit = int((prediction & truth).sum())\n            sweep_tp[k] += hit\n            sweep_fp[k] += int(prediction.sum()) - hit\n            sweep_fn[k] += int(truth.sum()) - hit\n        if truth.any():\n            prediction = prob >= UNET_THRESHOLD\n            hit = int((prediction & truth).sum())\n            group = stem.split(\"_\")[0].lower()\n            per_source.setdefault(group, []).append(hit / max(int((prediction | truth).sum()), 1))\n\nsweep_iou = sweep_tp / np.maximum(sweep_tp + sweep_fp + sweep_fn, 1)\nsource_iou = pd.Series({k: np.mean(v) for k, v in per_source.items()}).sort_values()\nsource_n = pd.Series({k: len(v) for k, v in per_source.items()})\n\nfig, axes = plt.subplots(1, 2, figsize=(14, 4.2))\naxes[0].plot(thresholds, sweep_iou, \"o-\")\naxes[0].axvline(UNET_THRESHOLD, color=\"crimson\", ls=\"--\", label=f\"mask-free choice ({UNET_THRESHOLD:.2f})\")\naxes[0].scatter([thresholds[sweep_iou.argmax()]], [sweep_iou.max()], color=\"green\", zorder=5,\n                label=f\"oracle best = {sweep_iou.max():.3f} @ {thresholds[sweep_iou.argmax()]:.2f}\")\naxes[0].set_xlabel(\"binarisation threshold\")\naxes[0].set_ylabel(\"dataset IoU\")\naxes[0].set_title(\"How much does the threshold matter?\")\naxes[0].legend(fontsize=8)\naxes[0].grid(alpha=0.3)\n\naxes[1].barh(source_iou.index, source_iou.values, color=\"#3182bd\")\nfor i, (name, value) in enumerate(source_iou.items()):\n    axes[1].text(value, i, f\"  {value:.2f}  (n={source_n[name]})\", va=\"center\", fontsize=8)\naxes[1].set_xlim(0, min(1.0, source_iou.max() * 1.45))\naxes[1].set_xlabel(\"mean image IoU\")\naxes[1].set_title(\"Where does it work? (per source sub-dataset)\")\nplt.tight_layout()\nplt.show()\n\nprint(f\"oracle would reach IoU {sweep_iou.max():.4f} at {thresholds[sweep_iou.argmax()]:.2f}, \"\nf\"the mask-free threshold gets {sweep_iou[np.argmin(np.abs(thresholds - UNET_THRESHOLD))]:.4f}\")\n"}, {"cell_type": "markdown", "id": "e61cf0d1", "metadata": {}, "source": "## 6. Analysis and conclusions\n\n### What each step gave us\n\nGoing down the ablation table, one change per row:\n\n1. Single scale CAM is the baseline. The location is right, the boundaries are not. At stride 8 a\n   3 pixel crack is a fraction of one cell.\n2. Multi-scale and flip TTA costs no extra training, only around 12x the inference time, and the\n   large scale is what recovers the thin structure. The U-Net gets a cheaper 4x flip average at inference.\n3. The guided filter is where the map stops looking like a blob. It is also the step where most\n   people would put a domain specific edge detector, and using a generic one is what keeps the\n   pipeline reusable. It reliably buys boundary F1, which is what it is for. Whether it also buys\n   IoU is less clear: it sharpens the map onto image edges, and on a textured surface some of those\n   edges are not the crack, so it can trade a little area agreement for a better placed contour.\n   The pseudo-masks are built from the refined map either way, because the U-Net cares much more\n   about where the boundary is than about a fraction of a percent of overlap.\n4. U-Net distillation with self-training is the biggest single jump, and the student ends up better\n   than the teacher it was trained on. The CAM errors are inconsistent between images while the real\n   signal is not, and the ignore band leaves the student room to correct the teacher.\n\nEverything in this section runs on the validation split. The test split is scored exactly once, by\nthe table in section 5, and hunting through test images for the interesting ones would quietly undo\nthat. `load_gt_mask` counts every mask it opens and section 5 asserts that the count for the\ntraining split is still zero.\n\n### The threshold experiment that did not work\n\nThe U-Net binarises at a flat 0.5 while every CAM variant uses a per image quantile that keeps about\n2% of the pixels. That looked like an inconsistency worth removing, so we tried picking the U-Net\nthreshold on validation as the one whose average predicted area also comes out at 2%. It is a legal\nchoice, it needs no masks, and it was clearly worse: the threshold landed at 0.99 and IoU fell from\n0.246 to 0.226.\n\nThe sweep below is what explains it. The best threshold on this model is around 0.2, not 0.99, so\nthe area prior was pushing in exactly the wrong direction. The reason is that the pseudo-masks are\ncut at the 2% quantile deliberately, to keep the teacher from over-claiming, and part of what the\nstudent gains over the teacher is drawing the crack wider and more complete than the teacher dared.\nPutting the same 2% back at the decision boundary throws that gain away. The size prior belongs in\nstage 2, where it stops a noisy CAM from flooding the mask, and not at the end of the pipeline.\n\nSo 0.5 stays. It is not arbitrary: it is the boundary the BCE half of the loss is trained around,\nand it needs nothing except the model. The oracle curve says a threshold chosen with the masks in\nhand would be worth roughly another 0.02 IoU, which is the honest size of what we are giving up by\nrefusing to look.\n\n### Where it fails\n\n* Thin low contrast cracks on textured concrete. The classifier can be right for the wrong reason\n  (the texture) and then the CAM covers a region instead of a line.\n* Shadows, joints and painted lines look the same as a crack at the image level. They are most of\n  our false positives and an image level label cannot tell them apart.\n* Cracks touching the border are under segmented. The puzzle term helps but the border tiles still\n  see the least context.\n* Recall is higher than precision everywhere. With only a bag level label the model learns where the\n  evidence is, and the evidence is thicker than the crack itself. The area prior threshold and the\n  asymmetric Tversky loss pull it back a long way, but they do not close the gap, they only stop it\n  from compounding at every stage.\n* The crack free images almost all come from one source sub-dataset, so the classifier can score well\n  by recognising the source instead of the crack. The heavy photometric augmentation and the balanced\n  sampler are there to make that shortcut harder, and the per source table is how we check whether it\n  worked.\n\n### Compared to full supervision\n\nA fully supervised U-Net on this dataset gets around 0.65-0.75 IoU. Getting a good part of that with\nno pixel annotations at all is the point of the exercise, and the annotation cost is one click per\nimage instead of a traced mask.\n\n### Using it on something else\n\nNothing above is crack specific. To run it on \"find the tumour cells given only healthy/sick slide\nlabels\":\n\n1. Point `DATA_ROOT` at the new dataset and give it `image_path, label`.\n2. Set `TOPK_RATIO` to a rough guess of how much of the image the target covers, and\n   `W_SPARSITY = 0` if the target is large.\n3. Nothing else changes.\n\n### What we would try next\n\n* A learned pixel affinity network (IRNet / AffinityNet) instead of the guided filter, to spread the\n  CAM seeds along learned boundaries instead of image gradients.\n* Contrastive scale equivariance (SEAM) next to the puzzle consistency.\n* Training the classifier at output stride 4 with a bigger crop. The stride 8 map is still the\n  tightest bottleneck in the whole pipeline.\n* A leave one source out run: train on every sub-dataset but one and test on the one left out. The\n  per source table says the classifier is right almost everywhere, but nine of the eleven sources\n  contain only cracked images, so that table cannot separate \"knows what a crack is\" from \"knows\n  which folder this came from\". Leave one source out is the experiment that could, and it costs one\n  full training run per source, which is why it is future work and not a result.\n* A domain adversarial head on the source label with a gradient reversal layer. We left it out on\n  purpose: it is the standard answer to a shortcut, but there is no measured gap for it to close\n  yet, and adding a regulariser against a problem we have not demonstrated would be method for its\n  own sake. It belongs after the leave one source out run, not before it."}, {"cell_type": "markdown", "metadata": {}, "source": "## 7. Kaggle submission\n\nThe competition test set is a separate folder of images with no masks and no labels, so here the\npipeline runs the way it would in deployment: the classifier decides whether there is a crack at all\nand only then does the U-Net say where it is.\n\nThe mask has to be sent at the native resolution of each test image, so we upsample the probability\nmap from 448x448 before thresholding. Images the classifier calls negative get an empty string,\nwhich is what the sample submission uses for \"no crack\".", "id": "c3c7b03f"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "COMP_DIR = Path(os.environ.get(\"COMP_DIR\", \"/kaggle/input/dl-2025-project-2-pro\"))\nTEST_IMAGE_DIR = COMP_DIR / \"test\" / \"images\"\nSAMPLE_CSV = COMP_DIR / \"sample_submission.csv\"\n\n# if the competition data is not attached, fall back to the same test images in the dataset\nif not TEST_IMAGE_DIR.is_dir():\n    TEST_IMAGE_DIR = DATA_ROOT / \"test\" / \"images\"\n\n\n# https://www.kaggle.com/paulorzp/rle-functions-run-lenght-encode-decode\ndef mask2rle(img):\n    \"\"\"Binary mask -> the run length string the competition expects.\"\"\"\n    pixels = img.T.flatten()\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    return \" \".join(str(x) for x in runs)\n\n\ndef rle2mask(rle, shape):\n    \"\"\"Inverse of mask2rle, only used to check that the encoding round-trips.\"\"\"\n    mask = np.zeros(shape[0] * shape[1], np.uint8)\n    if isinstance(rle, str) and rle.strip():\n        numbers = [int(x) for x in rle.split()]\n        for start, length in zip(numbers[::2], numbers[1::2], strict=True):\n            mask[start - 1:start - 1 + length] = 1\n    return mask.reshape(shape[1], shape[0]).T\n\n\ncheck = np.zeros((40, 30), np.uint8)\ncheck[5:12, 8:20] = 1\ncheck[30, :] = 1\nassert (rle2mask(mask2rle(check), check.shape) == check).all()\nassert mask2rle(np.zeros((10, 10), np.uint8)) == \"\"\nprint(\"rle round-trip ok\")", "id": "1e3be039"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "class CompetitionTestSet(Dataset):\n    \"\"\"Test images only - no mask, no label, we resize and remember the original size.\"\"\"\n\n    def __init__(self, files, image_dir, transform):\n        self.files = list(files)\n        self.image_dir = Path(image_dir)\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.files)\n\n    def __getitem__(self, idx):\n        name = self.files[idx]\n        image = np.asarray(Image.open(self.image_dir / name).convert(\"RGB\"))\n        height, width = image.shape[:2]\n        return {\"image\": self.transform(image=image)[\"image\"], \"file\": name, \"height\": height, \"width\": width}\n\n\n@torch.no_grad()\ndef predict_test_masks(files, batch_size=BATCH_SIZE):\n    \"\"\"Classifier gate first, then the U-Net mask at the image's own resolution.\"\"\"\n    model.eval()\n    unet.eval()\n    loader = DataLoader(\n        CompetitionTestSet(files, TEST_IMAGE_DIR, seg_eval_tf),\n        batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=torch.cuda.is_available(),\n    )\n\n    masks = {}\n    for batch in tqdm(loader, desc=\"predicting the test set\"):\n        images = batch[\"image\"].to(DEVICE)\n        gate = torch.sigmoid(model(images)[\"logit\"].squeeze(1)).cpu().numpy() >= GATE_THRESHOLD\n        probability = unet_probability(unet, images).cpu().numpy()\n\n        for i, name in enumerate(batch[\"file\"]):\n            height, width = int(batch[\"height\"][i]), int(batch[\"width\"][i])\n            if not gate[i]:\n                masks[name] = np.zeros((height, width), np.uint8)\n                continue\n            resized = cv2.resize(probability[i], (width, height), interpolation=cv2.INTER_LINEAR)\n            masks[name] = (resized >= UNET_THRESHOLD).astype(np.uint8)\n    return masks", "id": "6c5fea58"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "if not TEST_IMAGE_DIR.exists():\n    print(f\"{TEST_IMAGE_DIR} not found - attach the competition data to the notebook to submit\")\nelse:\n    # sample_submission.csv only has a few example rows, so all we take from it is the column\n    # names and the fact that the file column keeps the .jpg extension\n    if SAMPLE_CSV.exists():\n        print(pd.read_csv(SAMPLE_CSV).head(3).to_string(max_colwidth=40))\n\n    test_files = sorted(p.name for p in TEST_IMAGE_DIR.iterdir() if p.is_file())\n    print(f\"\\n{len(test_files)} test images\")\n\n    set_seed()\n    test_masks = predict_test_masks(test_files)\n\n    submission = pd.DataFrame({\n        \"file\": test_files,\n        \"crack\": [mask2rle(test_masks[name]) for name in test_files],\n    })\n    submission.to_csv(WORK_DIR / \"submission.csv\", index=False)\n\n    positives = int((submission.crack.str.len() > 0).sum())\n    coverage = float(np.mean([m.mean() for m in test_masks.values()]))\n    print(f\"wrote {WORK_DIR / 'submission.csv'}\")\n    print(f\"{positives}/{len(submission)} images predicted as cracked ({positives / len(submission):.1%})\")\n    print(f\"average predicted crack area: {coverage:.2%} of the image\")\n    stage_done(\"submission\")", "id": "dbb950e8"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# eyeball a few test predictions - there is no ground truth here, so this is all we can check\nif TEST_IMAGE_DIR.exists():\n    preview = list(test_masks)[:6]\n    fig, axes = plt.subplots(2, len(preview), figsize=(2.8 * len(preview), 6))\n    for j, name in enumerate(preview):\n        axes[0][j].imshow(Image.open(TEST_IMAGE_DIR / name).convert(\"RGB\"))\n        axes[0][j].set_title(name, fontsize=8)\n        axes[1][j].imshow(test_masks[name], cmap=\"gray\", vmin=0, vmax=1)\n        for row in (0, 1):\n            axes[row][j].axis(\"off\")\n    plt.tight_layout()\n    plt.show()", "id": "fad2dc91"}, {"cell_type": "markdown", "id": "64438c47-d501-4fab-9838-637edb1c65c0", "metadata": {}, "source": "## References\n\nMethods, in the order they appear in the notebook.\n\n* Jo, S. and Yu, I.-J. (2021). *Puzzle-CAM: improved localization via matching partial and full\n  features.* ICIP. The tiled consistency term in stage 1.\n* Zhou, Z.-H. (2004). *Multi-instance learning: a survey.* And Ilse, M., Tomczak, J. and Welling, M.\n  (2018), *Attention-based deep multiple instance learning*, ICML, for the bag/instance framing and\n  the case for pooling that is neither mean nor max.\n* Zhou, B. et al. (2016). *Learning deep features for discriminative localization.* CVPR. Class\n  activation maps, and the reason the 1x1 head carries no bias.\n* He, K., Sun, J. and Tang, X. (2013). *Guided image filtering.* TPAMI. The edge preserving filter in\n  stage 2.\n* Salehi, S. S. M., Erdogmus, D. and Gholipour, A. (2017). *Tversky loss function for image\n  segmentation using 3D fully convolutional deep networks.* MLMI. The asymmetric region loss.\n* Shit, S. et al. (2021). *clDice - a novel topology-preserving loss function for tubular structure\n  segmentation.* CVPR. The soft skeleton and the centreline loss.\n* Ronneberger, O., Fischer, P. and Brox, T. (2015). *U-Net: convolutional networks for biomedical\n  image segmentation.* MICCAI.\n* He, K. et al. (2016). *Deep residual learning for image recognition.* CVPR. The ResNet-34 encoder,\n  ImageNet weights from torchvision.\n* Chen, L.-C. et al. (2017). *Rethinking atrous convolution for semantic image segmentation.* The\n  dilated trunk that keeps the score map at stride 8.\n* Ahn, J., Cho, S. and Kwak, S. (2019). *Weakly supervised learning of instance segmentation with\n  inter-pixel relations.* CVPR. IRNet, listed above as the alternative to the guided filter.\n* Wang, Y. et al. (2020). *Self-supervised equivariant attention mechanism for weakly supervised\n  semantic segmentation.* CVPR. SEAM, listed above as future work.\n* Dataset: [Crack Segmentation Dataset](https://www.kaggle.com/datasets/lakshaymiddha/crack-segmentation-dataset),\n  a merge of CFD, DeepCrack, Rissbilder, GAPs, CrackTree, Volker, Eugen Muller, Sylvie Chambon,\n  forest and a set of crack free surfaces.\n"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python"}}, "nbformat": 4, "nbformat_minor": 5}