Evaluating IoU and Interpreting Segmentation Errors Without Retraining FlyWire

Welcome back. Your small U-Net has now been trained only on the Dataset A portion of the proxy dataset, while Dataset B has remained untouched. That separation matters: this lesson is the first time the model will be tested against genuinely held-out EM imagery.

You will load the saved checkpoint, convert its boundary probabilities into a binary prediction with a predeclared threshold, calculate boundary Intersection over Union (IoU), and inspect cases where the model missed boundaries or invented them. Finally, you will connect this compact experiment to FlyWire’s real workflow without confusing this notebook with a retraining of FlyWire itself.

Plan for roughly 45 minutes: 10 minutes on IoU, 20 minutes running evaluation, 10 minutes inspecting errors, and 5 minutes interpreting what the result does and does not establish.

---

IoU measures overlap where it matters

For a binary segmentation task, IoU compares the pixels predicted as belonging to the positive class with the pixels actually labeled as belonging to that class. It is also called the Jaccard index.

In the previous lesson, the positive class was boundary pixel. Therefore, this is boundary IoU, not neuron-instance IoU. Let:

• $TP$ be pixels labeled as boundaries that the model also predicts as boundaries;
• $FP$ be pixels the model predicts as boundaries but the target does not;
• $FN$ be target boundary pixels that the model misses.

Then:

$${\mathrm{IoU}}_{\text{boundary}}=\frac{TP}{TP+FP+FN}$$

The denominator is the union: every pixel that either mask regards as boundary. True negatives do not appear. That is why IoU is useful here: the huge number of easy “not boundary” pixels cannot dominate the score.

A perfect overlap gives $\mathrm{IoU}=1$. A prediction that has no overlap with the true boundary map gives $\mathrm{IoU}=0$.

Why not report pixel accuracy as the main result?

Suppose boundaries occupy only 10% of valid pixels. A model that predicts “non-boundary” everywhere receives roughly 90% pixel accuracy, despite failing the only class we care about. IoU exposes that failure:

$$TP=0,FP=0,FN>0$$ $${\mathrm{IoU}}_{\text{boundary}}=0$$

This is the same class-imbalance concern that motivated weighted loss during training, but now viewed as an evaluation issue.

---

Make the held-out evaluation notebook

Create a third notebook:

notebooks/
└── 03_evaluate_heldout_iou.ipynb

Your project should now look like this:

em-segmentation-proxy/
├── data/
│   └── prepared/
│       ├── train_arrays.npz
│       └── heldout_arrays.npz
├── models/
│   ├── boundary_unet_v1.pt
│   └── boundary_unet_v1_metadata.json
└── notebooks/
    ├── 02_train_boundary_unet.ipynb
    └── 03_evaluate_heldout_iou.ipynb

Install the two small packages used for tabulation and metric cross-checking if needed:

pip install pandas scikit-learn

Start with imports and paths. Use the same device-selection logic as in the training notebook.

from pathlib import Path

import matplotlib.pyplot as plt
from matplotlib.patches import Patch
import numpy as np
import pandas as pd
import torch
from sklearn.metrics import jaccard_score
from torch import nn


if torch.cuda.is_available():
    DEVICE = torch.device("cuda")
elif torch.backends.mps.is_available():
    DEVICE = torch.device("mps")
else:
    DEVICE = torch.device("cpu")

PROJECT = Path("..")
PREPARED = PROJECT / "data" / "prepared"
MODELS = PROJECT / "models"

HELDOUT_ARCHIVE = PREPARED / "heldout_arrays.npz"
CHECKPOINT = MODELS / "boundary_unet_v1.pt"

print("Using device:", DEVICE)
print("Held-out archive:", HELDOUT_ARCHIVE)
print("Checkpoint:", CHECKPOINT)

This notebook must not train, optimize, or overwrite the checkpoint. It has one purpose: inference and evaluation on Dataset B.

Restore the model exactly

A PyTorch checkpoint holds parameters, not the Python class definition. Copy the ConvBlock and SmallUNet definitions unchanged from 02_train_boundary_unet.ipynb, then run this cell:

class ConvBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()

        self.layers = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(x)


