Zarr I/O, val skeleton precompute, data-check visualization, lazy zarr label_aux

- read_volume/save_volume: add zarr format support (split store/subkey paths)
- LazyZarrVolumeDataset: add label_aux_paths, load non-zarr files eagerly
- data_factory: precompute val skeleton alongside train, val uses same dataset class
- VisualizationCallback: log image+label on first batch (data_check, no prediction)
- VisualizationCallback: head="all" visualizes all heads separately
- config_io: allow head="all" in validation

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: Donglai Wei

connectomics/data/io/__init__.py

@@ -14,6 +14,7 @@
     read_images,
     read_volume,
     save_volume,
+    volume_exists,
     write_hdf5,
 )
 from .transforms import (
@@ -32,6 +33,7 @@
     "read_volume",
     "save_volume",
     "get_vol_shape",
+    "volume_exists",
     "LoadVolumed",
     "SaveVolumed",
     "TileLoaderd",

connectomics/data/io/io.py

@@ -13,6 +13,7 @@
 import glob
 import logging
 import os
+from pathlib import Path
 from typing import List, Optional, Union
 
 import h5py
@@ -35,11 +36,9 @@ def _detect_format(filename: str) -> str:
     Returns canonical format string:
     'h5', 'tiff', 'png', 'nifti', 'zarr'.
     """
-    if ".zarr" in filename:
-        return "zarr"
     if filename.endswith(".nii.gz"):
         return "nifti"
-    suffix = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
+    suffix = Path(filename).suffix.lower().lstrip(".")
     _SUFFIX_MAP = {
         "h5": "h5",
         "hdf5": "h5",
@@ -49,12 +48,22 @@ def _detect_format(filename: str) -> str:
         "nii": "nifti",
     }
     fmt = _SUFFIX_MAP.get(suffix)
-    if fmt is None:
-        raise ValueError(
-            f"Unrecognizable file format for {filename}. "
-            f"Expected: h5, hdf5, tif, tiff, png, nii, nii.gz"
-        )
-    return fmt
+    if fmt is not None:
+        return fmt
+    if ".zarr" in filename:
+        return "zarr"
+    raise ValueError(
+        f"Unrecognizable file format for {filename}. "
+        f"Expected: h5, hdf5, tif, tiff, png, nii, nii.gz, zarr"
+    )
+
+
+def _split_zarr_path(filename: str) -> tuple[str, Optional[str]]:
+    """Split a zarr path into store path and optional subkey."""
+    zarr_idx = filename.index(".zarr")
+    zarr_path = filename[: zarr_idx + 5]
+    sub_key = filename[zarr_idx + 5 :].strip("/") or None
+    return zarr_path, sub_key
 
 
 # =============================================================================
@@ -346,10 +355,7 @@ def read_volume(
     elif fmt == "zarr":
         import zarr
 
-        # Path may be "dir.zarr/subkey" — split at .zarr boundary.
-        zarr_idx = filename.index(".zarr")
-        zarr_path = filename[: zarr_idx + 5]
-        sub_key = filename[zarr_idx + 5 :].strip("/") or None
+        zarr_path, sub_key = _split_zarr_path(filename)
         store = zarr.open(zarr_path, mode="r")
         arr = store[sub_key] if sub_key else store
         data = np.asarray(arr)
@@ -374,7 +380,7 @@ def save_volume(
     filename: str,
     volume: np.ndarray,
     dataset: str = "main",
-    file_format: str = "h5",
+    file_format: Optional[str] = None,
 ) -> None:
     """Save volumetric data in specified format.
 
@@ -384,9 +390,27 @@ def save_volume(
         dataset: Dataset name for HDF5 format.
         file_format: 'h5', 'tiff', 'png', 'nii', 'nii.gz'.
     """
+    file_format = file_format or _detect_format(filename)
+
     if file_format == "h5":
         write_hdf5(filename, volume, dataset=dataset)
 
+    elif file_format == "zarr":
+        import zarr
+
+        zarr_path, sub_key = _split_zarr_path(filename)
+        if sub_key:
+            group = zarr.open_group(zarr_path, mode="a")
+            group.create_dataset(sub_key, data=volume, overwrite=True)
+        else:
+            array = zarr.open(
+                zarr_path,
+                mode="w",
+                shape=volume.shape,
+                dtype=volume.dtype,
+            )
+            array[...] = volume
+
     elif file_format in ("tif", "tiff"):
         import tifffile
 
@@ -410,7 +434,7 @@ def save_volume(
 
     else:
         raise ValueError(
-            f"Unsupported format: {file_format}. " f"Expected: h5, tiff, png, nii, nii.gz"
+            f"Unsupported format: {file_format}. " f"Expected: h5, zarr, tiff, png, nii, nii.gz"
         )
 
 
@@ -436,17 +460,19 @@ def get_vol_shape(
     Returns shape consistent with what read_volume would
     produce: (D, H, W) or (C, D, H, W).
     """
-    if not os.path.exists(filename):
-        raise FileNotFoundError(f"File not found: {filename}")
-
     fmt = _detect_format(filename)
 
     if fmt == "zarr":
         try:
             import zarr
         except ModuleNotFoundError as exc:
             raise ModuleNotFoundError("zarr required. pip install zarr") from exc
-        obj = zarr.open(filename, mode="r")
+        zarr_path, sub_key = _split_zarr_path(filename)
+        if not os.path.exists(zarr_path):
+            raise FileNotFoundError(f"File not found: {zarr_path}")
+        obj = zarr.open(zarr_path, mode="r")
+        if sub_key:
+            return tuple(obj[sub_key].shape)
         if hasattr(obj, "shape"):
             return tuple(obj.shape)
         if dataset is not None:
@@ -456,6 +482,9 @@ def get_vol_shape(
             raise ValueError(f"No arrays in zarr group: {filename}")
         return tuple(obj[keys[0]].shape)
 
+    if not os.path.exists(filename):
+        raise FileNotFoundError(f"File not found: {filename}")
+
     if fmt == "h5":
         with h5py.File(filename, "r") as f:
             if dataset is None:
@@ -483,3 +512,15 @@ def get_vol_shape(
         return _get_nifti_shape(filename)
 
     raise ValueError(f"Unsupported format: {fmt}")
+
+
+def volume_exists(
+    filename: str,
+    dataset: Optional[str] = None,
+) -> bool:
+    """Return True when a volume path can be opened by this IO layer."""
+    try:
+        get_vol_shape(filename, dataset=dataset)
+    except (FileNotFoundError, KeyError, ValueError, OSError):
+        return False
+    return True

connectomics/data/processing/distance.py

@@ -2,6 +2,7 @@
 
 from typing import Dict, Optional, Tuple
 
+import cc3d
 import kimimaro
 import numpy as np
 from scipy.ndimage import binary_fill_holes, distance_transform_edt
@@ -13,8 +14,6 @@
     remove_small_holes,
 )
 
-import cc3d
-
 from .bbox_processor import BBoxInstanceProcessor, BBoxProcessorConfig
 from .quantize import energy_quantize
 
@@ -459,7 +458,7 @@ def kimimaro_config(label: np.ndarray, resolution: Tuple[float, ...]) -> dict:
 
     # --- dust threshold ---
     # Skip instances smaller than a 5³-voxel cube.
-    dust_threshold = max(5 ** label.ndim, 5)
+    dust_threshold = max(5**label.ndim, 5)
 
     # --- flags ---
     # fix_branching: improves branch-point accuracy but ~1.3x slower.
@@ -605,7 +604,10 @@ def precompute_sdt_volume(
 
     t0 = time.time()
     sdt = skeleton_aware_distance_transform(
-        label, resolution=resolution, alpha=alpha, bg_value=bg_value,
+        label,
+        resolution=resolution,
+        alpha=alpha,
+        bg_value=bg_value,
         max_parallel=parallel,
     )
     elapsed = time.time() - t0
@@ -665,7 +667,9 @@ def precompute_skeleton_volume(
                 skel_vol[verts[:, 0], verts[:, 1]] = inst_id
 
     n_skel_voxels = int((skel_vol > 0).sum())
-    print(f"  Skeleton volume: {n_skel_voxels} voxels ({n_skel_voxels / max(skel_vol.size, 1) * 100:.2f}%)")
+    print(
+        f"  Skeleton volume: {n_skel_voxels} voxels ({n_skel_voxels / max(skel_vol.size, 1) * 100:.2f}%)"
+    )
 
     save_volume(output_path, skel_vol)
     print(f"  Saved to {output_path}")
@@ -751,14 +755,24 @@ def compute_edt_with_skeleton(
 def sdt_path_for_label(label_path: str, mode: str = "sdt") -> str:
     """Derive the precomputed cache path from a label file path.
 
+    HDF5 labels produce sibling ``*.h5`` cache files. Zarr dataset paths such
+    as ``data.zarr/seg`` produce sibling arrays inside the same store, for
+    example ``data.zarr/seg_skeleton``. A bare ``data.zarr`` label path falls
+    back to a sibling store such as ``data_skeleton.zarr``.
+
     Args:
         mode: ``"sdt"`` for full SDT, ``"skeleton"`` for skeleton volume.
-
-    Examples:
-        ``train-labels.tif`` → ``train-labels_sdt.h5``
-        ``train-labels.tif`` → ``train-labels_skeleton.h5``
     """
     import os
 
+    if ".zarr" in label_path:
+        zarr_idx = label_path.index(".zarr")
+        zarr_path = label_path[: zarr_idx + 5]
+        sub_key = label_path[zarr_idx + 5 :].strip("/")
+        if sub_key:
+            return f"{zarr_path}/{sub_key}_{mode}"
+        base = zarr_path[: -len(".zarr")]
+        return f"{base}_{mode}.zarr"
+
     base, _ = os.path.splitext(label_path)
     return base + f"_{mode}.h5"

connectomics/training/lightning/callbacks.py

@@ -159,6 +159,9 @@ def on_train_batch_end(
         """Store first batch for epoch-end visualization."""
         if batch_idx == 0:
             self._last_train_batch = self._build_cached_batch(batch)
+            # Log image+label on the very first batch (no prediction) for data sanity check.
+            if trainer.current_epoch == 0 and trainer.logger is not None:
+                self._log_data_check(trainer, batch)
 
     def on_validation_batch_end(
         self,
@@ -324,6 +327,31 @@ def _log_visualization(
                 selected_channels=self.selected_channels,
             )
 
+    def _log_data_check(self, trainer, batch: Dict[str, torch.Tensor]) -> None:
+        """Log image + label from the first training batch (no prediction).
+
+        Runs once at the start of training so the user can visually verify
+        data loading, augmentation, and label transforms before waiting for
+        the first epoch to finish.
+        """
+        try:
+            writer = trainer.logger.experiment
+            image = batch["image"].cpu()
+            label = batch["label"].cpu()
+
+            self._log_visualization(
+                image=image,
+                label=label,
+                mask=None,
+                pred=label,  # show label in the pred slot too (no model output yet)
+                writer=writer,
+                iteration=0,
+                prefix="data_check",
+            )
+            logger.info("Logged data check visualization (image + label, no prediction)")
+        except Exception as e:
+            logger.warning("Data check visualization failed: %s", e)
+
     @staticmethod
     def _to_tensor(pred):
         """Extract a tensor from possible deep-supervision dict outputs."""

connectomics/training/lightning/data_factory.py

@@ -16,7 +16,7 @@
     build_val_transforms,
 )
 from ...data.datasets import create_data_dicts_from_paths
-from ...data.io import get_vol_shape
+from ...data.io import get_vol_shape, volume_exists
 from .data import ConnectomicsDataModule, SimpleDataModule
 from .path_utils import expand_file_paths
 
@@ -76,12 +76,10 @@ def _maybe_precompute_label_aux(
 
     print(f"label_aux_type={mode} ({split_name}): " f"resolution={list(resolution)}, alpha={alpha}")
 
-    import os
-
     paths = []
     for lp in label_paths:
         sp = sdt_path_for_label(lp, mode=mode)
-        if not os.path.exists(sp):
+        if not volume_exists(sp):
             if mode == "sdt":
                 precompute_sdt_volume(lp, sp, resolution=resolution, alpha=alpha, bg_value=bg_value)
             else:
@@ -616,7 +614,7 @@ def create_datamodule(
             logger.info("Auto-computing iter_num from volume size...")
 
             from ...data.datasets.sampling import compute_total_samples
-            from ...data.io import get_vol_shape
+            from ...data.io import get_vol_shape, volume_exists
 
             # Get volume sizes
             volume_sizes = []