Fix detach_sample dropping channels of multi-channel views

packages/viscy-utils/src/viscy_utils/log_images.py

@@ -17,13 +17,19 @@ def detach_sample(imgs: Sequence[Tensor], log_samples_per_batch: int) -> list[li
     """Extract middle-Z slices from a batch for image grid logging.
 
     Layout: one row per sample, columns ordered as
-    ``anchor_ch0, anchor_ch1, ..., positive_ch0, positive_ch1, ..., negative_ch0, ...``
-    Channels expand horizontally within each view, which suits landscape monitors.
+    ``view0_ch0, view0_ch1, ..., view1_ch0, view1_ch1, ...``.
+    Each view contributes all of its own channels, so views with
+    different channel counts (e.g. a 1-channel ``source`` with a
+    2-channel ``target``/``pred``) are all logged in full. Channels
+    expand horizontally within each view, which suits landscape monitors.
 
     Parameters
     ----------
     imgs : Sequence[Tensor]
-        One ``(B, C, Z, Y, X)`` tensor per view (anchor, positive, negative).
+        One ``(B, C, Z, Y, X)`` tensor per view. Examples: ``(anchor,
+        positive, negative)`` for contrastive learning, or ``(source,
+        target, pred)`` for virtual staining. Views may have different
+        channel counts.
     log_samples_per_batch : int
         Number of samples from the batch to include (first N).
 
@@ -34,14 +40,13 @@ def detach_sample(imgs: Sequence[Tensor], log_samples_per_batch: int) -> list[li
         ``(view, channel)`` pairs in view-major order.
     """
     num_samples = min(imgs[0].shape[0], log_samples_per_batch)
-    n_channels = imgs[0].shape[1]
     rows = []
     for i in range(num_samples):
         row = []
         for img in imgs:
             patch = to_numpy(img[i])
             mid_z = patch.shape[1] // 2
-            for c in range(n_channels):
+            for c in range(patch.shape[0]):
                 row.append(patch[c, mid_z])
         rows.append(row)
     return rows

packages/viscy-utils/src/viscy_utils/log_images_test.py

@@ -0,0 +1,37 @@
+"""Tests for image grid logging helpers."""
+
+import torch
+
+from viscy_utils.log_images import detach_sample, render_images
+
+
+def test_detach_sample_logs_all_channels_per_view():
+    """Each view contributes all of its own channels, even when counts differ.
+
+    Regression: a single ``n_channels`` taken from the first view truncated
+    multi-channel ``target``/``pred`` to the 1-channel ``source`` count,
+    dropping the membrane channel in virtual-staining grids.
+    """
+    # phase2fluor shapes: source 1ch, target/pred 2ch, (B, C, Z, Y, X).
+    source = torch.rand(2, 1, 5, 16, 16)
+    target = torch.rand(2, 2, 5, 16, 16)
+    pred = torch.rand(2, 2, 5, 16, 16)
+
+    rows = detach_sample((source, target, pred), log_samples_per_batch=2)
+
+    assert len(rows) == 2
+    # 1 (source) + 2 (target) + 2 (pred) = 5 columns per sample.
+    assert all(len(row) == 5 for row in rows)
+
+
+def test_render_images_grid_shape_matches_channels():
+    """Rendered grid width covers every channel of every view."""
+    source = torch.rand(3, 1, 5, 64, 64)
+    target = torch.rand(3, 2, 5, 64, 64)
+    pred = torch.rand(3, 2, 5, 64, 64)
+
+    rows = detach_sample((source, target, pred), log_samples_per_batch=3)
+    grid = render_images(rows)
+
+    # 3 samples * 64 px tall; 5 channels * 64 px wide; RGB.
+    assert grid.shape == (192, 320, 3)