Ironing out issues with parameterized patch product export

extensions/chaeo/tests/test_zstack.py

@@ -5,6 +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.params import RoiSetExportParams
 from extensions.chaeo.products import export_patches_from_zstack, export_multichannel_patches_from_zstack, export_patch_masks_from_zstack
 from extensions.chaeo.workflows import infer_object_map_from_zstack
 from extensions.chaeo.zmask import build_zmask_from_object_mask
@@ -207,12 +208,21 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
             self.pxmodel,
             DummyInstanceSegmentationModel(),
         ]
+        export_params = RoiSetExportParams(**{
+            'pixel_probabilities': True,
+            'patches_3d': {},
+            # 'patches_2d_for_annotation': {
+            #     'draw_bounding_box': True,
+            #     'rgb_overlay_channels': [4, 2, None]
+            # }
+        })
         infer_object_map_from_zstack(
             multichannel_zstack['path'],
-            output_path,
+            output_path / 'workflow',
             models,
             pxmap_foreground_channel=pp['pxmap_channel'],
             pxmap_threshold=pp['pxmap_threshold'],
             segmentation_channel=pp['segmentation_channel'],
             patches_channel=pp['patches_channel'],
+            # exports=export_params,
         )
\ No newline at end of file

extensions/chaeo/workflows.py

@@ -283,7 +283,7 @@ def infer_object_map_from_zstack(
     # make zmask
     rois = RoiSet(
         obmask.get_one_channel_data(pxmap_foreground_channel),
-        stack.get_one_channel_data(segmentation_channel),
+        stack,
         mask_type=zmask_type,
         filters=zmask_filters,
         expand_box_by=kwargs.get('zmask_expand_box_by', (0, 0)),

extensions/chaeo/zmask.py

@@ -152,7 +152,7 @@ class RoiSet(object):
         return projected
 
     def get_raw_patches(self, channel):
-        return get_patches_from_zmask_meta(self.acc_raw, self.zmask_meta)
+        return get_patches_from_zmask_meta(self.acc_raw, self.zmask_meta).get_one_channel_data(channel)
 
     def get_patch_masks(self):
         return get_patch_masks_from_zmask_meta(self.acc_raw, self.zmask_meta)
@@ -224,7 +224,7 @@ def build_zmask_from_object_mask(
     """
 
     # validate inputs
-    assert zstack.chroma == 1
+    # assert zstack.chroma == 1
     assert mask_type in ('contours', 'boxes'), mask_type
     assert obmask.is_mask()
     assert obmask.chroma == 1