SegNeuron

https://github.com/yanchaoz/segneuron
Admin
SegNeuron is a generalist 3D neuron instance segmentation model designed to perform high-quality...
Tokens:9,021
Snippets:38
Trust Score:6
License:MIT
Update:2 months ago
Show doc for...
Tokens:
Raw
### Install Dependencies and Verify Environment

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Installs required libraries including python-elf for postprocessing and verifies PyTorch GPU availability.

```python
# Install dependencies that are not included in the default Colab runtime.
# python-elf provides the FRMC/multicut postprocessing used by SegNeuron.
%pip install -q "python-elf==0.9.1" "imageio>=2.31" "tifffile>=2023.7" "scikit-image>=0.20" tqdm

import sys, torch, numpy as np
print("Python:", sys.version.split()[0])
print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
if not torch.cuda.is_available():
    print("WARNING: No GPU is available. 3D inference will be very slow; switch to a GPU runtime.")
```

--------------------------------

### Run Pretraining

Source: https://github.com/yanchaoz/segneuron/blob/main/README.md

Navigate to the Pretrain directory and execute the pretraining script.

```bash
cd Pretrain
```

```bash
python pretrain.py
```

--------------------------------

### Configure Input Source

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Sets up file upload from the browser or mounts Google Drive based on user-defined boolean flags.

```python
#@title Input source
UPLOAD_FROM_BROWSER = False #@param {type:"boolean"}
USE_GOOGLE_DRIVE = False #@param {type:"boolean"}

if UPLOAD_FROM_BROWSER:
    from google.colab import files
    uploaded = files.upload()
    print("Uploaded:", list(uploaded))

if USE_GOOGLE_DRIVE:
    from google.colab import drive
    drive.mount("/content/drive")
    print("Google Drive mounted at /content/drive")
```

--------------------------------

### Run Instance Segmentation

Source: https://github.com/yanchaoz/segneuron/blob/main/README.md

Navigate to the post-processing directory and execute the instance segmentation script.

```bash
cd Postprocess
```

```bash
python FRMC_post.py
```

--------------------------------

### Run Supervised Training

Source: https://github.com/yanchaoz/segneuron/blob/main/README.md

Navigate to the training directory and execute the supervised training script.

```bash
cd Train_and_Inference
```

```bash
python supervised_train.py
```

--------------------------------

### Define Paths and Inference Parameters

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Configures file paths, crop sizes, and model parameters for the inference process.

```python
#@title Paths and parameters (edit, then run)
RAW_PATH = "/content/raw.tif" #@param {type:"string"}
CHECKPOINT_PATH = "/content/SegNeuronModel.ckpt" #@param {type:"string"}
OUTPUT_DIR = "/content/segneuron_results" #@param {type:"string"}

# Values are ordered as z/y/x. The y and x crop sizes should be multiples of 16.
# A stride must not exceed the corresponding crop size.
CROP_SIZE = (20, 128, 128)
STRIDE = (10, 64, 64)
BATCH_SIZE = 1 #@param {type:"integer"}
NORMALIZATION = "auto" #@param ["auto", "minmax", "none"]
WEIGHT_SIGMA = 0.5 #@param {type:"number"}
BETA = 0.25 #@param {type:"number"}

from pathlib import Path
Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
AFFINITY_OUTPUT = str(Path(OUTPUT_DIR) / "affinities.npy")
BOUNDARY_OUTPUT = str(Path(OUTPUT_DIR) / "boundaries.tif")
SEGMENTATION_OUTPUT = str(Path(OUTPUT_DIR) / "instances.npy")
print("Results will be written to", OUTPUT_DIR)
```

--------------------------------

### Archive and Download Results

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Packages the output directory into a ZIP file and triggers a browser download in Google Colab.

```python
#@title Download ZIP
DOWNLOAD_RESULTS = True #@param {type:"boolean"}

import shutil
archive = shutil.make_archive("/content/segneuron_results", "zip", OUTPUT_DIR)
print("Created", archive, round(Path(archive).stat().st_size / 1024**2, 1), "MiB")
if DOWNLOAD_RESULTS:
    from google.colab import files
    files.download(archive)
```

