Diagnose stalled mito_mitoEM_H train

Author: Donglai Wei

connectomics/decoding/__init__.py

@@ -1,6 +1,7 @@
 """Decoding package for PyTorch Connectomics."""
 
 from .base import DecodeStep
+from .abiss import decode_abiss
 from .pipeline import (
     apply_decode_mode,
     apply_decode_pipeline,
@@ -67,6 +68,7 @@
     "decode_instance_binary_contour_distance",
     "decode_affinity_cc",
     "decode_distance_watershed",
+    "decode_abiss",
     # Auto-tuning
     "optimize_threshold",
     "optimize_parameters",

connectomics/decoding/abiss.py

@@ -0,0 +1,156 @@
+"""ABISS external wrapper decoder.
+
+This module exposes a decoder that bridges prediction tensors to an external
+ABISS (or ABISS-compatible) command-line pipeline.
+
+The wrapper writes predictions to temporary files, runs a user-specified
+command, then reads back an instance-label segmentation.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any, Dict, List, Mapping, Optional, Sequence
+import os
+import subprocess
+
+import numpy as np
+
+from connectomics.data.io import read_hdf5, write_hdf5
+
+from .utils import cast2dtype
+
+__all__ = ["decode_abiss"]
+
+
+def _format_command(
+    command: str | Sequence[str],
+    mapping: Mapping[str, str],
+) -> tuple[str | List[str], bool]:
+    """Format command placeholders for shell/list execution."""
+    if isinstance(command, str):
+        return command.format(**mapping), True
+    if isinstance(command, Sequence):
+        return [str(part).format(**mapping) for part in command], False
+    raise TypeError(f"`command` must be str or sequence[str], got {type(command).__name__}.")
+
+
+def _load_output(output_h5: Path, output_npy: Path, output_dataset: str) -> np.ndarray:
+    """Load decoded segmentation output from file."""
+    if output_h5.exists():
+        seg = read_hdf5(str(output_h5), dataset=output_dataset)
+    elif output_npy.exists():
+        seg = np.load(output_npy)
+    else:
+        raise FileNotFoundError(
+            "decode_abiss did not produce output file. "
+            f"Expected one of: {output_h5}, {output_npy}"
+        )
+
+    seg = np.asarray(seg)
+    if seg.ndim == 4 and seg.shape[0] == 1:
+        seg = seg[0]
+    if seg.ndim != 3:
+        raise ValueError(
+            "decode_abiss output must be 3D label volume (Z, Y, X) "
+            f"or singleton-channel 4D; got shape {seg.shape}."
+        )
+
+    if not np.issubdtype(seg.dtype, np.integer):
+        seg = np.rint(seg).astype(np.uint64, copy=False)
+
+    return cast2dtype(seg)
+
+
+def decode_abiss(
+    predictions: np.ndarray,
+    command: str | Sequence[str],
+    *,
+    input_dataset: str = "main",
+    output_dataset: str = "main",
+    channels: Optional[Sequence[int]] = None,
+    workdir: Optional[str] = None,
+    keep_workspace: bool = False,
+    timeout_sec: Optional[int] = None,
+    env: Optional[Dict[str, Any]] = None,
+    check: bool = True,
+) -> np.ndarray:
+    """Decode instance segmentation with an external ABISS command.
+
+    Args:
+        predictions: Model output, typically shape ``(C, Z, Y, X)``.
+        command: External command to execute. Supports placeholders:
+            - ``{workspace}``: working directory path
+            - ``{input_h5}``, ``{input_npy}``: prediction file paths
+            - ``{output_h5}``, ``{output_npy}``: expected output file paths
+            - ``{input_dataset}``, ``{output_dataset}``: HDF5 dataset names
+        input_dataset: Dataset name when writing input HDF5.
+        output_dataset: Dataset name when reading output HDF5.
+        channels: Optional channel indices to select before saving input.
+        workdir: Optional fixed workspace directory. If None, uses temp dir.
+        keep_workspace: Keep temp workspace when using auto temp dir.
+        timeout_sec: Optional subprocess timeout in seconds.
+        env: Optional extra environment variables for subprocess.
+        check: Raise on non-zero return code if True.
+
+    Returns:
+        3D instance label volume ``(Z, Y, X)``.
+    """
+    pred = np.asarray(predictions)
+    if pred.ndim not in (3, 4):
+        raise ValueError(
+            f"decode_abiss expects 3D/4D predictions, got shape {pred.shape}."
+        )
+
+    if channels is not None:
+        if pred.ndim != 4:
+            raise ValueError("`channels` can only be used for 4D predictions (C, Z, Y, X).")
+        pred = pred[np.asarray(channels)]
+
+    if workdir is not None:
+        workspace_path = Path(workdir).resolve()
+        workspace_path.mkdir(parents=True, exist_ok=True)
+        temp_ctx = None
+    else:
+        temp_ctx = TemporaryDirectory(prefix="decode_abiss_")
+        workspace_path = Path(temp_ctx.name).resolve()
+
+    try:
+        input_h5 = workspace_path / "predictions.h5"
+        input_npy = workspace_path / "predictions.npy"
+        output_h5 = workspace_path / "segmentation.h5"
+        output_npy = workspace_path / "segmentation.npy"
+
+        # Save both formats so external command can choose the easiest input.
+        write_hdf5(str(input_h5), pred, dataset=input_dataset)
+        np.save(input_npy, pred)
+
+        mapping = {
+            "workspace": str(workspace_path),
+            "input_h5": str(input_h5),
+            "input_npy": str(input_npy),
+            "output_h5": str(output_h5),
+            "output_npy": str(output_npy),
+            "input_dataset": input_dataset,
+            "output_dataset": output_dataset,
+        }
+        cmd, use_shell = _format_command(command, mapping)
+
+        proc_env = os.environ.copy()
+        if env:
+            proc_env.update({str(k): str(v) for k, v in env.items()})
+
+        subprocess.run(
+            cmd,
+            shell=use_shell,
+            env=proc_env,
+            cwd=str(workspace_path),
+            check=check,
+            timeout=timeout_sec,
+        )
+
+        return _load_output(output_h5=output_h5, output_npy=output_npy, output_dataset=output_dataset)
+    finally:
+        if temp_ctx is not None and not keep_workspace:
+            temp_ctx.cleanup()

connectomics/decoding/registry.py

@@ -62,6 +62,7 @@ def list_decoders() -> List[str]:
 
 def register_builtin_decoders() -> None:
     """Populate registry with built-in decoders."""
+    from .abiss import decode_abiss
     from .segmentation import (
         decode_affinity_cc,
         decode_distance_watershed,
@@ -76,5 +77,5 @@ def register_builtin_decoders() -> None:
     )
     register_decoder("decode_affinity_cc", decode_affinity_cc, overwrite=True)
     register_decoder("decode_distance_watershed", decode_distance_watershed, overwrite=True)
+    register_decoder("decode_abiss", decode_abiss, overwrite=True)
     register_decoder("polarity2instance", polarity2instance, overwrite=True)
-

tests/unit/test_decode_abiss_wrapper.py

@@ -0,0 +1,62 @@
+"""Tests for decode_abiss external wrapper."""
+
+from __future__ import annotations
+
+import sys
+
+import numpy as np
+import pytest
+
+from connectomics.decoding import decode_abiss
+
+
+def test_decode_abiss_with_list_command_writes_npy_output():
+    pred = np.zeros((3, 6, 8, 10), dtype=np.float32)
+    pred[0, 1:4, 2:6, 3:8] = 0.9
+
+    command = [
+        sys.executable,
+        "-c",
+        (
+            "import h5py, numpy as np; "
+            "x = h5py.File('{input_h5}', 'r')['{input_dataset}'][:]; "
+            "y = (x[0] > 0.5).astype(np.uint64); "
+            "np.save('{output_npy}', y)"
+        ),
+    ]
+
+    seg = decode_abiss(pred, command=command)
+    assert seg.shape == (6, 8, 10)
+    assert np.issubdtype(seg.dtype, np.integer)
+    assert seg.max() == 1
+    assert seg[2, 3, 4] == 1
+
+
+def test_decode_abiss_with_string_command_writes_h5_output():
+    pred = np.zeros((3, 5, 7, 9), dtype=np.float32)
+    pred[1, 1:4, 2:5, 3:7] = 1.0
+
+    command = (
+        f"{sys.executable} -c \""
+        "import h5py, numpy as np; "
+        "x = h5py.File('{input_h5}', 'r')['{input_dataset}'][:]; "
+        "y = (x[1] > 0.5).astype(np.uint64); "
+        "f = h5py.File('{output_h5}', 'w'); "
+        "f.create_dataset('{output_dataset}', data=y); "
+        "f.close()\""
+    )
+
+    seg = decode_abiss(pred, command=command)
+    assert seg.shape == (5, 7, 9)
+    assert seg.max() == 1
+    assert seg[2, 3, 4] == 1
+
+
+def test_decode_abiss_raises_if_output_missing():
+    pred = np.zeros((3, 4, 4, 4), dtype=np.float32)
+
+    command = [sys.executable, "-c", "print('no output written')"]
+
+    with pytest.raises(FileNotFoundError, match="did not produce output file"):
+        decode_abiss(pred, command=command)
+

tests/unit/test_decoding_pipeline.py

@@ -38,6 +38,7 @@ def test_builtin_decoders_are_registered():
     assert "decode_affinity_cc" in names
     assert "decode_distance_watershed" in names
     assert "decode_instance_binary_contour_distance" in names
+    assert "decode_abiss" in names
 
 
 def test_decode_pipeline_dict_mode_matches_direct_decoder(affinity_with_redundant_channels):
@@ -62,4 +63,3 @@ def test_decode_pipeline_unknown_decoder_raises(affinity_with_redundant_channels
     decode_modes = [{"name": "decode_not_exists", "kwargs": {}}]
     with pytest.raises(ValueError, match="Unknown decode function"):
         apply_decode_pipeline(affinity_with_redundant_channels, decode_modes)
-

tutorials/neuron_snemi.yaml

@@ -195,6 +195,20 @@ test:
   - name: decode_affinity_cc
     kwargs:
       threshold: 0.5
+  # Example: external ABISS wrapper (replace decode_affinity_cc above)
+  # - name: decode_abiss
+  #   kwargs:
+  #     command:
+  #       - python
+  #       - scripts/run_abiss_single.py
+  #       - --input
+  #       - "{input_h5}"
+  #       - --input-dataset
+  #       - "{input_dataset}"
+  #       - --output
+  #       - "{output_h5}"
+  #       - --output-dataset
+  #       - "{output_dataset}"
   evaluation:
     enabled: true
     metrics:
@@ -228,6 +242,10 @@ inference:
   - name: decode_affinity_cc
     kwargs:
       threshold: 0.5
+  # Example (same as test.decoding):
+  # - name: decode_abiss
+  #   kwargs:
+  #     command: "python scripts/run_abiss_single.py --input {input_h5} --output {output_h5}"
   evaluation:
     enabled: false
     metrics: