Chapter 46: Segmentation Architectures

import numpy as np
from scipy import ndimage

from kaggle_companion.activities._common import clean

SEED = 46
DRAWS = 10                 # independent image sets; every estimate is a mean over draws
SIZE = 32                  # images are 32 x 32 pixels
N_DEV, N_HELD = 150, 300   # development images (choose the rule) and held-out images (score it)
THRESHOLDS = np.linspace(0.2, 0.9, 15)
AREAS = (0, 5, 10, 20)     # minimum object sizes in pixels tried by the post-processing rule


def generate(n, empty_share, rng):
    """Masks with one disc each (radius 2 to 8 px; a share of the images is empty) and a stand-in model's probability map.
    The map is the blurred mask plus spatially smooth noise, so it misses small objects and speckles over empty images."""
    yy, xx = np.mgrid[:SIZE, :SIZE]
    masks = np.zeros((n, SIZE, SIZE), bool)
    for i in np.where(rng.random(n) >= empty_share)[0]:
        radius = rng.choice([2, 3, 4, 6, 8], p=[0.25, 0.25, 0.2, 0.15, 0.15])
        cy, cx = rng.integers(radius + 1, SIZE - radius - 1, 2)
        masks[i] = (yy - cy) ** 2 + (xx - cx) ** 2 <= radius ** 2
    blur = ndimage.gaussian_filter(masks.astype(float), (0, 1.2, 1.2))
    field = ndimage.gaussian_filter(rng.normal(size=masks.shape), (0, 1.5, 1.5))
    field /= field.std()
    return masks, 1 / (1 + np.exp(-(9 * (blur - 0.35) + 1.44 * field - 1.0)))


def pooled_dice(pred, target):
    """Intersection and support summed over the whole set, then one ratio."""
    return (2 * (pred & target).sum() + 1e-6) / (pred.sum() + target.sum() + 1e-6)


def image_dice(pred, target):
    """One ratio per image, averaged later. An empty target with an empty prediction scores 1 (the rule stated in the text)."""
    inter = (pred & target).sum((1, 2))
    total = pred.sum((1, 2)) + target.sum((1, 2))
    return np.where(total == 0, 1.0, 2 * inter / np.maximum(total, 1))


def remove_small(pred, area):
    """Delete connected components smaller than `area` pixels."""
    if area == 0:
        return pred
    out = pred.copy()
    for i in np.where(pred.any((1, 2)))[0]:
        labels, count = ndimage.label(pred[i])
        for k, size in enumerate(ndimage.sum(pred[i], labels, range(1, count + 1)), 1):
            if size < area:
                out[i][labels == k] = False
    return out


def run(empty_share):
    rng = np.random.default_rng(SEED)
    keys = ("pooled_rule", "image_rule", "area_rule")
    rows = {k: {"pooled": [], "image": []} for k in keys}
    chosen = {"pooled_rule": [], "image_rule": [], "area_rule": [], "area": []}
    curve_pooled, curve_image = [], []
    for _ in range(DRAWS):
        dev_mask, dev_prob = generate(N_DEV, empty_share, rng)
        held_mask, held_prob = generate(N_HELD, empty_share, rng)
        # Rule 1: the threshold that maximises pooled Dice on the development images.
        t_pooled = THRESHOLDS[np.argmax([pooled_dice(dev_prob > t, dev_mask) for t in THRESHOLDS])]
        # Rule 2: the threshold that maximises mean per-image Dice (empty rule included).
        t_image = THRESHOLDS[np.argmax([image_dice(dev_prob > t, dev_mask).mean() for t in THRESHOLDS])]
        # Rule 3: threshold and minimum object size that maximise mean per-image Dice.
        best = max(((image_dice(remove_small(dev_prob > t, a), dev_mask).mean(), t, a) for a in AREAS for t in THRESHOLDS[::2]),
                   key=lambda s: s[0])
        preds = {"pooled_rule": held_prob > t_pooled, "image_rule": held_prob > t_image,
                 "area_rule": remove_small(held_prob > best[1], best[2])}
        for k, pred in preds.items():
            rows[k]["pooled"].append(pooled_dice(pred, held_mask))
            rows[k]["image"].append(image_dice(pred, held_mask).mean())
        chosen["pooled_rule"].append(t_pooled)
        chosen["image_rule"].append(t_image)
        chosen["area_rule"].append(best[1])
        chosen["area"].append(best[2])
        curve_pooled.append([pooled_dice(held_prob > t, held_mask) for t in THRESHOLDS])
        curve_image.append([image_dice(held_prob > t, held_mask).mean() for t in THRESHOLDS])
    gain = np.array(rows["area_rule"]["image"]) - np.array(rows["pooled_rule"]["image"])
    return clean({
        "empty_share": empty_share,
        "rules": {k: {"pooled": np.mean(v["pooled"]), "image": np.mean(v["image"])} for k, v in rows.items()},
        "threshold": {k: np.mean(chosen[k]) for k in keys}, "min_area": np.mean(chosen["area"]),
        "gain_image": gain.mean(), "gain_se": gain.std(ddof=1) / np.sqrt(DRAWS), "win_rate": float((gain > 0).mean()),
        "thresholds": THRESHOLDS, "curve_pooled": np.mean(curve_pooled, 0), "curve_image": np.mean(curve_image, 0),
        "draws": DRAWS, "dev_images": N_DEV, "held_images": N_HELD,
    })

import pandas as pd
from kaggle_companion.activities.ch46 import explain

VALUES = [0, 0.25, 0.5, 0.75]
results = {value: run(value) for value in VALUES}
table = pd.DataFrame({str(v): explain(results[v], v)['metrics'] for v in VALUES})
print(table.to_string())

import matplotlib.pyplot as plt
from kaggle_companion.activities.ch46 import draw, NCOLS, HEIGHT

DEFAULT = 0.5
fig, axes = plt.subplots(1, NCOLS, figsize=(6.6 if NCOLS == 1 else 10.4, HEIGHT), layout='constrained')
draw(axes, results[DEFAULT], DEFAULT)
plt.show()