--------------------------------

### Run Affinity Inference

Source: https://github.com/yanchaoz/segneuron/blob/main/README.md

Navigate to the inference directory and execute the affinity inference script.

```bash
cd Train_and_Inference
```

```bash
python inference.py
```

--------------------------------

### Execute Sliding-Window Inference

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Performs inference over a volume using a sliding window approach, accumulating results and saving outputs to disk.

```python
@torch.inference_mode()
def predict(raw_path, checkpoint_path, affinity_output, boundary_output,
            crop_size=(20,128,128), stride=(10,64,64), batch_size=1,
            normalization="auto", weight_sigma=0.5, device=None):
    device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
    dataset = SlidingWindowVolume(raw_path, crop_size, stride, normalization, weight_sigma)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False,
                        num_workers=0, pin_memory=device.type == "cuda")
    model = load_model(checkpoint_path, device)
    affinity_sum = affinity_weight = boundary_sum = boundary_weight = None
    for inputs, indices in tqdm(loader, desc="SegNeuron inference"):
        affinities, boundaries = model(inputs.to(device, non_blocking=True))
        affinities, boundaries = affinities.cpu().numpy(), boundaries.cpu().numpy()
        if affinity_sum is None:
            affinity_sum, affinity_weight = dataset.accumulator(affinities.shape[1])
            boundary_sum, boundary_weight = dataset.accumulator(boundaries.shape[1])
        for affinity, boundary, index in zip(affinities, boundaries, indices.tolist()):
            dataset.add(affinity_sum, affinity_weight, affinity, index)
            dataset.add(boundary_sum, boundary_weight, boundary, index)
    affinity = dataset.finish(affinity_sum, affinity_weight)
    boundary = dataset.finish(boundary_sum, boundary_weight)
    np.save(affinity_output, affinity)
    tifffile.imwrite(boundary_output, boundary[0], photometric="minisblack")
    print("Affinity:", affinity.shape, affinity.dtype, "->", affinity_output)
    print("Boundary:", boundary.shape, boundary.dtype, "->", boundary_output)
    return affinity, boundary
```

--------------------------------

### Extract and Load Model State

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Utility functions to normalize checkpoint state dictionaries and initialize the MNet model architecture.

```python
def checkpoint_state_dict(checkpoint):
    if isinstance(checkpoint, dict):
        for key in ("model_weights", "model_state_dict", "state_dict", "model"):
            if key in checkpoint and isinstance(checkpoint[key], dict):
                checkpoint = checkpoint[key]
                break
    if not isinstance(checkpoint, dict):
        raise ValueError("Checkpoint does not contain a state dictionary")
    return OrderedDict((key.removeprefix("module."), value) for key, value in checkpoint.items())


def load_model(checkpoint_path, device):
    checkpoint_path = Path(checkpoint_path).expanduser()
    if not checkpoint_path.is_file():
        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
    model = MNet(1, kn=(32, 64, 96, 128, 256), FMU="sub")
    state = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(checkpoint_state_dict(state))
    return model.to(device).eval()
```

--------------------------------

### Initialize MNet Model Layers

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Defines the structural components of the MNet model, including upsampling, downsampling, and output convolution layers.

```python
        self.up31_ = Up(fct * (kn[3] + kn[4]), kn[3], ('both', 'both'), FMU)
        self.up32_ = Up(fct * (kn[2] + kn[3]), kn[2], ('both', '3d'), FMU)

        self.down41 = Down(kn[2], kn[3], ('3d', 'both'), FMU)
        self.bottleneck4 = Down(fct * kn[3], kn[4], ('both', 'both'), FMU)
        self.up41 = Up(fct * (kn[3] + kn[4]), kn[3], ('both', '3d'), FMU)

        self.up41_ = Up(fct * (kn[3] + kn[4]), kn[3], ('both', '3d'), FMU)

        self.bottleneck5 = Down(kn[3], kn[4], ('3d', '3d'))

        self.outputs = nn.ModuleList(
            [nn.Conv3d(c, 3, kernel_size=(1, 1, 1), stride=1, padding=0, bias=False)
             for c in [kn[0], kn[1], kn[1], kn[2], kn[2], kn[3], kn[3]]]
        )

        self.outputs2 = nn.ModuleList(
            [nn.Conv3d(c, 1, kernel_size=(1, 1, 1), stride=1, padding=0, bias=False)
             for c in [kn[0], kn[1], kn[1], kn[2], kn[2], kn[3], kn[3]]]
        )
        self.sigmoid = nn.Sigmoid()
```