class SmallUNet(nn.Module):
    def __init__(self):
        super().__init__()

        self.enc1 = ConvBlock(1, 8)
        self.pool1 = nn.MaxPool2d(kernel_size=2)

        self.enc2 = ConvBlock(8, 16)
        self.pool2 = nn.MaxPool2d(kernel_size=2)

        self.bottleneck = ConvBlock(16, 32)

        self.up2 = nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2)
        self.dec2 = ConvBlock(32, 16)

        self.up1 = nn.ConvTranspose2d(16, 8, kernel_size=2, stride=2)
        self.dec1 = ConvBlock(16, 8)

        self.output = nn.Conv2d(8, 1, kernel_size=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        skip1 = self.enc1(x)
        x = self.pool1(skip1)

        skip2 = self.enc2(x)
        x = self.pool2(skip2)

        x = self.bottleneck(x)

        x = self.up2(x)
        x = torch.cat([x, skip2], dim=1)
        x = self.dec2(x)

        x = self.up1(x)
        x = torch.cat([x, skip1], dim=1)
        x = self.dec1(x)

        return self.output(x)

Load the checkpoint and inspect its saved assumptions before making predictions:

checkpoint = torch.load(CHECKPOINT, map_location=DEVICE)

print("Checkpoint keys:", list(checkpoint))
print("Task:", checkpoint["training_config"]["task"])
print("Held-out policy:", checkpoint["training_config"]["heldout_dataset"])
print("Training epochs:", checkpoint["training_config"]["epochs"])

model = SmallUNet().to(DEVICE)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()

Now load the held-out volume. This should have the same aligned array structure as the training archive, but it comes from Dataset B.

with np.load(HELDOUT_ARCHIVE) as archive:
    image = archive["image"].astype(np.float32)
    boundary = archive["boundary"].astype(np.float32)
    valid = archive["valid"].astype(np.float32)

    heldout_provenance = {
        "source_file": str(archive["source_file"].item()),
        "crop_origin_zyx": archive["origin_zyx"].tolist(),
        "resolution_nm_zyx": archive["resolution_nm_zyx"].tolist(),
    }

print("Held-out image shape:", image.shape)
print("Held-out target shape:", boundary.shape)
print("Held-out valid-mask shape:", valid.shape)
print("Held-out provenance:", heldout_provenance)

assert image.ndim == 3
assert image.shape == boundary.shape == valid.shape
assert set(np.unique(boundary)).issubset({0.0, 1.0})
assert set(np.unique(valid)).issubset({0.0, 1.0})
assert np.isfinite(image).all()

The expected axis order remains:

$$(z,y,x)$$

The model expects a batch and a channel dimension, so a batch of $n$ sections must have shape:

$$(n,1,y,x)$$

---

Predict probabilities, then make a fixed binary decision

Your model produces logits. Apply sigmoid to obtain a probability-like boundary score between zero and one.

A threshold is required to calculate IoU. For this baseline, use a fixed threshold of $0.5$:

THRESHOLD = 0.5
INFERENCE_BATCH_SIZE = 4

There is an important evaluation discipline here. Do not try many thresholds on Dataset B and report only the one that looks best. That would tune your evaluation data and make the final result optimistic.

Because the prior lesson did not select a threshold on the development monitor, $0.5$ is a reasonable predeclared baseline convention. In a later iteration, you could choose a threshold on Dataset A development sections, freeze that choice, and then perform one final Dataset B evaluation.

@torch.inference_mode()
def predict_volume_probability(
    model: nn.Module,
    volume: np.ndarray,
    batch_size: int = 4,
) -> np.ndarray:
    """
    Run 2D inference section by section in small batches.

    Input volume shape:  (z, y, x)
    Output volume shape: (z, y, x), probabilities in [0, 1]
    """
    model.eval()

    probabilities = np.empty_like(volume, dtype=np.float32)

    for start in range(0, volume.shape[0], batch_size):
        stop = min(start + batch_size, volume.shape[0])

        x = torch.from_numpy(volume[start:stop, None, ...]).to(DEVICE)
        logits = model(x)

        probabilities[start:stop] = (
            torch.sigmoid(logits)
            .squeeze(1)
            .cpu()
            .numpy()
        )

    return probabilities


probability = predict_volume_probability(
    model,
    image,
    batch_size=INFERENCE_BATCH_SIZE,
)

prediction = probability >= THRESHOLD
target = boundary.astype(bool)
valid_mask = valid.astype(bool)

print("Probability range:", probability.min(), "to", probability.max())
print("Prediction boundary fraction:",
      prediction[valid_mask].mean())
print("Target boundary fraction:",
      target[valid_mask].mean())

assert probability.shape == image.shape
assert prediction.shape == target.shape

Compare the last two printed fractions. They need not be equal, but they should be plausible. If the predicted boundary fraction is exactly zero or nearly one, pause before trusting the metric. The model may have loaded incorrectly, the inputs may not match the training normalization, or the selected threshold may be unsuitable.

---

Calculate global and per-section IoU

There are two complementary summaries worth reporting.

1. Global boundary IoU pools every valid held-out pixel into one large comparison. Large sections or sections with more boundary pixels contribute more.
2. Mean per-section IoU calculates IoU separately for each EM section and then averages the sections. Each section receives equal weight.

Neither is inherently “the” correct score. Report both, label them clearly, and preserve the per-section table so the aggregate does not hide difficult locations.

def binary_counts(
    prediction: np.ndarray,
    target: np.ndarray,
    valid_mask: np.ndarray,
) -> dict[str, int]:
    """Count binary outcomes only within valid labeled pixels."""
    pred = prediction[valid_mask]
    truth = target[valid_mask]

    tp = int(np.logical_and(pred, truth).sum())
    fp = int(np.logical_and(pred, ~truth).sum())
    fn = int(np.logical_and(~pred, truth).sum())
    tn = int(np.logical_and(~pred, ~truth).sum())

    return {"tp": tp, "fp": fp, "fn": fn, "tn": tn}


def iou_from_counts(counts: dict[str, int]) -> float:
    union = counts["tp"] + counts["fp"] + counts["fn"]

    # A section with neither predicted nor target boundaries has no
    # meaningful positive-class overlap to assess.
    if union == 0:
        return np.nan

    return counts["tp"] / union


section_rows = []

for z in range(image.shape[0]):
    counts = binary_counts(
        prediction[z],
        target[z],
        valid_mask[z],
    )

    union = counts["tp"] + counts["fp"] + counts["fn"]
    valid_pixels = int(valid_mask[z].sum())

    section_rows.append(
        {
            "z": z,
            "valid_pixels": valid_pixels,
            "target_boundary_pixels": counts["tp"] + counts["fn"],
            "predicted_boundary_pixels": counts["tp"] + counts["fp"],
            "tp": counts["tp"],
            "fp": counts["fp"],
            "fn": counts["fn"],
            "tn": counts["tn"],
            "union": union,
            "boundary_iou": iou_from_counts(counts),
        }
    )

per_section = pd.DataFrame(section_rows)

global_counts = binary_counts(prediction, target, valid_mask)
global_iou = iou_from_counts(global_counts)

mean_section_iou = per_section["boundary_iou"].mean(skipna=True)

valid_pixels_total = int(valid_mask.sum())
pixel_accuracy = (
    (global_counts["tp"] + global_counts["tn"])
    / valid_pixels_total
)

print(f"Fixed threshold:         {THRESHOLD:.2f}")
print(f"Global boundary IoU:     {global_iou:.4f}")
print(f"Mean per-section IoU:    {mean_section_iou:.4f}")
print(f"Pixel accuracy:          {pixel_accuracy:.4f}")
print()
print("Global confusion counts:")
print(global_counts)

Use scikit-learn only as a cross-check. It should agree with your explicit global count calculation:

sklearn_iou = jaccard_score(
    target[valid_mask].astype(np.uint8),
    prediction[valid_mask].astype(np.uint8),
    average="binary",
    pos_label=1,
    zero_division=0,
)

print(f"scikit-learn Jaccard score: {sklearn_iou:.4f}")

assert np.isclose(global_iou, sklearn_iou)

The assertion is useful engineering practice. Two independent implementations that agree reduce the chance that a masking, dimension, or aggregation mistake has silently shaped the result.

Save the results rather than leaving them only in notebook output:

RESULTS = PROJECT / "results"
RESULTS.mkdir(exist_ok=True)

per_section.to_csv(
    RESULTS / "heldout_boundary_iou_per_section.csv",
    index=False,
)

summary = {
    "experiment_name": checkpoint["training_config"]["experiment_name"],
    "heldout_source": heldout_provenance,
    "threshold": THRESHOLD,
    "global_boundary_iou": float(global_iou),
    "mean_per_section_boundary_iou": float(mean_section_iou),
    "pixel_accuracy": float(pixel_accuracy),
    "global_counts": global_counts,
}

print(summary)

At this stage, avoid interpreting a single IoU value as a verdict on segmentation quality. The number is a compact summary of the overlap of thresholded boundary maps on one held-out crop. It cannot reveal whether errors are a few catastrophic locations, many slight boundary shifts, or annotation ambiguities.

---

Turn aggregate errors into inspectable EM cases

Start by examining the per-section table. Sort by IoU, but also look at false-positive and false-negative counts separately.

display_columns = [
    "z",
    "boundary_iou",
    "target_boundary_pixels",
    "predicted_boundary_pixels",
    "tp",
    "fp",
    "fn",
]

worst_sections = (
    per_section
    .dropna(subset=["boundary_iou"])
    .sort_values("boundary_iou")
    [display_columns]
    .head(5)
)

print("Five lowest-IoU held-out sections:")
display(worst_sections)

print("\nSections with the most missed target boundary pixels:")
display(
    per_section
    .sort_values("fn", ascending=False)
    [display_columns]
    .head(5)
)

print("\nSections with the most invented boundary pixels:")
display(
    per_section
    .sort_values("fp", ascending=False)
    [display_columns]
    .head(5)
)

Choose one or two representative sections. Do not automatically select only the worst numerical case. A useful small error set includes:

• a section with many false negatives;
• a section with many false positives;
• if available, a case showing a near-boundary displacement, where both error types trace the same structure.

The following helper makes these distinctions visible. Green pixels are correctly predicted boundaries; red pixels are false positives; cyan pixels are false negatives.

def plot_heldout_case(
    z: int,
    image: np.ndarray,
    target: np.ndarray,
    probability: np.ndarray,
    prediction: np.ndarray,
    valid_mask: np.ndarray,
    threshold: float,
) -> None:
    valid_z = valid_mask[z]
    target_z = target[z]
    pred_z = prediction[z]

    tp = valid_z & pred_z & target_z
    fp = valid_z & pred_z & ~target_z
    fn = valid_z & ~pred_z & target_z

    error_rgb = np.zeros((*target_z.shape, 3), dtype=np.float32)
    error_rgb[tp] = [0.0, 1.0, 0.0]   # green
    error_rgb[fp] = [1.0, 0.0, 0.0]   # red
    error_rgb[fn] = [0.0, 0.8, 1.0]   # cyan
    error_rgb[~valid_z] = [0.25, 0.25, 0.25]  # invalid region

    fig, axes = plt.subplots(1, 4, figsize=(18, 5))

    axes[0].imshow(image[z], cmap="gray", interpolation="nearest")
    axes[0].set_title(f"Held-out raw EM, z={z}")

    axes[1].imshow(
        target_z,
        cmap="magma",
        vmin=0,
        vmax=1,
        interpolation="nearest",
    )
    axes[1].set_title("Annotated boundary target")

    probability_plot = axes[2].imshow(
        probability[z],
        cmap="magma",
        vmin=0,
        vmax=1,
        interpolation="nearest",
    )
    axes[2].set_title(f"Predicted probability, threshold={threshold:.2f}")
    fig.colorbar(probability_plot, ax=axes[2], fraction=0.046)

    axes[3].imshow(error_rgb, interpolation="nearest")
    axes[3].set_title("Boundary comparison")
    axes[3].legend(
        handles=[
            Patch(color="lime", label="True positive"),
            Patch(color="red", label="False positive"),
            Patch(color="cyan", label="False negative"),
            Patch(color="gray", label="Invalid / ignored"),
        ],
        loc="lower right",
        fontsize=8,
    )

    for ax in axes:
        ax.axis("off")

    plt.tight_layout()
    plt.show()

Plot the lowest-IoU section first:

z_case = int(
    per_section
    .dropna(subset=["boundary_iou"])
    .sort_values("boundary_iou")
    .iloc[0]["z"]
)

plot_heldout_case(
    z_case,
    image,
    target,
    probability,
    prediction,
    valid_mask,
    THRESHOLD,
)

Read the error image as evidence, not merely colored pixels

Use the raw EM panel as the primary evidence. The target and prediction panels tell you what each system asserted, but they cannot by themselves determine whether the visual evidence was intrinsically clear.

Pattern in the comparison panel	Likely interpretation	Connectomics risk if propagated into segmentation
Cyan trace follows a visible membrane	The model missed a real boundary	Neighboring neurites may be joined, creating a false merge
Red trace occurs on internal texture, vesicles, or a dark organelle boundary	The model mistook texture for a cell boundary	A neurite may be fragmented, creating a false split
Red and cyan traces sit side by side along the same membrane	The boundary is slightly displaced	IoU can drop substantially even if the broad structure is recognized
A crowded or low-contrast region produces mixed red and cyan errors	Ambiguous local image evidence or imperfect target geometry	Treat the interpretation as uncertain; inspect adjacent sections
Errors concentrate at crop borders or invalid regions	Evaluation or preprocessing artifact may be involved	Do not convert this directly into a biological conclusion

For each representative case, look at its neighboring held-out sections. The model is 2D, but the EM data are serial. A membrane that is obscure in one section may become clear above or below it.

for z_neighbor in range(
    max(0, z_case - 1),
    min(image.shape[0], z_case + 2),
):
    plot_heldout_case(
        z_neighbor,
        image,
        target,
        probability,
        prediction,
        valid_mask,
        THRESHOLD,
    )

Keep a concise note in the notebook beneath each selected case. For example:

Case: Dataset B, z=12
Dominant error: false negatives along a faint oblique membrane.
Evidence: target boundary continues through z=11, z=12, and z=13;
the raw contrast is weakest at z=12.
Interpretation: likely image ambiguity rather than a global coordinate error.
Potential reconstruction consequence: local false-merge risk.
Confidence: moderate.

This habit is directly transferable to FlyWire proofreading: record location, evidence across sections, uncertainty, and likely consequence. The difference is that, in FlyWire, the next step is normally inspection of a specific reconstruction in the interface rather than retraining a local notebook model.

---

What this experiment says about FlyWire, and what it does not

The proxy pipeline is valuable because it exposes the basic causal chain:

1. Human-created labels define an image interpretation.
2. A model learns statistical patterns associated with those labels.
3. Held-out imagery tests whether the learned pattern transfers beyond training examples.
4. Prediction errors identify image situations likely to create reconstruction risk.
5. Human review remains necessary where image evidence is ambiguous or mistakes would be consequential.

That logic is relevant to FlyWire. But the notebook you built is not a reduced implementation of FlyWire’s production system, and its checkpoint cannot affect FlyWire data.

The important distinctions are concrete:

This proxy experiment	FlyWire production context
A compact 2D U-Net predicts boundary likelihood for one $256\times 256$ EM section at a time	Automated reconstruction operates on a much larger electron-microscopy volume and requires three-dimensional continuity
A small pre-labeled CREMI subset provides supervision	Production segmentation depends on a far larger, specialized data and systems pipeline
Output is a binary boundary map	Reconstruction requires forming segments or fragments and making decisions about their connectivity
No fragment generation or agglomeration is implemented	Automated segmentation must make or support connectivity decisions beyond pixelwise boundary prediction
Dataset B IoU evaluates a simplified proxy target	Real reconstruction quality also involves topology, branch continuity, merge and split errors, and biological plausibility
A local PyTorch checkpoint changes only local files	FlyWire data, segmentation layers, and versioned materializations are not altered by this notebook
You inspect 2D predictions manually	Proofreading uses 2D imagery, segmentation overlays, adjacent sections, and 3D morphology together

A particularly important caution follows from your task choice:

A boundary-model false negative is not automatically a FlyWire false merge, and a boundary-model false positive is not automatically a FlyWire false split.

Those are plausible risks, not one-to-one conclusions. A real reconstruction depends on subsequent segmentation and connectivity decisions, three-dimensional evidence, and human proofreading. Conversely, a small pixel displacement may produce both false-positive and false-negative boundary pixels, lowering IoU without necessarily causing a severe topological error.

Your result is therefore best described honestly as:

“A held-out evaluation of a lightweight 2D EM boundary-prediction baseline, used to study how local image ambiguity can create signals associated with split and merge risk.”

It is not evidence that you retrained FlyWire’s AI, improved its released segmentation, or measured the quality of an entire FlyWire neuron reconstruction.

---

Key takeaways

You have completed the evaluation stage of the segmentation proxy:

• IoU, or Jaccard index, measures overlap of predicted and annotated boundary pixels:

$$\mathrm{IoU}=\frac{TP}{TP+FP+FN}$$

• Boundary IoU is more informative than pixel accuracy when non-boundary pixels dominate.
• You evaluated the untouched Dataset B volume with a fixed threshold of $0.5$, avoiding threshold selection based on held-out results.
• You reported both global boundary IoU and mean per-section IoU, then saved a per-section results table.
• False negatives suggest missed boundaries and possible false-merge risk; false positives suggest invented boundaries and possible false-split risk.
• Visual inspection, including adjacent EM sections, is essential because one score cannot distinguish boundary shifts, texture confusion, and genuine ambiguity.
• The notebook illustrates the learning-and-evaluation logic behind automated segmentation, but it does not retrain, modify, or replace FlyWire’s production segmentation and proofreading workflow.

Next, you will begin the capstone module by formulating a narrowly scoped connectomics question that can be answered reproducibly with FlyWire data.