Preparing Electron Microscopy Segmentation Training and Held-Out Datasets

Hello. In the preceding lesson, you treated a segmentation as a hypothesis grounded in the underlying EM imagery: difficult membranes, thin processes, crossings, and image defects can lead to false splits or false merges. We also separated FlyWire’s production pipeline from the later human proofreading and connectome-materialization stages.

Now we move to a hands-on proxy for the model portion of that pipeline. You will prepare a small, labeled EM dataset for a lightweight segmentation experiment. The aim is not to retrain FlyWire’s production system. It is to create correctly aligned, versioned arrays that a local training pipeline can consume and that you can evaluate honestly in the next lesson.

By the end, you will have:

• a training EM block from CREMI Dataset A;
• a fully separate held-out EM block from CREMI Dataset B;
• raw grayscale arrays, preserved neuron-instance IDs, and a derived binary boundary target;
• compressed NumPy archives and HDF5 files, with crop and resolution provenance.

---

What “pre-labeled EM data” means here

The CREMI dataset contains paired three-dimensional volumes:

1. a grayscale EM image volume, where each voxel is an observed pixel intensity;
2. a neuron-label volume, where each voxel contains an integer identifying the annotated neuron instance.

The source data use the spatial axis order:

$$(z,y,x)$$

Here, $z$ is the image-section index, while $y$ and $x$ are the within-section pixel coordinates. Do not casually transpose these axes. A model, a visual inspection step, and any later interpretation of anisotropic voxel resolution all depend on consistent axis meaning.

The colors in a segmentation viewer can be misleading: the stored label volume is not a colored image. It is an instance-label array such as:

$$[0,0,17,17,17,402,402]$$

The numeric values are identifiers, not ordered classes. Neuron ID 402 is not “more similar” to neuron ID 17 than it is to neuron ID 900. Therefore:

• never normalize or interpolate instance IDs;
• preserve the IDs as integer data;
• use nearest-neighbor behavior if you ever resample labels;
• derive a model target appropriate to the model’s task.

For the lightweight experiment in this course, we will preserve the instance IDs and additionally derive a binary boundary target. A target voxel is $1$ where two nonzero neighboring neuron IDs differ, and $0$ elsewhere. This gives a baseline model a clear, modest task: predict likely boundaries between labeled neuronal regions.

---

Inspect the dataset format and choose a sound split

Read the CREMI data specification before downloading. It identifies the exact HDF5 dataset paths and distinguishes the raw image volume from neuron annotations.

Why use two volumes rather than randomly splitting slices?

A tempting shortcut is to take one 3D volume and randomly assign individual $z$-slices to training and validation. Do not do that.

Adjacent EM sections are highly correlated. A held-out slice may differ from a training slice by only one $40$-nm step in depth, showing almost the same membranes and organelles. That produces an overly optimistic evaluation: the model has effectively seen the local tissue already.

Instead, use this split:

Role	Source	Use
Training	CREMI Dataset A	Fit model parameters in the next lesson
Held-out	CREMI Dataset B	Evaluate only after training decisions are complete
Reserved	CREMI Dataset C	Leave untouched for later experiments

Dataset B is not a true blind benchmark in the strict competition sense, because you possess its labels. It is, however, a much better held-out evaluation volume than a random selection of slices from Dataset A.

Create a small project layout

Make a fresh local project directory. Keep source data separate from generated data so that you can always rebuild the prepared arrays.

em-segmentation-proxy/
├── data/
│   ├── source/
│   │   ├── sample_A_20160501.hdf
│   │   └── sample_B_20160501.hdf
│   └── prepared/
├── notebooks/
│   └── 01_prepare_cremi.ipynb
└── requirements.txt

Your minimal notebook environment needs:

pip install numpy h5py matplotlib

The h5py package reads the HDF5 format used by CREMI. NumPy will hold the arrays, and Matplotlib provides a quick visual inspection before you commit to training.

---

Build paired training and held-out arrays

