update fiber.yaml

Author: Donglai Wei

connectomics/data/io/__init__.py

@@ -14,17 +14,23 @@
 # Core I/O functions
 from .io import (
     # HDF5 I/O
-    read_hdf5, write_hdf5, list_hdf5_datasets,
-
+    read_hdf5,
+    write_hdf5,
+    list_hdf5_datasets,
     # Image I/O
-    read_image, read_images, read_image_as_volume,
-    save_image, save_images, SUPPORTED_IMAGE_FORMATS,
-
+    read_image,
+    read_images,
+    read_image_as_volume,
+    save_image,
+    save_images,
+    SUPPORTED_IMAGE_FORMATS,
     # Pickle I/O
-    read_pickle_file, write_pickle_file,
-
+    read_pickle_file,
+    write_pickle_file,
     # High-level volume I/O
-    read_volume, save_volume, get_vol_shape,
+    read_volume,
+    save_volume,
+    get_vol_shape,
 )
 
 # Tile operations
@@ -35,7 +41,8 @@
 
 # Utilities
 from .utils import (
-    vast_to_segmentation,
+    seg_to_rgb,
+    rgb_to_seg,
     normalize_data_range,
     convert_to_uint8,
     split_multichannel_mask,
@@ -51,25 +58,35 @@
 
 __all__ = [
     # HDF5 I/O
-    'read_hdf5', 'write_hdf5', 'list_hdf5_datasets',
-
+    "read_hdf5",
+    "write_hdf5",
+    "list_hdf5_datasets",
     # Image I/O
-    'read_image', 'read_images', 'read_image_as_volume',
-    'save_image', 'save_images', 'SUPPORTED_IMAGE_FORMATS',
-
+    "read_image",
+    "read_images",
+    "read_image_as_volume",
+    "save_image",
+    "save_images",
+    "SUPPORTED_IMAGE_FORMATS",
     # Pickle I/O
-    'read_pickle_file', 'write_pickle_file',
-
+    "read_pickle_file",
+    "write_pickle_file",
     # High-level volume I/O
-    'read_volume', 'save_volume', 'get_vol_shape',
-
+    "read_volume",
+    "save_volume",
+    "get_vol_shape",
     # Tile operations
-    'create_tile_metadata', 'reconstruct_volume_from_tiles',
-
+    "create_tile_metadata",
+    "reconstruct_volume_from_tiles",
     # Utilities
-    'vast_to_segmentation', 'normalize_data_range', 'convert_to_uint8',
-    'split_multichannel_mask', 'squeeze_arrays',
-
+    "seg_to_rgb",
+    "rgb_to_seg",
+    "normalize_data_range",
+    "convert_to_uint8",
+    "split_multichannel_mask",
+    "squeeze_arrays",
     # MONAI transforms
-    'LoadVolumed', 'SaveVolumed', 'TileLoaderd',
+    "LoadVolumed",
+    "SaveVolumed",
+    "TileLoaderd",
 ]

connectomics/data/io/tiles.py

@@ -12,7 +12,7 @@
 from scipy.ndimage import zoom
 
 from .io import read_image
-from .utils import vast_to_segmentation
+from .utils import rgb_to_seg
 
 
 def create_tile_metadata(
@@ -170,8 +170,8 @@ def reconstruct_volume_from_tiles(
                         result[z - z0,
                                y_actual_start - y0:y_actual_end - y0,
                                x_actual_start - x0:x_actual_end - x0] = \
-                            vast_to_segmentation(patch[y_actual_start - y_patch_start:y_actual_end - y_patch_start,
-                                                        x_actual_start - x_patch_start:x_actual_end - x_patch_start])
+                            rgb_to_seg(patch[y_actual_start - y_patch_start:y_actual_end - y_patch_start,
+                                             x_actual_start - x_patch_start:x_actual_end - x_patch_start])
 
     # Apply padding for chunks touching the border of the large input volume
     if max(boundary_diffs) > 0:

connectomics/data/io/utils.py

@@ -9,7 +9,24 @@
 import numpy as np
 
 
-def vast_to_segmentation(segmentation_data: np.ndarray) -> np.ndarray:
+def seg_to_rgb(seg):
+    """
+    Convert a segmentation map to an RGB image.
+
+    Args:
+        seg (numpy.ndarray): The input segmentation map.
+
+    Returns:
+        numpy.ndarray: The RGB image representation of the segmentation map.
+
+    Notes:
+        - The function converts a segmentation map to an RGB image, where each unique segment ID is assigned a unique color.
+        - The RGB image is represented as a numpy array.
+    """
+    return np.stack([seg // 65536, seg // 256, seg % 256], axis=2).astype(np.uint8)
+
+
+def rgb_to_seg(rgb: np.ndarray) -> np.ndarray:
     """Convert VAST segmentation format to standard format.
 
     VAST format uses RGB encoding where each pixel's RGB values are combined
@@ -22,20 +39,25 @@ def vast_to_segmentation(segmentation_data: np.ndarray) -> np.ndarray:
         Converted segmentation with proper ID encoding
     """
     # Convert to 24 bits
-    if segmentation_data.ndim == 2 or segmentation_data.shape[-1] == 1:
-        return np.squeeze(segmentation_data)
-    elif segmentation_data.ndim == 3:  # Single RGB image
-        return (segmentation_data[:, :, 0].astype(np.uint32) * 65536 +
-                segmentation_data[:, :, 1].astype(np.uint32) * 256 +
-                segmentation_data[:, :, 2].astype(np.uint32))
-    elif segmentation_data.ndim == 4:  # Multiple RGB images
-        return (segmentation_data[:, :, :, 0].astype(np.uint32) * 65536 +
-                segmentation_data[:, :, :, 1].astype(np.uint32) * 256 +
-                segmentation_data[:, :, :, 2].astype(np.uint32))
-
-
-def normalize_data_range(data: np.ndarray, target_min: float = 0.0,
-                        target_max: float = 1.0, ignore_uint8: bool = True) -> np.ndarray:
+    if rgb.ndim == 2 or rgb.shape[-1] == 1:
+        return np.squeeze(rgb)
+    elif rgb.ndim == 3:  # Single RGB image
+        return (
+            rgb[:, :, 0].astype(np.uint32) * 65536
+            + rgb[:, :, 1].astype(np.uint32) * 256
+            + rgb[:, :, 2].astype(np.uint32)
+        )
+    elif rgb.ndim == 4:  # Multiple RGB images
+        return (
+            rgb[:, :, :, 0].astype(np.uint32) * 65536
+            + rgb[:, :, :, 1].astype(np.uint32) * 256
+            + rgb[:, :, :, 2].astype(np.uint32)
+        )
+
+
+def normalize_data_range(
+    data: np.ndarray, target_min: float = 0.0, target_max: float = 1.0, ignore_uint8: bool = True
+) -> np.ndarray:
     """Normalize array values to a target range.
 
     Args:
@@ -119,9 +141,10 @@ def squeeze_arrays(*arrays):
 
 
 __all__ = [
-    'vast_to_segmentation',
-    'normalize_data_range',
-    'convert_to_uint8',
-    'split_multichannel_mask',
-    'squeeze_arrays',
-]
+    "seg_to_rgb",
+    "rgb_to_seg",
+    "normalize_data_range",
+    "convert_to_uint8",
+    "split_multichannel_mask",
+    "squeeze_arrays",
+]

connectomics/data/process/monai_transforms.py

@@ -602,7 +602,7 @@ class MultiTaskLabelTransformd(MapTransform):
     }
     _TASK_DEFAULTS: Dict[str, Dict[str, Any]] = {
         "binary": {},
-        "affinity": {"offsets": ["1-1-0", "1-0-0", "0-1-0", "0-0-1"]},
+        "affinity": {"offsets": ["1-0-0", "0-1-0", "0-0-1"]},  # Default: 3 short-range affinities (z, y, x)
         "instance_boundary": {"thickness": 1, "edge_mode": "seg-all", "mode": "3d"},
         "instance_edt": {"mode": "2d", "quantize": False},
         "skeleton_aware_edt": {

connectomics/data/process/target.py

@@ -244,74 +244,111 @@ def seg_to_binary(label, segment_id=[]):
     return fg_mask
 
 
-def seg_to_affinity(seg: np.ndarray, target_opt: List[str]) -> np.ndarray:
+def seg_to_affinity(
+    seg: np.ndarray,
+    offsets: List[str] = None,
+    long_range: int = None,
+) -> np.ndarray:
     """
-    Compute affinities from a segmentation based on target options.
+    Compute affinity maps from segmentation.
+
+    Supports two modes:
+    1. DeepEM/SNEMI style: Provide `offsets` as list of strings (e.g., ["0-0-1", "0-1-0", "1-0-0"])
+    2. BANIS style: Provide `long_range` as int for 6-channel output (3 short + 3 long range)
 
     Args:
         seg: The segmentation to compute affinities from. Shape: (z, y, x).
-        target_opt: List of strings defining affinity offsets.
-            Can be either:
-            - Legacy format: ['1', '0-0-1', '0-1-0', ...] (first element is type indicator)
-            - Modern format: ['0-0-1', '0-1-0', ...] (direct offset list)
+             0 indicates background.
+        offsets: List of offset strings in "z-y-x" format (e.g., ["0-0-1", "0-1-0", "1-0-0"]).
+                 Each string defines one affinity channel.
+        long_range: BANIS-style: offset for long-range affinities. Produces 6 channels:
+                    - Channel 0-2: Short-range (offset 1) for z, y, x
+                    - Channel 3-5: Long-range (offset long_range) for z, y, x
 
     Returns:
-        The affinities. Shape: (num_offsets, z, y, x).
+        The affinities. Shape: (num_channels, z, y, x).
     """
-    if len(target_opt) == 0:
-        # Default short-range affinities
-        offsets = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
-    else:
-        # Detect format: check if first element is a type indicator or an offset
-        start_idx = 0
-        if len(target_opt) > 0 and "-" not in target_opt[0]:
-            # Legacy format: first element is type indicator (e.g., '1')
-            start_idx = 1
-
-        # Parse offsets from target_opt
-        offsets = []
-        for opt_str in target_opt[start_idx:]:
-            if "-" in opt_str:
-                offset = [int(x) for x in opt_str.split("-")]
-                offsets.append(offset)
-
-        # Fallback to default if no valid offsets found
-        if len(offsets) == 0:
-            offsets = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
-
-    num_offsets = len(offsets)
-    affinities = np.zeros((num_offsets, *seg.shape), dtype=np.float32)
-
-    for i, offset in enumerate(offsets):
-        dz, dy, dx = offset
-
-        # Create slices for the offset
+    # BANIS mode: use long_range parameter (takes precedence if specified)
+    if long_range is not None:
+        affinities = np.zeros((6, *seg.shape), dtype=np.float32)
+
+        # Short range affinities (offset 1)
+        affinities[0, :-1] = (seg[:-1] == seg[1:]) & (seg[1:] > 0)
+        affinities[1, :, :-1] = (seg[:, :-1] == seg[:, 1:]) & (seg[:, 1:] > 0)
+        affinities[2, :, :, :-1] = (seg[:, :, :-1] == seg[:, :, 1:]) & (seg[:, :, 1:] > 0)
+
+        # Long range affinities
+        affinities[3, :-long_range] = (seg[:-long_range] == seg[long_range:]) & (seg[long_range:] > 0)
+        affinities[4, :, :-long_range] = (seg[:, :-long_range] == seg[:, long_range:]) & (seg[:, long_range:] > 0)
+        affinities[5, :, :, :-long_range] = (seg[:, :, :-long_range] == seg[:, :, long_range:]) & (seg[:, :, long_range:] > 0)
+
+        return affinities
+
+    # DeepEM/SNEMI mode: use offsets parameter
+    if offsets is None:
+        # Default: short-range affinities for z, y, x
+        offsets = ["1-0-0", "0-1-0", "0-0-1"]
+
+    # Parse offsets from strings
+    parsed_offsets = []
+    for offset_str in offsets:
+        parts = offset_str.split("-")
+        if len(parts) == 3:
+            parsed_offsets.append([int(parts[0]), int(parts[1]), int(parts[2])])
+        else:
+            raise ValueError(f"Invalid offset format: {offset_str}. Expected 'z-y-x' format.")
+
+    num_channels = len(parsed_offsets)
+    affinities = np.zeros((num_channels, *seg.shape), dtype=np.float32)
+
+    for i, (dz, dy, dx) in enumerate(parsed_offsets):
+        # Handle each axis independently
+        # For positive offset: compare seg[:-offset] with seg[offset:]
+        # For negative offset: compare seg[-offset:] with seg[:offset]
+
+        if dz == 0 and dy == 0 and dx == 0:
+            # Zero offset: all foreground pixels are 1
+            affinities[i] = (seg > 0).astype(np.float32)
+            continue
+
+        # Build source and destination slices for each axis
         if dz > 0:
-            src_slice = (slice(None, -dz), slice(None), slice(None))
-            dst_slice = (slice(dz, None), slice(None), slice(None))
+            z_src = slice(None, -dz)
+            z_dst = slice(dz, None)
         elif dz < 0:
-            src_slice = (slice(-dz, None), slice(None), slice(None))
-            dst_slice = (slice(None, dz), slice(None), slice(None))
+            z_src = slice(-dz, None)
+            z_dst = slice(None, dz)
         else:
-            src_slice = (slice(None), slice(None), slice(None))
-            dst_slice = (slice(None), slice(None), slice(None))
+            z_src = slice(None)
+            z_dst = slice(None)
 
         if dy > 0:
-            src_slice = (src_slice[0], slice(None, -dy), src_slice[2])
-            dst_slice = (dst_slice[0], slice(dy, None), dst_slice[2])
+            y_src = slice(None, -dy)
+            y_dst = slice(dy, None)
         elif dy < 0:
-            src_slice = (src_slice[0], slice(-dy, None), src_slice[2])
-            dst_slice = (dst_slice[0], slice(None, dy), dst_slice[2])
+            y_src = slice(-dy, None)
+            y_dst = slice(None, dy)
+        else:
+            y_src = slice(None)
+            y_dst = slice(None)
 
         if dx > 0:
-            src_slice = (src_slice[0], src_slice[1], slice(None, -dx))
-            dst_slice = (dst_slice[0], dst_slice[1], slice(dx, None))
+            x_src = slice(None, -dx)
+            x_dst = slice(dx, None)
         elif dx < 0:
-            src_slice = (src_slice[0], src_slice[1], slice(-dx, None))
-            dst_slice = (dst_slice[0], dst_slice[1], slice(None, dx))
+            x_src = slice(-dx, None)
+            x_dst = slice(None, dx)
+        else:
+            x_src = slice(None)
+            x_dst = slice(None)
+
+        src_slice = (z_src, y_src, x_src)
+        dst_slice = (z_dst, y_dst, x_dst)
 
-        # Compute affinity
-        affinities[i][dst_slice] = (seg[src_slice] == seg[dst_slice]) & (seg[dst_slice] > 0)
+        # Compute affinity: same segment ID and not background
+        affinities[i][dst_slice] = (
+            (seg[src_slice] == seg[dst_slice]) & (seg[dst_slice] > 0)
+        ).astype(np.float32)
 
     return affinities