--------------------------------

### Sliding-Window Inference Implementation

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Defines the SlidingWindowVolume dataset and helper functions for processing 3D TIFF volumes with overlapping windows.

```python
from collections import OrderedDict
from pathlib import Path
import math
import imageio.v2 as imageio
import numpy as np
import tifffile
import torch
from torch.utils.data import Dataset, DataLoader
from tqdm.auto import tqdm


def _triple(value, name):
    if len(value) != 3 or any(int(v) <= 0 for v in value):
        raise ValueError(f"{name} must contain three positive integers (z, y, x): {value}")
    return tuple(int(v) for v in value)


def window_starts(length, window, step):
    if step > window:
        raise ValueError(f"stride {step} cannot exceed crop size {window}")
    if length <= window:
        return [0]
    starts = list(range(0, length - window + 1, step))
    if starts[-1] != length - window:
        starts.append(length - window)
    return starts


def normalize_volume(raw, mode="auto"):
    raw = np.asarray(raw)
    if mode == "none":
        return raw.astype(np.float32, copy=False)
    if mode == "minmax":
        out = raw.astype(np.float32)
        lo, hi = float(np.nanmin(out)), float(np.nanmax(out))
        return (out - lo) / (hi - lo) if hi > lo else np.zeros_like(out)
    if mode != "auto":
        raise ValueError(f"unknown normalization: {mode}")
    if np.issubdtype(raw.dtype, np.integer):
        return raw.astype(np.float32) / float(np.iinfo(raw.dtype).max)
    out = raw.astype(np.float32, copy=False)
    lo, hi = float(np.nanmin(out)), float(np.nanmax(out))
    if 0 <= lo and hi <= 1:
        return out
    return (out - lo) / (hi - lo) if hi > lo else np.zeros_like(out)


class SlidingWindowVolume(Dataset):
    def __init__(self, raw_path, crop_size=(20, 128, 128), stride=(10, 64, 64),
                 normalization="auto", weight_sigma=0.5):
        raw_path = Path(raw_path).expanduser()
        if not raw_path.is_file():
            raise FileNotFoundError(f"Input TIFF not found: {raw_path}")
        raw = np.asarray(tifffile.imread(raw_path))
        if raw.ndim != 3:
            raise ValueError(f"Expected TIFF shape (z,y,x), got {raw.shape}")
        self.original_shape = tuple(raw.shape)
        self.crop_size = _triple(crop_size, "crop_size")
        self.stride = _triple(stride, "stride")
        if self.crop_size[1] % 16 or self.crop_size[2] % 16:
            raise ValueError("crop y/x must be multiples of 16 for MNet")
        pad_after = tuple(max(crop - size, 0) for size, crop in zip(raw.shape, self.crop_size))
        if any(pad_after):
            raw = np.pad(raw, tuple((0, p) for p in pad_after), mode="edge")
        self.raw = normalize_volume(raw, normalization)
        self.padded_shape = tuple(self.raw.shape)
        axes = [window_starts(n, c, s) for n, c, s in
                zip(self.padded_shape, self.crop_size, self.stride)]
        self.positions = [(z, y, x) for z in axes[0] for y in axes[1] for x in axes[2]]
        coords = [np.linspace(-1, 1, n, dtype=np.float32) for n in self.crop_size]
        zz, yy, xx = np.meshgrid(*coords, indexing="ij")
        if weight_sigma > 0:
            weight = np.exp(-(zz*zz + yy*yy + xx*xx) / (2 * weight_sigma**2))
            weight = np.maximum(weight, 1e-3)
        else:
            weight = np.ones(self.crop_size, np.float32)
        self.weight = weight[None].astype(np.float32)

    def __len__(self): return len(self.positions)

    def __getitem__(self, index):
        z, y, x = self.positions[index]
        dz, dy, dx = self.crop_size
        patch = self.raw[z:z+dz, y:y+dy, x:x+dx]
        return np.ascontiguousarray(patch[None], dtype=np.float32), index

    def accumulator(self, channels):
        return (np.zeros((channels,) + self.padded_shape, np.float32),
                np.zeros((1,) + self.padded_shape, np.float32))

    def add(self, total, weights, prediction, index):
        z, y, x = self.positions[int(index)]
        dz, dy, dx = self.crop_size
        if tuple(prediction.shape[1:]) != self.crop_size:
            raise ValueError(f"Model output {prediction.shape[1:]} != crop {self.crop_size}")
        total[:, z:z+dz, y:y+dy, x:x+dx] += prediction * self.weight
        weights[:, z:z+dz, y:y+dy, x:x+dx] += self.weight

    def finish(self, total, weights):
        if np.any(weights == 0):
            raise RuntimeError("Sliding windows left uncovered voxels")
        z, y, x = self.original_shape
        return (total / weights)[:, :z, :y, :x]
```