Create notebooks/01_prepare_cremi.ipynb. The following cells form a compact, reproducible preparation workflow.

1. Load and inspect a CREMI file

Start by listing the relevant dataset paths, shapes, data types, and spatial resolutions. This is a cheap check that can prevent training on the wrong array or silently misaligning labels.

from pathlib import Path

import h5py
import matplotlib.pyplot as plt
import numpy as np

PROJECT = Path("..")
SOURCE = PROJECT / "data" / "source"
PREPARED = PROJECT / "data" / "prepared"
PREPARED.mkdir(parents=True, exist_ok=True)

A_PATH = SOURCE / "sample_A_20160501.hdf"
B_PATH = SOURCE / "sample_B_20160501.hdf"

RAW_KEY = "volumes/raw"
LABEL_KEY = "volumes/labels/neuron_ids"


def inspect_cremi(path: Path) -> None:
    with h5py.File(path, "r") as f:
        raw = f[RAW_KEY]
        labels = f[LABEL_KEY]

        print(f"File: {path.name}")
        print(f"Raw:    shape={raw.shape}, dtype={raw.dtype}")
        print(f"Labels: shape={labels.shape}, dtype={labels.dtype}")
        print(f"Resolution (z, y, x) nm: {raw.attrs.get('resolution')}")
        print(f"Raw range: {raw[:].min()} to {raw[:].max()}")
        print(f"Label IDs in a small probe: {np.unique(labels[:4, :64, :64])[:20]}")


inspect_cremi(A_PATH)
inspect_cremi(B_PATH)

For the cropped files, the raw and label volumes should have matching shapes. The padded files can contain extra raw-image context and require offset handling, so the workflow deliberately excludes them for now.

The expected invariants are:

with h5py.File(A_PATH, "r") as f:
    assert f[RAW_KEY].shape == f[LABEL_KEY].shape
    assert np.issubdtype(f[LABEL_KEY].dtype, np.integer)

If either assertion fails, stop and determine whether you downloaded the padded file or selected the wrong HDF5 key. Do not “fix” a shape mismatch by trimming one array arbitrarily: images and labels must represent precisely the same physical voxels.

2. Normalize only the raw image

The raw EM array is image intensity, so converting it to float32 in the range $[0,1]$ is appropriate. The neuron IDs are categorical instance identifiers, so they must remain integers.

def normalize_raw(raw: np.ndarray) -> np.ndarray:
    """Convert integer EM intensities to float32 in the range [0, 1]."""
    if not np.issubdtype(raw.dtype, np.integer):
        raise TypeError(f"Expected integer raw EM data, got {raw.dtype}")

    max_value = np.iinfo(raw.dtype).max
    return raw.astype(np.float32) / max_value

This is deliberately different from independently applying min-max normalization to Dataset A and Dataset B. Per-volume min-max scaling would let the held-out volume influence its own preprocessing statistics and can conceal intensity-distribution differences. For the CREMI integer grayscale data, scaling by the known maximum integer value is deterministic and identical for both volumes.

3. Select the same spatially coherent crop from each volume

A block of shape $(64,256,256)$ is large enough to contain several structures but small enough for a laptop-scale experiment. It contains:

$$64\times 256\times 256=4,194,304$$

voxels per volume.

Because CREMI voxels are anisotropic, this is not physically cubic. The $z$-axis has much coarser resolution than $x$ and $y$. That is normal for serial-section EM; retain the source resolution metadata rather than pretending the voxels are isotropic.

CROP_SHAPE = (64, 256, 256)  # (z, y, x)


