Restore RGB mosaic-view save (closes the untracked R7-cited regression)

software/control/widgets_mosaic.py

@@ -12,6 +12,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import List, Optional, Tuple
 
+import imageio
 import numpy as np
 import tifffile
 import yaml
@@ -107,6 +108,19 @@ def blit_tiles_to_canvas(
         canvas[dst_y_start:dst_y_end, dst_x_start:dst_x_end] = tile[src_y_start:src_y_end, src_x_start:src_x_end]
 
 
+def _rgb_to_uint8(data: np.ndarray) -> np.ndarray:
+    """Convert an (H, W, 3) RGB canvas to uint8 for PNG saving. RGB layers keep
+    the source dtype (uint8 from 8-bit cameras, uint16 from RGB48), so a
+    higher-depth canvas is scaled by its dtype's full range; floats are assumed
+    display-normalized in [0, 1]."""
+    if data.dtype == np.uint8:
+        return data
+    if np.issubdtype(data.dtype, np.integer):
+        max_val = np.iinfo(data.dtype).max
+        return (data.astype(np.float32) * (255.0 / max_val)).astype(np.uint8)
+    return (np.clip(data, 0.0, 1.0) * 255.0).astype(np.uint8)
+
+
 class UnifiedMosaicWidget(QWidget):
     """Single widget for mosaic and plate view display.
 
@@ -862,19 +876,25 @@ def _snapshot_for_save(self) -> Optional[dict]:
             self._log.warning("Save skipped: viewer_pixel_size_mm is unset.")
             return None
         channels = []
+        rgb_channels = []
         for layer in self._image_layers():
             if layer.data.ndim == 3 and layer.data.shape[2] == 3:
-                # Defer RGB save support — see plan R7.
-                self._log.warning(f"Skipping RGB layer '{layer.name}' from save (not supported yet).")
+                # RGB overview save restored (plan R7 done for the overview path):
+                # saved as a colored PNG + rgb_channel_names/rgb_view_files sidecar
+                # keys, matching what older Squid wrote. Downstream consumers
+                # (e.g. SquidXplorer's stain-color reconstruction) read these to
+                # recover BF color when frames were saved grayscale.
+                rgb_channels.append((layer.name, np.array(layer.data, copy=True)))
                 continue
             channels.append((layer.name, np.array(layer.data, copy=True)))
-        if not channels:
-            self._log.warning("Save skipped: no monochrome image layers present.")
+        if not channels and not rgb_channels:
+            self._log.warning("Save skipped: no image layers present.")
             return None
         snapshot = {
             "mode": self.mode.value,
             "resolution_um": resolution_um,
             "channels": channels,
+            "rgb_channels": rgb_channels,
             "saved_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
             # Capture the flag values now so toggling them between snapshot and
             # write doesn't leave the sidecar describing one thing and the
@@ -908,14 +928,15 @@ def _write_save_snapshot(self, target_dir: str, snapshot: dict) -> None:
             resolution_um = snapshot["resolution_um"]
             res_tag = f"{int(round(resolution_um))}um"
             channels = snapshot["channels"]
+            rgb_channels = snapshot.get("rgb_channels") or []
             channel_names = [name for name, _ in channels]
-            sidecar = {k: v for k, v in snapshot.items() if k != "channels"}
+            sidecar = {k: v for k, v in snapshot.items() if k not in ("channels", "rgb_channels")}
             sidecar["channel_names"] = channel_names
 
             save_overview = snapshot["save_overview"]
             save_per_well = snapshot["save_per_well"] and mode == DisplayMode.PLATE.value
 
-            if save_overview:
+            if save_overview and channels:
                 stack = np.stack([data for _, data in channels], axis=0)  # (C, H, W)
                 whole_path = os.path.join(target_dir, f"mosaic_{mode}_{res_tag}.ome.tiff")
                 tifffile.imwrite(
@@ -935,7 +956,23 @@ def _write_save_snapshot(self, target_dir: str, snapshot: dict) -> None:
                 sidecar["whole_view_file"] = os.path.basename(whole_path)
                 self._log.info(f"Saved whole view: {whole_path} ({stack.shape}, {stack.nbytes/1e6:.1f} MB)")
 
-            if save_per_well:
+            if save_overview and rgb_channels:
+                # Colored overview PNGs + sidecar keys, spelled exactly as older
+                # Squid wrote them (rgb_channel_names / rgb_view_files) so
+                # existing consumers read them unchanged.
+                rgb_names = []
+                rgb_files = []
+                for name, data in rgb_channels:
+                    safe_name = name.replace(" ", "_")
+                    png_path = os.path.join(target_dir, f"mosaic_{mode}_{res_tag}_{safe_name}.png")
+                    imageio.imwrite(png_path, _rgb_to_uint8(data))
+                    rgb_names.append(name)
+                    rgb_files.append(os.path.basename(png_path))
+                    self._log.info(f"Saved RGB view: {png_path} ({data.shape})")
+                sidecar["rgb_channel_names"] = rgb_names
+                sidecar["rgb_view_files"] = rgb_files
+
+            if save_per_well and channels:
                 self._write_per_well_tiffs(target_dir, snapshot, res_tag)
                 sidecar["per_well_dir"] = "wells"
 

software/tests/control/test_unified_mosaic_widget.py

@@ -155,3 +155,95 @@ def test_plate_view_still_uses_integer_downsample(self, qtbot, monkeypatch):
         )
         # Integer factor 3 -> 2.22 um, NOT the exact target 2.0 um.
         assert widget.viewer_pixel_size_mm == pytest.approx(0.00222, abs=1e-5)
+
+
+class TestRGBMosaicSave:
+    """Plan R7's overview half: RGB layers are saved as colored PNGs with the
+    rgb_channel_names / rgb_view_files sidecar keys older Squid wrote, which
+    SquidXplorer's stain-color reconstruction consumes."""
+
+    @pytest.fixture
+    def save_flags(self, monkeypatch):
+        monkeypatch.setattr(control._def, "SAVE_DOWNSAMPLED_OVERVIEW", True)
+        monkeypatch.setattr(control._def, "SAVE_DOWNSAMPLED_WELL_IMAGES", False)
+
+    def test_snapshot_collects_rgb_alongside_mono(self, mosaic_widget, save_flags):
+        widget, _ = mosaic_widget
+        widget.updateTile(_tile_update(np.full((100, 100), 200, dtype=np.uint8), 10.0, 10.0))
+        rgb = np.zeros((50, 50, 3), dtype=np.uint8)
+        rgb[..., 0] = 255
+        widget.viewer.add_image(rgb, rgb=True, name="BF LED matrix full RGB")
+
+        snapshot = widget._snapshot_for_save()
+        assert snapshot is not None
+        assert [name for name, _ in snapshot["channels"]] == ["BF"]
+        assert [name for name, _ in snapshot["rgb_channels"]] == ["BF LED matrix full RGB"]
+        # The mono path is untouched: mono arrays stay 2-D.
+        assert snapshot["channels"][0][1].ndim == 2
+
+    def test_write_saves_png_and_sidecar_keys(self, mosaic_widget, save_flags, tmp_path):
+        import imageio
+        import yaml
+
+        widget, _ = mosaic_widget
+        widget.updateTile(_tile_update(np.full((100, 100), 200, dtype=np.uint8), 10.0, 10.0))
+        rgb = np.zeros((50, 50, 3), dtype=np.uint8)
+        rgb[..., 1] = 128
+        widget.viewer.add_image(rgb, rgb=True, name="BF LED matrix full")
+
+        snapshot = widget._snapshot_for_save()
+        target = tmp_path / "mosaic_view"
+        widget._write_save_snapshot(str(target), snapshot)
+
+        png = target / "mosaic_mosaic_2um_BF_LED_matrix_full.png"
+        assert png.is_file()
+        assert (target / "mosaic_mosaic_2um.ome.tiff").is_file()  # mono overview untouched
+        with open(target / "mosaic_mosaic_2um.yaml") as f:
+            sidecar = yaml.safe_load(f)
+        assert sidecar["channel_names"] == ["BF"]
+        assert sidecar["rgb_channel_names"] == ["BF LED matrix full"]
+        assert sidecar["rgb_view_files"] == ["mosaic_mosaic_2um_BF_LED_matrix_full.png"]
+        # The PNG round-trips as 8-bit RGB with the layer's pixels.
+        back = imageio.imread(png)
+        assert back.dtype == np.uint8
+        assert back.shape == (50, 50, 3)
+        assert back[0, 0, 1] == 128
+
+    def test_rgb_only_snapshot_still_saves(self, mosaic_widget, save_flags, tmp_path):
+        import imageio
+        import yaml
+
+        widget, _ = mosaic_widget
+        widget.viewer_pixel_size_mm = 0.002
+        rgb = (np.ones((20, 20, 3)) * 65535).astype(np.uint16)  # RGB48 canvas
+        widget.viewer.add_image(rgb, rgb=True, name="BF LED matrix full RGB")
+
+        snapshot = widget._snapshot_for_save()
+        assert snapshot is not None  # no-layers early return needs BOTH lists empty
+        target = tmp_path / "mosaic_view"
+        widget._write_save_snapshot(str(target), snapshot)
+
+        png = target / "mosaic_mosaic_2um_BF_LED_matrix_full_RGB.png"
+        assert png.is_file()
+        assert not (target / "mosaic_mosaic_2um.ome.tiff").exists()  # no mono channels
+        with open(target / "mosaic_mosaic_2um.yaml") as f:
+            sidecar = yaml.safe_load(f)
+        assert sidecar["channel_names"] == []
+        assert sidecar["rgb_view_files"] == ["mosaic_mosaic_2um_BF_LED_matrix_full_RGB.png"]
+        back = imageio.imread(png)
+        assert back.dtype == np.uint8
+        assert int(back.max()) == 255  # uint16 full-scale maps to 255
+
+    def test_mono_only_sidecar_carries_no_rgb_keys(self, mosaic_widget, save_flags, tmp_path):
+        import yaml
+
+        widget, _ = mosaic_widget
+        widget.updateTile(_tile_update(np.full((100, 100), 200, dtype=np.uint8), 10.0, 10.0))
+        snapshot = widget._snapshot_for_save()
+        target = tmp_path / "mosaic_view"
+        widget._write_save_snapshot(str(target), snapshot)
+        with open(target / "mosaic_mosaic_2um.yaml") as f:
+            sidecar = yaml.safe_load(f)
+        assert "rgb_channel_names" not in sidecar
+        assert "rgb_view_files" not in sidecar
+        assert "rgb_channels" not in sidecar  # arrays never leak into the sidecar