--------------------------------

### Validate Input Files and Metadata

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Verifies the existence of raw data and checkpoint files, checks the dimensionality of the TIFF volume, and estimates the number of sliding windows required for inference.

```python
import os, tifffile

for name, path in {"raw": RAW_PATH, "checkpoint": CHECKPOINT_PATH}.items():
    if not Path(path).is_file():
        raise FileNotFoundError(f"{name} file not found: {path}")
raw_info = tifffile.imread(RAW_PATH)
if raw_info.ndim != 3:
    raise ValueError(f"RAW must be 3-D, got {raw_info.shape}")
estimated = (len(window_starts(max(raw_info.shape[0], CROP_SIZE[0]), CROP_SIZE[0], STRIDE[0])) *
             len(window_starts(max(raw_info.shape[1], CROP_SIZE[1]), CROP_SIZE[1], STRIDE[1])) *
             len(window_starts(max(raw_info.shape[2], CROP_SIZE[2]), CROP_SIZE[2], STRIDE[2])))
print("RAW:", raw_info.shape, raw_info.dtype,
      "range", (float(np.nanmin(raw_info)), float(np.nanmax(raw_info))))
print("Checkpoint:", round(Path(CHECKPOINT_PATH).stat().st_size / 1024**2, 1), "MiB")
print("Sliding windows:", estimated)
```

--------------------------------

### Visualize Segmentation Results

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Displays the raw image, boundary map, and instance segmentation for a specific z-slice using matplotlib.

```python
import matplotlib.pyplot as plt

z = instances.shape[0] // 2
raw = tifffile.imread(RAW_PATH)
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
axes[0].imshow(raw[z], cmap="gray")
axes[0].set_title(f"Raw z={z}")
axes[1].imshow(boundaries[0, z], cmap="magma", vmin=0, vmax=1)
axes[1].set_title("Boundary")
axes[2].imshow(instances[z], cmap="nipy_spectral")
axes[2].set_title("Instances")
for ax in axes: ax.axis("off")
plt.tight_layout()
```

--------------------------------

### Up Module Implementation

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Defines the upsampling module for the MNet architecture, handling feature fusion and stream management.

