From 489bc836bd7a5462a4d0c917e4f1ba6037d01701 Mon Sep 17 00:00:00 2001 From: maragall <126172415+maragall@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:17:22 -0400 Subject: [PATCH] fix: restore RGB mosaic-view save (plan R7, overview path) The unified mosaic/plate refactor dropped RGB layers from the mosaic-view save entirely (_snapshot_for_save skipped them with 'not supported yet', deferred as plan R7). With MULTIPOINT_BF_SAVING_OPTION set to RGB2GRAY or Green Channel Only, the per-frame color is already destroyed at save time, so the colored mosaic-view PNG that older Squid wrote was the ONLY surviving record of BF stain color. Downstream consumers, e.g. SquidXplorer's stain-color reconstruction, read the sidecar's rgb_channel_names / rgb_view_files keys to restore it; current Squid acquisitions gave them nothing to read. RGB layers are now collected alongside the monochrome channels and, when SAVE_DOWNSAMPLED_OVERVIEW is on, written as 8-bit RGB PNGs (mosaic__um_.png, spaces to underscores like elsewhere in the repo) with rgb_channel_names / rgb_view_files recorded in the yaml sidecar under the exact key spelling older Squid used, so existing consumers read the new sidecars unchanged. The monochrome path is untouched: same ome.tiff, same sidecar keys, and per-well TIFFs stay monochrome-only. A snapshot is now skipped only when BOTH lists are empty. Tests: TestRGBMosaicSave in test_unified_mosaic_widget.py covers snapshot collection, PNG + sidecar writes, the RGB-only case (uint16 RGB48 scaled to 8-bit), and that mono-only sidecars carry no rgb keys. Co-Authored-By: Claude Fable 5 --- software/control/widgets_mosaic.py | 51 ++++++++-- .../control/test_unified_mosaic_widget.py | 92 +++++++++++++++++++ 2 files changed, 136 insertions(+), 7 deletions(-) diff --git a/software/control/widgets_mosaic.py b/software/control/widgets_mosaic.py index 17a6ae662..55d02c8a0 100644 --- a/software/control/widgets_mosaic.py +++ b/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" diff --git a/software/tests/control/test_unified_mosaic_widget.py b/software/tests/control/test_unified_mosaic_widget.py index 62a53aa17..2d752d6fc 100644 --- a/software/tests/control/test_unified_mosaic_widget.py +++ b/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