def centered_crop(array: np.ndarray, crop_shape: tuple[int, int, int]):
    """Return a centered 3D crop and its source-space origin."""
    if array.ndim != 3:
        raise ValueError(f"Expected a 3D array, got shape {array.shape}")

    if any(crop > size for crop, size in zip(crop_shape, array.shape)):
        raise ValueError(
            f"Crop {crop_shape} does not fit inside source shape {array.shape}"
        )

    origin = tuple((size - crop) // 2 for size, crop in zip(array.shape, crop_shape))
    slices = tuple(
        slice(start, start + size)
        for start, size in zip(origin, crop_shape)
    )

    return array[slices], origin

A centered crop is not biologically privileged; it is simply deterministic. The important part is recording the crop origin. If you later spot a poor prediction, you can locate its source coordinates rather than treating the prepared data as an anonymous tensor.

4. Preserve instances and derive a boundary target

A segmentation pipeline may predict affinities, distances, embeddings, or boundaries. Our small proxy will use boundaries, but preserve the original labels so that later experiments can derive a different target without re-downloading the data.

def instance_boundaries(
    instance_ids: np.ndarray,
    valid_mask: np.ndarray,
) -> np.ndarray:
    """
    Mark boundaries between distinct, nonzero neighboring instance IDs.

    A voxel is a boundary voxel when it differs from a valid immediate
    neighbor along z, y, or x. Unknown or unlabeled voxels are excluded.
    """
    boundary = np.zeros(instance_ids.shape, dtype=bool)

    for axis in range(3):
        lower = [slice(None)] * 3
        upper = [slice(None)] * 3
        lower[axis] = slice(0, -1)
        upper[axis] = slice(1, None)

        lower = tuple(lower)
        upper = tuple(upper)

        different_instances = (
            (instance_ids[lower] != instance_ids[upper])
            & valid_mask[lower]
            & valid_mask[upper]
        )

        boundary[lower] |= different_instances
        boundary[upper] |= different_instances

    return boundary.astype(np.uint8)

The valid_mask makes an explicit policy choice: label ID 0 is treated as unlabeled or background, not as a neuronal instance. This avoids turning every boundary around an unlabeled region into a positive training target. Keep that policy visible in the notebook; it is part of the dataset definition, not an incidental implementation detail.

Now combine loading, cropping, validation, and target creation:

def prepare_volume(path: Path, crop_shape=CROP_SHAPE) -> dict:
    with h5py.File(path, "r") as f:
        raw_full = f[RAW_KEY][:]
        ids_full = f[LABEL_KEY][:]
        resolution_nm = np.asarray(f[RAW_KEY].attrs["resolution"])

    if raw_full.shape != ids_full.shape:
        raise ValueError(
            f"Raw/label mismatch: raw={raw_full.shape}, labels={ids_full.shape}"
        )

    raw_crop, origin_zyx = centered_crop(raw_full, crop_shape)
    ids_crop, label_origin_zyx = centered_crop(ids_full, crop_shape)

    assert origin_zyx == label_origin_zyx
    assert np.issubdtype(ids_crop.dtype, np.integer)

    image = normalize_raw(raw_crop)
    valid = ids_crop != 0
    boundary = instance_boundaries(ids_crop, valid)

    return {
        "image": image,
        "instance_ids": ids_crop,
        "valid": valid.astype(np.uint8),
        "boundary": boundary,
        "origin_zyx": np.asarray(origin_zyx, dtype=np.int32),
        "resolution_nm_zyx": resolution_nm,
        "source_file": path.name,
    }


train = prepare_volume(A_PATH)
heldout = prepare_volume(B_PATH)

for name, sample in {"train": train, "heldout": heldout}.items():
    print(f"\n{name}")
    print("image shape:", sample["image"].shape)
    print("image dtype/range:", sample["image"].dtype,
          sample["image"].min(), sample["image"].max())
    print("instance-ID dtype:", sample["instance_ids"].dtype)
    print("valid-label fraction:", sample["valid"].mean())
    print("boundary fraction:", sample["boundary"].mean())
    print("crop origin (z, y, x):", sample["origin_zyx"])
    print("resolution nm (z, y, x):", sample["resolution_nm_zyx"])

    assert sample["image"].shape == sample["instance_ids"].shape
    assert sample["image"].shape == sample["boundary"].shape
    assert sample["boundary"].dtype == np.uint8
    assert set(np.unique(sample["boundary"])).issubset({0, 1})

The boundary fraction will usually be small. That is expected: most voxels are not precisely on a boundary. It also previews an issue for the training lesson: the model will see a class imbalance between non-boundary and boundary voxels.

---

Inspect alignment before writing files

Array shape checks are necessary, but they cannot prove that an image and its label are spatially aligned. View a central section with the binary target overlaid on the raw EM image.

def show_section(sample: dict, z: int | None = None) -> None:
    image = sample["image"]
    ids = sample["instance_ids"]
    boundary = sample["boundary"]

    if z is None:
        z = image.shape[0] // 2

    fig, axes = plt.subplots(1, 3, figsize=(15, 5))

    axes[0].imshow(image[z], cmap="gray", interpolation="nearest")
    axes[0].set_title(f"Raw EM, z={z}")
    axes[0].axis("off")

    axes[1].imshow(ids[z], interpolation="nearest")
    axes[1].set_title("Neuron instance IDs")
    axes[1].axis("off")

    axes[2].imshow(image[z], cmap="gray", interpolation="nearest")
    axes[2].imshow(
        boundary[z],
        cmap="magma",
        alpha=0.65,
        interpolation="nearest",
    )
    axes[2].set_title("Derived boundary overlay")
    axes[2].axis("off")

    plt.tight_layout()


show_section(train)
show_section(heldout)

Inspect several values of z, especially near the beginning and end of each crop. In a sensible overlay:

• boundary pixels generally follow transitions between differently labeled cellular regions;
• the overlay should not appear globally shifted relative to membranes;
• the label visualization should have crisp regions, not blended intermediate values;
• a noisy-looking label map may be normal because instance IDs receive arbitrary display colors.

This inspection is the bridge back to the previous lesson. A boundary target contains an annotated interpretation of the EM evidence; it does not eliminate ambiguity in faint, crowded, or damaged regions.

---

Save NumPy arrays and pipeline-friendly HDF5 files

Use NumPy archives for convenient notebook work and separate HDF5 files for a configuration-driven connectomics pipeline. The source labels remain in the files under a familiar CREMI-style path.

def save_prepared_sample(name: str, sample: dict) -> None:
    # Convenient single-file archive for notebook experiments.
    np.savez_compressed(
        PREPARED / f"{name}_arrays.npz",
        image=sample["image"],
        instance_ids=sample["instance_ids"],
        valid=sample["valid"],
        boundary=sample["boundary"],
        origin_zyx=sample["origin_zyx"],
        resolution_nm_zyx=sample["resolution_nm_zyx"],
        source_file=sample["source_file"],
    )

    # Separate raw-image HDF5 file.
    with h5py.File(PREPARED / f"{name}_raw.h5", "w") as f:
        raw_ds = f.create_dataset(
            RAW_KEY,
            data=sample["image"],
            compression="gzip",
        )
        raw_ds.attrs["resolution"] = sample["resolution_nm_zyx"]
        raw_ds.attrs["crop_origin_zyx"] = sample["origin_zyx"]
        raw_ds.attrs["source_file"] = sample["source_file"]
        raw_ds.attrs["normalization"] = "integer maximum scaling to [0, 1]"

    # Label HDF5 file: retain source instance IDs and proxy-training targets.
    with h5py.File(PREPARED / f"{name}_labels.h5", "w") as f:
        ids_ds = f.create_dataset(
            LABEL_KEY,
            data=sample["instance_ids"],
            compression="gzip",
        )
        ids_ds.attrs["resolution"] = sample["resolution_nm_zyx"]
        ids_ds.attrs["crop_origin_zyx"] = sample["origin_zyx"]
        ids_ds.attrs["source_file"] = sample["source_file"]

        f.create_dataset(
            "volumes/labels/boundary",
            data=sample["boundary"],
            compression="gzip",
        )
        f.create_dataset(
            "volumes/labels/valid",
            data=sample["valid"],
            compression="gzip",
        )


save_prepared_sample("train", train)
save_prepared_sample("heldout", heldout)

Verify that saved files can be read back. This guards against accidentally training on an in-memory variable that was never correctly persisted.

for name in ("train", "heldout"):
    with np.load(PREPARED / f"{name}_arrays.npz") as archive:
        print(name, archive["image"].shape, archive["boundary"].shape)

    with h5py.File(PREPARED / f"{name}_labels.h5", "r") as f:
        print(name, f[LABEL_KEY].shape, f["volumes/labels/boundary"].shape)

Your output directory should now contain:

data/prepared/
├── train_arrays.npz
├── train_raw.h5
├── train_labels.h5
├── heldout_arrays.npz
├── heldout_raw.h5
└── heldout_labels.h5

---

Connect the prepared data to a segmentation pipeline

The PyTorch Connectomics repository documents the general pattern: copy a close tutorial configuration, then point its training and validation image and label paths at your own files. Its broader view of the workflow also helps locate what this preparation stage does and does not accomplish.

Conceptually, your configuration should distinguish Dataset A from Dataset B:

data.train.image = data/prepared/train_raw.h5
data.train.label = data/prepared/train_labels.h5

data.val.image = data/prepared/heldout_raw.h5
data.val.label = data/prepared/heldout_labels.h5

Use the configuration closest to the specific model tutorial you run next, because pipelines differ in the exact HDF5 internal dataset key and target representation they expect. Your files contain:

HDF5 path	Meaning
volumes/raw	Float32 grayscale EM input, scaled to $[0,1]$
volumes/labels/neuron_ids	Original integer neuron-instance labels
volumes/labels/boundary	Derived binary boundary target for the proxy task
volumes/labels/valid	Mask indicating which voxels have usable nonzero instance labels

If the baseline is a boundary-prediction model, configure it to use volumes/labels/boundary and, if supported, volumes/labels/valid to exclude unlabeled regions from loss and metrics. Keep neuron_ids untouched for traceability and possible future instance-oriented experiments.

For a 2D baseline, each $z$-section can later be treated as one training item. For a 3D baseline, retain the current $(z,y,x)$ volume and sample smaller 3D patches during loading. Do not reshape data merely to make it look like a conventional RGB image dataset; EM stack geometry is part of the problem.

---

Preparation checklist

Before leaving the notebook, record these facts in a Markdown cell:

Dataset provenance
- Training source: CREMI cropped Dataset A
- Held-out source: CREMI cropped Dataset B
- Reserved source: CREMI Dataset C, unused

Spatial convention
- Array axis order: (z, y, x)
- Resolution: copied from source HDF5 metadata
- Crop shape: (64, 256, 256)
- Crop selection: centered, with stored source origin

Targets
- Source labels: integer neuron instance IDs
- Proxy target: binary boundary between distinct nonzero instances
- Unlabeled policy: neuron ID 0 is excluded by valid mask

Normalization
- Raw image: integer maximum scaling to float32 [0, 1]
- Labels: preserved as integers; never normalized or interpolated

This is lightweight provenance, but it is enough to reproduce the exact arrays later. It also makes it clear what is experimental design and what comes from the source dataset.

---

Key takeaways

A trainable EM dataset is not just “images plus masks.” Its reliability depends on keeping image and label voxels aligned, preserving instance-label semantics, recording coordinate conventions, and avoiding leakage between training and evaluation data.

You have prepared:

• a spatially coherent Dataset A training crop;
• a separate Dataset B held-out crop;
• grayscale EM inputs normalized consistently;
• integer neuron-instance labels kept intact;
• a derived binary boundary target and valid-label mask;
• NumPy and HDF5 outputs with resolution and crop-origin metadata.

Next, you will train a lightweight baseline segmenter on these arrays. You will then compute intersection-over-union on the held-out volume and inspect where its failures overlap with the image ambiguities discussed earlier: weak boundaries, crowded processes, and difficult continuity evidence.