```python
class Up(BasicNet):
    def __init__(self, in_channels, out_channels, mode: tuple, FMU='sub'):
        """
        basic module at upsampling stage
        Args:
            in_channels:
            out_channels:
            mode: represent the streams coming in and out. e.g., ('2d', 'both'): one input stream (2d) and two output streams (2d and 3d)
            FMU: determine the type of feature fusion if there are two input streams
        """
        super().__init__()
        self.mode_in, self.mode_out = mode
        self.FMU = FMU
        norm_args = (self.norm_kwargs, self.norm_kwargs)
        activation_args = (self.activation_kwargs, self.activation_kwargs)

        if self.mode_out == '2d' or self.mode_out == 'both':
            self.CB2d = CB3d(in_channels=in_channels, out_channels=out_channels,
                             kSize=((1, 3, 3), (1, 3, 3)), stride=(1, 1), padding=(0, 1, 1),
                             norm_args=norm_args, activation_args=activation_args)

        if self.mode_out == '3d' or self.mode_out == 'both':
            self.CB3d = CB3d(in_channels=in_channels, out_channels=out_channels,
                             kSize=(3, 3), stride=(1, 1), padding=(1, 1, 1),
                             norm_args=norm_args, activation_args=activation_args)

    def forward(self, x):
        x2d, xskip2d, x3d, xskip3d = x

        tarSize = xskip2d.shape[2:]
        up2d = F.interpolate(x2d, size=tarSize, mode='trilinear', align_corners=False)
        up3d = F.interpolate(x3d, size=tarSize, mode='trilinear', align_corners=False)

        cat = torch.cat([FMU(xskip2d, xskip3d, self.FMU), FMU(up2d, up3d, self.FMU)], dim=1)

        if self.mode_out == '2d':
            return self.CB2d(cat)
        elif self.mode_out == '3d':
            return self.CB3d(cat)
        elif self.mode_out == 'both':
            return self.CB2d(cat), self.CB3d(cat)
```

--------------------------------

### MNet Architecture Definition

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Initializes the MNet model with multiple downsampling and upsampling stages.

```python
class MNet(BasicNet):
    def __init__(self, in_channels, kn=(32, 48, 64, 80, 96), FMU='sub'):
        """

        Args:
            in_channels: channels of input
            num_classes: output classes
            kn: the number of kernels
            ds: deep supervision
            FMU: type of feature merging unit
        """
        super().__init__()

        channel_factor = {'sum': 1, 'sub': 1, 'cat': 2}
        fct = channel_factor[FMU]

        self.down11 = Down(in_channels, kn[0], ('/', 'both'), downsample=False)
        self.down12 = Down(kn[0], kn[1], ('2d', 'both'))
        self.down13 = Down(kn[1], kn[2], ('2d', 'both'))
        self.down14 = Down(kn[2], kn[3], ('2d', 'both'))
        self.bottleneck1 = Down(kn[3], kn[4], ('2d', '2d'))
        self.up11 = Up(fct * (kn[3] + kn[4]), kn[3], ('both', '2d'), FMU)
        self.up12 = Up(fct * (kn[2] + kn[3]), kn[2], ('both', '2d'), FMU)
        self.up13 = Up(fct * (kn[1] + kn[2]), kn[1], ('both', '2d'), FMU)
        self.up14 = Up(fct * (kn[0] + kn[1]), kn[0], ('both', 'both'), FMU)

        self.up11_ = Up(fct * (kn[3] + kn[4]), kn[3], ('both', '2d'), FMU)
        self.up12_ = Up(fct * (kn[2] + kn[3]), kn[2], ('both', '2d'), FMU)
        self.up13_ = Up(fct * (kn[1] + kn[2]), kn[1], ('both', '2d'), FMU)
        self.up14_ = Up(fct * (kn[0] + kn[1]), kn[0], ('both', 'both'), FMU)

        self.down21 = Down(kn[0], kn[1], ('3d', 'both'))
        self.down22 = Down(fct * kn[1], kn[2], ('both', 'both'), FMU)
        self.down23 = Down(fct * kn[2], kn[3], ('both', 'both'), FMU)
        self.bottleneck2 = Down(fct * kn[3], kn[4], ('both', 'both'), FMU)
        self.up21 = Up(fct * (kn[3] + kn[4]), kn[3], ('both', 'both'), FMU)
        self.up22 = Up(fct * (kn[2] + kn[3]), kn[2], ('both', 'both'), FMU)
        self.up23 = Up(fct * (kn[1] + kn[2]), kn[1], ('both', '3d'), FMU)

        self.up21_ = Up(fct * (kn[3] + kn[4]), kn[3], ('both', 'both'), FMU)
        self.up22_ = Up(fct * (kn[2] + kn[3]), kn[2], ('both', 'both'), FMU)
        self.up23_ = Up(fct * (kn[1] + kn[2]), kn[1], ('both', '3d'), FMU)

        self.down31 = Down(kn[1], kn[2], ('3d', 'both'))
        self.down32 = Down(fct * kn[2], kn[3], ('both', 'both'), FMU)
        self.bottleneck3 = Down(fct * kn[3], kn[4], ('both', 'both'), FMU)
        self.up31 = Up(fct * (kn[3] + kn[4]), kn[3], ('both', 'both'), FMU)
        self.up32 = Up(fct * (kn[2] + kn[3]), kn[2], ('both', '3d'), FMU)
```

