update read_images

Author: Donglai Wei

connectomics/data/io/io.py

@@ -9,20 +9,18 @@
 """
 
 from __future__ import annotations
-from typing import Optional, List, Union
-import os
+
 import glob
+import os
 import pickle
+from typing import List, Optional, Union
+
 import h5py
-import numpy as np
 import imageio
 import nibabel as nib
+import numpy as np
 
-# Avoid PIL "IOError: image file truncated"
-from PIL import ImageFile
-
-ImageFile.LOAD_TRUNCATED_IMAGES = True
-
+from .utils import rgb_to_seg
 
 # =============================================================================
 # HDF5 I/O
@@ -107,7 +105,9 @@ def list_hdf5_datasets(filename: str) -> List[str]:
 SUPPORTED_IMAGE_FORMATS = ["png", "tif", "tiff", "jpg", "jpeg"]
 
 
-def read_image(filename: str, add_channel: bool = False) -> Optional[np.ndarray]:
+def read_image(
+    filename: str, add_channel: bool = False, image_type: str = "image"
+) -> Optional[np.ndarray]:
     """Read a single image file.
 
     Args:
@@ -121,12 +121,14 @@ def read_image(filename: str, add_channel: bool = False) -> Optional[np.ndarray]
         return None
 
     image = imageio.imread(filename)
+    if image_type == "seg" and image.ndim == 3:
+        image = rgb_to_seg(image)
     if add_channel and image.ndim == 2:
         image = image[:, :, None]
     return image
 
 
-def read_images(filename_pattern: str) -> np.ndarray:
+def read_images(filename_pattern: str, image_type: str = "image") -> np.ndarray:
     """Read multiple images from a filename pattern.
 
     Args:
@@ -143,17 +145,11 @@ def read_images(filename_pattern: str) -> np.ndarray:
         raise ValueError(f"No files found matching pattern: {filename_pattern}")
 
     # Determine array shape from first image
-    first_image = imageio.imread(file_list[0])
-    if first_image.ndim == 2:
-        data = np.zeros((len(file_list), *first_image.shape), dtype=first_image.dtype)
-    elif first_image.ndim == 3:
-        data = np.zeros((len(file_list), *first_image.shape), dtype=first_image.dtype)
-    else:
-        raise ValueError(f"Unsupported image dimensions: {first_image.ndim}D")
-
+    first_image = read_image(file_list[0], image_type=image_type)
+    data = np.zeros((len(file_list), *first_image.shape), dtype=first_image.dtype)
     # Load all images
     for i, filepath in enumerate(file_list):
-        data[i] = imageio.imread(filepath)
+        data[i] = read_image(filepath, image_type=image_type)
 
     return data
 
@@ -171,7 +167,7 @@ def read_image_as_volume(filename: str, drop_channel: bool = False) -> np.ndarra
     Raises:
         ValueError: If file format is not supported
     """
