Support balanced multichannel overlay as RGB export

extensions/chaeo/products.py

@@ -10,9 +10,16 @@ from skimage.measure import shannon_entropy
 from tifffile import imwrite
 
 from extensions.chaeo.annotators import draw_box_on_patch
-from model_server.accessors import GenericImageDataAccessor
+from model_server.accessors import GenericImageDataAccessor, InMemoryDataAccessor
 from model_server.process import pad, rescale, resample_to_8bit
 
+def _make_rgb(zs):
+    h, w, c, nz = zs.shape
+    assert c <= 3
+    outdata = np.zeros((h, w, 3, nz), dtype=zs.dtype)
+    outdata[:, :, 0:c, :] = zs[:, :, :, :]
+    return outdata
+
 def _focus_metrics():
     return {
         'max_intensity': lambda x: np.max(x),
@@ -29,9 +36,15 @@ def _write_patch_to_file(where, fname, data):
 
     if ext == 'PNG':
         assert data.dtype == 'uint8', f'Invalid data type {data.dtype}'
-        assert data.shape[2] == 1, f'Cannot export multichannel images as PNGs; RGB not supported'
+        assert data.shape[2] <= 3, f'Cannot export images with more than 3 channels as PNGs'
         assert data.shape[3] == 1, f'Cannot export z-stacks as PNGs'
-        imsave(where / fname, data[:, :, 0, 0], check_contrast=False)
+        if data.shape[2] == 1:
+            outdata = data[:, :, 0, 0]
+        elif data.shape[2] == 2: # add a blank blue channel
+            outdata = _make_rgb(data)
+        else: # preserve RGB order
+            outdata = data[:, :, :, 0]
+        imsave(where / fname, outdata, check_contrast=False)
         return True
 
     elif ext in ['TIF', 'TIFF']:
@@ -74,7 +87,7 @@ def export_patches_from_zstack(
         focus_metric: str = None,
         **kwargs
 ):
-    assert stack.chroma == 1, 'Expecting monochromatic image data'
+    # assert stack.chroma == 1, 'Expecting monochromatic image data'
     assert stack.nz > 1, 'Expecting z-stack'
 
     exported = []
@@ -121,9 +134,13 @@ def export_patches_from_zstack(
             patch = rescale(patch, rescale_clip)
 
         if kwargs.get('draw_bounding_box') is True:
+            bci = kwargs.get('bounding_box_channel', 0)
+            assert bci < 3
+            if bci > 0:
+                patch = _make_rgb(patch)
             for zi in range(0, patch.shape[3]):
-                patch[:, :, 0, zi] = draw_box_on_patch(
-                    patch[:, :, 0, zi],
+                patch[:, :, bci, zi] = draw_box_on_patch(
+                    patch[:, :, bci, zi],
                     ((x0, y0), (x1, y1)),
                 )
 
@@ -212,3 +229,49 @@ def export_3d_patches_with_focus_metrics(
         exported.append(fstem + '.csv')
 
     return exported
+
+def export_multichannel_patches_from_zstack(
+    where: Path,
+    stack: GenericImageDataAccessor,
+    zmask_meta: list,
+    ch_rgb_overlay: tuple = None,
+    overlay_gain: tuple = (1.0, 1.0, 1.0),
+    ch_white: int = None,
+    **kwargs
+):
+    def _safe_add(a, g, b):
+        assert a.dtype == b.dtype
+        assert a.shape == b.shape
+        assert g >= 0.0
+
+        return np.clip(
+            a.astype('uint32') + g * b.astype('uint32'),
+            0,
+            np.iinfo(a.dtype).max
+        ).astype(a.dtype)
+
+    idata = stack.data
+    if ch_white:
+        assert ch_white < stack.chroma
+        mdata = idata[:, :, [ch_white, ch_white, ch_white], :]
+    else:
+        mdata = idata
+
+    if ch_rgb_overlay:
+        assert len(ch_rgb_overlay) == 3
+        assert len(overlay_gain) == 3
+        for ii, ci in enumerate(ch_rgb_overlay):
+            if ci is None:
+                continue
+            assert isinstance(ci, int)
+            assert ci < stack.chroma
+            mdata[:, :, ii, :] = _safe_add(
+                mdata[:, :, ii, :],
+                overlay_gain[ii],
+                idata[:, :, ci, :]
+            )
+
+    mstack = InMemoryDataAccessor(mdata)
+    return export_patches_from_zstack(
+        where, mstack, zmask_meta, **kwargs
+    )
\ No newline at end of file

extensions/chaeo/tests/test_zstack.py

@@ -5,7 +5,7 @@ import numpy as np
 from conf.testing import output_path
 
 from extensions.chaeo.conf.testing import multichannel_zstack, pixel_classifier, pipeline_params
-from extensions.chaeo.products import export_patches_from_zstack
+from extensions.chaeo.products import export_patches_from_zstack, export_multichannel_patches_from_zstack
 from extensions.chaeo.zmask import build_zmask_from_object_mask
 from model_server.accessors import generate_file_accessor, InMemoryDataAccessor, write_accessor_data_to_file
 from extensions.ilastik.models import IlastikObjectClassifierModel, IlastikPixelClassifierModel
@@ -115,4 +115,21 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
         write_accessor_data_to_file(
             output_path / 'flattened.tif',
             InMemoryDataAccessor(img)
-        )
\ No newline at end of file
+        )
+
+    def test_make_multichannel_2d_patches_from_zmask(self):
+        zmask, meta = self.test_zmask_makes_correct_boxes(
+            filters={'area': (1e3, 1e4)},
+            expand_box_by=(128, 2)
+        )
+        files = export_multichannel_patches_from_zstack(
+            output_path / '2d_patches_chlorophyl',
+            InMemoryDataAccessor(self.stack.data),
+            meta,
+            ch_white=4,
+            ch_rgb_overlay=(3, None, None),
+            draw_bounding_box=True,
+            bounding_box_channel=1,
+            overlay_gain=(0.1, 1.0, 1.0)
+        )
+        self.assertGreaterEqual(len(files), 1)
\ No newline at end of file