--------------------------------

### Execute Segmentation Inference

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Performs GPU-based inference and FRMC postprocessing to generate affinity, boundary, and instance segmentation maps.

```python
affinities, boundaries, instances = run_segmentation(
    RAW_PATH,
    CHECKPOINT_PATH,
    OUTPUT_DIR,
    crop_size=CROP_SIZE,
    stride=STRIDE,
    batch_size=BATCH_SIZE,
    normalization=NORMALIZATION,
    weight_sigma=WEIGHT_SIGMA,
    beta=BETA,
)
print("Unique instance labels:", len(np.unique(instances)))
```

--------------------------------

### Implement FRMC/Multicut Postprocessing and Segmentation

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

The postprocess function performs watershed segmentation and multicut optimization on affinity arrays. The run_segmentation function orchestrates the full prediction and postprocessing pipeline.

```python
def postprocess(affinities, boundaries=None, output=None, beta=0.25):
    import elf.segmentation.features as feats
    import elf.segmentation.multicut as mc
    import elf.segmentation.watershed as ws

    affinities = np.asarray(affinities, dtype=np.float32)
    if affinities.ndim != 4 or affinities.shape[0] != 3:
        raise ValueError(f"Affinities must be (3,z,y,x), got {affinities.shape}")
    if boundaries is not None:
        boundaries = np.asarray(boundaries, dtype=np.float32)
        if boundaries.ndim == 4 and boundaries.shape[0] == 1:
            boundaries = boundaries[0]
        if boundaries.shape != affinities.shape[1:]:
            raise ValueError(f"Boundary {boundaries.shape} != {affinities.shape[1:]}")
        affinities = np.minimum(affinities, boundaries[None])

    inverted = 1.0 - affinities
    boundary_input = np.maximum(inverted[1], inverted[2])
    watershed = np.zeros_like(boundary_input, dtype=np.uint64)
    offset = 0
    for z in tqdm(range(watershed.shape[0]), desc="Watershed"):
        wsz, max_id = ws.distance_transform_watershed(
            boundary_input[z], threshold=0.25, sigma_seeds=2.0)
        wsz = wsz.astype(np.uint64, copy=False)
        wsz[wsz != 0] += offset
        offset += int(max_id)
        watershed[z] = wsz

    rag = feats.compute_rag(watershed)
    offsets = [[-1,0,0], [0,-1,0], [0,0,-1]]
    probabilities = feats.compute_affinity_features(rag, inverted, offsets)[:, 0]
    edge_sizes = feats.compute_boundary_mean_and_length(rag, boundary_input)[:, 1]
    costs = mc.transform_probabilities_to_costs(
        probabilities, edge_sizes=edge_sizes, beta=float(beta))
    node_labels = mc.multicut_kernighan_lin(rag, costs)
    segmentation = feats.project_node_labels_to_pixels(rag, node_labels)
    if output:
        np.save(output, segmentation)
        print("Instances:", segmentation.shape, segmentation.dtype, "->", output)
    return segmentation


def run_segmentation(raw_path, checkpoint_path, output_dir,
                     crop_size=(20,128,128), stride=(10,64,64), batch_size=1,
                     normalization="auto", weight_sigma=0.5, beta=0.25):
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    affinities, boundaries = predict(
        raw_path, checkpoint_path,
        str(output_dir / "affinities.npy"), str(output_dir / "boundaries.tif"),
        crop_size, stride, batch_size, normalization, weight_sigma)
    instances = postprocess(
        affinities, boundaries, str(output_dir / "instances.npy"), beta)
    return affinities, boundaries, instances
```