-    image_suffix = filename[filename.rfind(".") + 1:].lower()
+    image_suffix = filename[filename.rfind(".") + 1 :].lower()
     if image_suffix not in SUPPORTED_IMAGE_FORMATS:
         raise ValueError(
             f"Unsupported format: {image_suffix}. Supported formats: {SUPPORTED_IMAGE_FORMATS}"
@@ -281,7 +277,7 @@ def read_volume(
     if filename.endswith(".nii.gz"):
         image_suffix = "nii.gz"
     else:
-        image_suffix = filename[filename.rfind(".") + 1:].lower()
+        image_suffix = filename[filename.rfind(".") + 1 :].lower()
 
     if image_suffix in ["h5", "hdf5"]:
         data = read_hdf5(filename, dataset)
@@ -420,7 +416,7 @@ def get_vol_shape(filename: str, dataset: Optional[str] = None) -> tuple:
     if filename.endswith(".nii.gz"):
         image_suffix = "nii.gz"
     else:
-        image_suffix = filename[filename.rfind(".") + 1:].lower()
+        image_suffix = filename[filename.rfind(".") + 1 :].lower()
 
     if image_suffix in ["h5", "hdf5"]:
         # HDF5: Read shape from metadata (no data loading)

connectomics/metrics/metrics_seg.py

@@ -48,24 +48,55 @@ class AdaptedRandError(torchmetrics.Metric):
 
     This wrapper lets us accumulate scores during Lightning `test_step` without
     manual numpy↔torch conversions in the training loop.
+
+    Args:
+        return_all_stats: If True, also compute and return precision and recall
+        dist_sync_on_step: Whether to sync across distributed processes on each step
     """
 
     full_state_update: bool = False
 
-    def __init__(self, dist_sync_on_step: bool = False) -> None:
+    def __init__(self, return_all_stats: bool = False, dist_sync_on_step: bool = False) -> None:
         super().__init__(dist_sync_on_step=dist_sync_on_step)
+        self.return_all_stats = return_all_stats
+
         self.add_state("total", default=torch.tensor(0.0), dist_reduce_fx="sum")
         self.add_state("count", default=torch.tensor(0), dist_reduce_fx="sum")
 
+        if return_all_stats:
+            self.add_state("total_precision", default=torch.tensor(0.0), dist_reduce_fx="sum")
+            self.add_state("total_recall", default=torch.tensor(0.0), dist_reduce_fx="sum")
+
     def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
         # Move to CPU and numpy for the underlying implementation
         preds_np = preds.detach().cpu().numpy()
         target_np = target.detach().cpu().numpy()
-        score = float(adapted_rand(preds_np, target_np))
-        self.total += torch.tensor(score, device=self.total.device)
+
+        if self.return_all_stats:
+            are, precision, recall = adapted_rand(preds_np, target_np, all_stats=True)
+            self.total += torch.tensor(are, device=self.total.device)
+            self.total_precision += torch.tensor(precision, device=self.total_precision.device)
+            self.total_recall += torch.tensor(recall, device=self.total_recall.device)
+        else:
+            score = float(adapted_rand(preds_np, target_np, all_stats=False))
+            self.total += torch.tensor(score, device=self.total.device)
+
         self.count += 1
 
     def compute(self) -> torch.Tensor:
         if self.count == 0:
+            if self.return_all_stats:
+                return {
+                    "adapted_rand_error": torch.tensor(0.0, device=self.total.device),
+                    "adapted_rand_precision": torch.tensor(0.0, device=self.total.device),
+                    "adapted_rand_recall": torch.tensor(0.0, device=self.total.device),
+                }
             return torch.tensor(0.0, device=self.total.device)
+
+        if self.return_all_stats:
+            return {
+                "adapted_rand_error": self.total / self.count,
+                "adapted_rand_precision": self.total_precision / self.count,
+                "adapted_rand_recall": self.total_recall / self.count,
+            }
         return self.total / self.count

connectomics/training/lit/model.py

@@ -194,7 +194,8 @@ def _setup_test_metrics(self):
                 self.test_accuracy = torchmetrics.Accuracy(task='multiclass', num_classes=num_classes).to(self.device)
         if 'adapted_rand' in metrics:
             from ...metrics.metrics_seg import AdaptedRandError
-            self.test_adapted_rand = AdaptedRandError().to(self.device)
+            # Enable all_stats to also compute precision and recall
+            self.test_adapted_rand = AdaptedRandError(return_all_stats=True).to(self.device)
 
     def _invert_save_prediction_transform(self, data: np.ndarray) -> np.ndarray:
         """
@@ -349,14 +350,30 @@ def _compute_test_metrics(self, decoded_predictions: np.ndarray, labels: torch.T
             # Adapted Rand Error is for instance segmentation
             if hasattr(self, "test_adapted_rand") and isinstance(self.test_adapted_rand, torchmetrics.Metric):
                 from ...metrics.metrics_seg import AdaptedRandError
-                per_volume_metric = AdaptedRandError().to(self.device)
+                # Use return_all_stats=True to get precision and recall
+                per_volume_metric = AdaptedRandError(return_all_stats=True).to(self.device)
                 per_volume_metric.update(pred_instances.cpu(), labels_instances.cpu())
-                adapted_rand_value = per_volume_metric.compute()
-                print(f"  {volume_prefix}Adapted Rand Error: {adapted_rand_value.item():.6f}")
+                adapted_rand_stats = per_volume_metric.compute()
+
+                # Print per-volume metrics
+                if isinstance(adapted_rand_stats, dict):
+                    print(f"  {volume_prefix}Adapted Rand Error: {adapted_rand_stats['adapted_rand_error'].item():.6f}")
+                    print(f"  {volume_prefix}Adapted Rand Precision: {adapted_rand_stats['adapted_rand_precision'].item():.6f}")
+                    print(f"  {volume_prefix}Adapted Rand Recall: {adapted_rand_stats['adapted_rand_recall'].item():.6f}")
+                else:
+                    print(f"  {volume_prefix}Adapted Rand Error: {adapted_rand_stats.item():.6f}")
 
                 # Update running metric for epoch-level aggregation
                 self.test_adapted_rand.update(pred_instances.cpu(), labels_instances.cpu())
-                self.log("test_adapted_rand", self.test_adapted_rand, on_step=False, on_epoch=True, prog_bar=True, logger=True)
+
+                # Log metrics - handle both dict and tensor return values
+                epoch_stats = self.test_adapted_rand.compute()
+                if isinstance(epoch_stats, dict):
+                    self.log("test_adapted_rand", epoch_stats['adapted_rand_error'], on_step=False, on_epoch=True, prog_bar=True, logger=True)
+                    self.log("test_adapted_rand_precision", epoch_stats['adapted_rand_precision'], on_step=False, on_epoch=True, prog_bar=True, logger=True)
+                    self.log("test_adapted_rand_recall", epoch_stats['adapted_rand_recall'], on_step=False, on_epoch=True, prog_bar=True, logger=True)
+                else:
+                    self.log("test_adapted_rand", epoch_stats, on_step=False, on_epoch=True, prog_bar=True, logger=True)
 
         else:
             # For binary/semantic segmentation: binarize predictions

scripts/images_to_h5.py

@@ -25,16 +25,18 @@
 Note: Use quotes around the input pattern to prevent shell expansion.
 """
 
-import sys
 import os
-from pathlib import Path
-from connectomics.data.io import read_volume, write_hdf5
+import sys
+
+from connectomics.data.io import read_images, write_hdf5
 
 
 def main():
     """Main conversion function."""
     if len(sys.argv) < 3:
-        print("Usage: python scripts/images_to_h5.py <input_pattern> <output_file.h5> [dataset_key]")
+        print(
+            "Usage: python scripts/images_to_h5.py <input_pattern> <output_file.h5> [dataset_key]"
+        )
         print("")
         print("Examples:")
         print('  python scripts/images_to_h5.py "datasets/images/*.tiff" output.h5')
@@ -46,48 +48,27 @@ def main():
 
     input_pattern = sys.argv[1]
     output_file = sys.argv[2]
-    dataset_key = sys.argv[3] if len(sys.argv) > 3 else "main"
+    image_type = sys.argv[3] if len(sys.argv) > 3 else "image"
+    dataset_key = sys.argv[4] if len(sys.argv) > 4 else "main"
 
     # Ensure output directory exists
     output_dir = os.path.dirname(output_file)
     if output_dir and not os.path.exists(output_dir):
         print(f"Creating output directory: {output_dir}")
         os.makedirs(output_dir, exist_ok=True)
 
-    # Detect file format from pattern
-    pattern_lower = input_pattern.lower()
-    if any(ext in pattern_lower for ext in ['.tif', '.tiff']):
-        format_name = "TIFF"
-    elif '.png' in pattern_lower:
-        format_name = "PNG"
-    elif any(ext in pattern_lower for ext in ['.jpg', '.jpeg']):
-        format_name = "JPEG"
-    else:
-        format_name = "image"
-
-    print(f"Reading {format_name} files matching: {input_pattern}")
-    print("This may take a while for large volumes...")
-
     # Read all image files as a 3D volume
-    try:
-        volume = read_volume(input_pattern)
-    except Exception as e:
-        print(f"Error reading images: {e}")
-        print("\nTips:")
-        print("  - Check that the file pattern is correct")
-        print("  - Ensure all images have the same dimensions")
-        print("  - Verify the image files are readable")
-        sys.exit(1)
+    volume = read_images(input_pattern, image_type=image_type)
 
-    print(f"\n{'='*60}")
-    print(f"Volume Information:")
-    print(f"{'='*60}")
+    print(f"\n{'=' * 60}")
+    print("Volume Information:")
+    print(f"{'=' * 60}")
     print(f"  Shape:      {volume.shape}")
     print(f"  Data type:  {volume.dtype}")
     print(f"  Size:       {volume.nbytes / (1024**3):.2f} GB")
     print(f"  Min value:  {volume.min()}")
     print(f"  Max value:  {volume.max()}")
-    print(f"{'='*60}")
+    print(f"{'=' * 60}")
 
     print(f"\nSaving to: {output_file}")
     print(f"Dataset key: '{dataset_key}'")