--------------------------------

### Define MNet Architecture Components

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Contains the core classes for 3D convolution blocks, feature fusion, and downsampling logic required for the MNet model.

```python
from torch import nn
from torch import nn
import torch
import torch.nn.functional as F


class CNA3d(nn.Module):  # conv + norm + activation
    def __init__(self, in_channels, out_channels, kSize, stride, padding=(1, 1, 1), bias=True, norm_args=None,
                 activation_args=None):
        super().__init__()
        self.norm_args = norm_args
        self.activation_args = activation_args

        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=kSize, stride=stride, padding=padding, bias=bias)

        if norm_args is not None:
            self.norm = nn.InstanceNorm3d(out_channels, **norm_args)

        if activation_args is not None:
            self.activation = nn.LeakyReLU(**activation_args)

    def forward(self, x):
        x = self.conv(x)

        if self.norm_args is not None:
            x = self.norm(x)

        if self.activation_args is not None:
            x = self.activation(x)
        return x


class CB3d(nn.Module):  # conv block 3d
    def __init__(self, in_channels, out_channels, kSize=(3, 3), stride=(1, 1), padding=(1, 1, 1), bias=True,
                 norm_args: tuple = (None, None), activation_args: tuple = (None, None)):
        super().__init__()

        self.conv1 = CNA3d(in_channels, out_channels, kSize=kSize[0], stride=stride[0],
                           padding=padding, bias=bias, norm_args=norm_args[0], activation_args=activation_args[0])

        self.conv2 = CNA3d(out_channels, out_channels, kSize=kSize[1], stride=stride[1],
                           padding=padding, bias=bias, norm_args=norm_args[1], activation_args=activation_args[1])

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        return x


class BasicNet(nn.Module):
    norm_kwargs = {'affine': True}
    activation_kwargs = {'negative_slope': 1e-2, 'inplace': True}

    def __init__(self):
        super(BasicNet, self).__init__()

    def parameter_count(self):
        print("model have {} paramerters in total".format(sum(x.numel() for x in self.parameters()) / 1e6))


def FMU(x1, x2, mode='sub'):
    """
    feature merging unit
    Args:
        x1:
        x2:
        mode: type of fusion
    Returns:
    """
    if mode == 'sum':
        return torch.add(x1, x2)
    elif mode == 'sub':
        return torch.abs(x1 - x2)
    elif mode == 'cat':
        return torch.cat((x1, x2), dim=1)
    else:
        raise Exception('Unexpected mode')


class Down(BasicNet):
    def __init__(self, in_channels, out_channels, mode: tuple, FMU='sub', downsample=True, min_z=8):
        """
        basic module at downsampling stage
        Args:
            in_channels:
            out_channels:
            mode: represent the streams coming in and out. e.g., ('2d', 'both'): one input stream (2d) and two output streams (2d and 3d)
            FMU: determine the type of feature fusion if there are two input streams
            downsample: determine whether to downsample input features (only the first module of MNet do not downsample)
            min_z: if the size of z-axis < min_z, maxpooling won't be applied along z-axis
        """
        super().__init__()
        self.mode_in, self.mode_out = mode
        self.downsample = downsample
        self.FMU = FMU
        self.min_z = min_z
        norm_args = (self.norm_kwargs, self.norm_kwargs)
        activation_args = (self.activation_kwargs, self.activation_kwargs)

        if self.mode_out == '2d' or self.mode_out == 'both':
            self.CB2d = CB3d(in_channels=in_channels, out_channels=out_channels,
                             kSize=((1, 3, 3), (1, 3, 3)), stride=(1, 1), padding=(0, 1, 1),
                             norm_args=norm_args, activation_args=activation_args)

        if self.mode_out == '3d' or self.mode_out == 'both':
            self.CB3d = CB3d(in_channels=in_channels, out_channels=out_channels,
                             kSize=(3, 3), stride=(1, 1), padding=(1, 1, 1),
                             norm_args=norm_args, activation_args=activation_args)

    def forward(self, x):
        if self.downsample:
            if self.mode_in == 'both':
                x2d, x3d = x
                p2d = F.max_pool3d(x2d, kernel_size=(1, 2, 2), stride=(1, 2, 2))
                if x3d.shape[2] >= self.min_z:
                    p3d = F.max_pool3d(x3d, kernel_size=(2, 2, 2), stride=(2, 2, 2))
                else:
                    p3d = F.max_pool3d(x3d, kernel_size=(1, 2, 2), stride=(1, 2, 2))

                x = FMU(p2d, p3d, mode=self.FMU)

            elif self.mode_in == '2d':
                x = F.max_pool3d(x, kernel_size=(1, 2, 2), stride=(1, 2, 2))
```

--------------------------------

### Implement MNet Forward Pass

Source: https://github.com/yanchaoz/segneuron/blob/main/SegNeuron_Colab_Inference.ipynb

Executes the forward pass of the MNet model, managing feature propagation through downsampling, bottleneck, and upsampling stages.

```python
    def forward(self, x):
        down11 = self.down11(x)
        down12 = self.down12(down11[0])
        down13 = self.down13(down12[0])
        down14 = self.down14(down13[0])
        bottleNeck1 = self.bottleneck1(down14[0])

        down21 = self.down21(down11[1])
        down22 = self.down22([down21[0], down12[1]])
        down23 = self.down23([down22[0], down13[1]])
        bottleNeck2 = self.bottleneck2([down23[0], down14[1]])

        down31 = self.down31(down21[1])
        down32 = self.down32([down31[0], down22[1]])
        bottleNeck3 = self.bottleneck3([down32[0], down23[1]])

        down41 = self.down41(down31[1])
        bottleNeck4 = self.bottleneck4([down41[0], down32[1]])

        bottleNeck5 = self.bottleneck5(down41[1])

        up41 = self.up41([bottleNeck4[0], down41[0], bottleNeck5, down41[1]])

        up31 = self.up31([bottleNeck3[0], down32[0], bottleNeck4[1], down32[1]])
        up32 = self.up32([up31[0], down31[0], up41, down31[1]])

        up21 = self.up21([bottleNeck2[0], down23[0], bottleNeck3[1], down23[1]])
        up22 = self.up22([up21[0], down22[0], up31[1], down22[1]])
        up23 = self.up23([up22[0], down21[0], up32, down21[1]])

        up11 = self.up11([bottleNeck1, down14[0], bottleNeck2[1], down14[1]])
        up12 = self.up12([up11, down13[0], up21[1], down13[1]])
        up13 = self.up13([up12, down12[0], up22[1], down12[1]])
        up14 = self.up14([up13, down11[0], up23, down11[1]])

        up41_ = self.up41_([bottleNeck4[0], down41[0], bottleNeck5, down41[1]])

        up31_ = self.up31_([bottleNeck3[0], down32[0], bottleNeck4[1], down32[1]])
        up32_ = self.up32_([up31_[0], down31[0], up41_, down31[1]])

        up21_ = self.up21_([bottleNeck2[0], down23[0], bottleNeck3[1], down23[1]])
        up22_ = self.up22_([up21_[0], down22[0], up31_[1], down22[1]])
        up23_ = self.up23_([up22_[0], down21[0], up32_, down21[1]])

        up11_ = self.up11_([bottleNeck1, down14[0], bottleNeck2[1], down14[1]])
        up12_ = self.up12_([up11_, down13[0], up21_[1], down13[1]])
        up13_ = self.up13_([up12_, down12[0], up22_[1], down12[1]])
        up14_ = self.up14_([up13_, down11[0], up23_, down11[1]])

        return self.sigmoid(self.outputs[0](up14[0] + up14[1])), self.sigmoid(self.outputs2[0](up14_[0] + up14_[1]))
```
Tokens:
Raw