Tests pass

model_server/extensions/chaeo/annotators.py

@@ -30,9 +30,10 @@ def draw_boxes_on_2d_image(yx_img, boxes, **kwargs):
 
 
 def draw_boxes_on_3d_image(yxcz_img, boxes, draw_full_depth=False, **kwargs):
-    assert len(yxcz_img.shape) == 4
-    nz = yxcz_img.shape[3]
-    assert yxcz_img.shape[2] == 1
+    # assert len(yxcz_img.shape) == 4
+    # nz = yxcz_img.shape[3]
+    _, _, chroma, nz = yxcz_img.shape
+    # assert yxcz_img.shape[2] == 1
 
     annotated = np.zeros(yxcz_img.shape, dtype=yxcz_img.dtype)
 
@@ -41,7 +42,10 @@ def draw_boxes_on_3d_image(yxcz_img, boxes, draw_full_depth=False, **kwargs):
             zi_boxes = boxes
         else:
             zi_boxes = [bb for bb in boxes if bb['info'].zi == zi]
-        annotated[:, :, 0, zi] = draw_boxes_on_2d_image(yxcz_img[:, :, 0, zi], zi_boxes, **kwargs)
+        # annotated[:, :, 0, zi] = draw_boxes_on_2d_image(yxcz_img[:, :, 0, zi], zi_boxes, **kwargs)
+        for c in range(0, chroma):
+            annotated[:, :, c, zi] = draw_boxes_on_2d_image(yxcz_img[:, :, c, zi], zi_boxes, **kwargs)
+
 
     if clip := kwargs.get('rescale_clip'):
         assert clip >= 0.0 and clip <= 1.0

model_server/extensions/chaeo/tests/test_zstack.py

@@ -48,7 +48,8 @@ class TestRoiSetMonoProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
             self.stack_ch_pa,
             params=RoiSetMetaParams(
                 mask_type=mask_type,
-                filters=kwargs.get('filters', {'area': {'min': 1e3, 'max': 1e4}})
+                filters=kwargs.get('filters', {'area': {'min': 1e3, 'max': 1e4}}),
+                expand_box_by=(64, 2)
             )
         )
         return roiset
@@ -100,7 +101,7 @@ class TestRoiSetMonoProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
         return self._make_roi_set()
 
     def test_zmask_makes_correct_expanded_boxes(self):
-        return self._make_roi_set(expand_box_by=(64, 2))
+        return self._make_roi_set()
 
     def test_zmask_slices_are_valid(self):
         roiset = self._make_roi_set()
@@ -120,10 +121,7 @@ class TestRoiSetMonoProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
             self.assertTrue(np.all([si >= 1 for si in rbb.shape]))
 
     def test_make_2d_patches_from_zmask(self):
-        roiset = self._make_roi_set(
-            # filters={'area': {'min': 1e3, 'max': 1e4}},
-            expand_box_by=(64, 2)
-        )
+        roiset = self._make_roi_set()
         files = export_patches_from_zstack(
             output_path / '2d_patches',
             roiset,
@@ -132,10 +130,7 @@ class TestRoiSetMonoProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
         self.assertGreaterEqual(len(files), 1)
 
     def test_make_3d_patches_from_zmask(self):
-        roiset = self._make_roi_set(
-            # filters={'area': {'min': 1e3, 'max': 1e4}},
-            expand_box_by=(64, 2),
-        )
+        roiset = self._make_roi_set()
         files = export_patches_from_zstack(
             output_path / '3d_patches',
             roiset,
@@ -143,10 +138,7 @@ class TestRoiSetMonoProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
         self.assertGreaterEqual(len(files), 1)
 
     def test_export_annotated_zstack(self):
-        roiset = self._make_roi_set(
-            # filters={'area': {'min': 1e3, 'max': 1e4}},
-            expand_box_by=(64, 2),
-        )
+        roiset = self._make_roi_set()
         file = roiset.export_annotated_zstack(
             output_path / 'annotated_stack',
         )
@@ -177,10 +169,7 @@ class TestRoiSetMonoProducts(BaseTestRoiSetMonoProducts, unittest.TestCase):
         )
 
     def test_make_binary_masks_from_zmask(self):
-        roiset = self._make_roi_set(
-            # filters={'area': {'min': 1e3, 'max': 1e4}},
-            expand_box_by=(128, 2)
-        )
+        roiset = self._make_roi_set()
         files = roiset.export_patch_masks(output_path / '2d_mask_patches', )
         self.assertGreaterEqual(len(files), 1)
 

model_server/extensions/chaeo/workflows.py

@@ -63,7 +63,7 @@ def infer_object_map_from_zstack(
 
     return {
         'timer_results': ti.events,
-        'dataframe': rois.df,
+        'dataframe': rois.get_df(),
         'interm': {},
         'output_path': output_folder_path,
     }

model_server/extensions/chaeo/zmask.py

@@ -78,7 +78,7 @@ class RoiSet(object):
 
         self.count = len(self.zmask_meta)
         self.object_id_labels = self.interm['label_map']
-        self.object_class_map = {}  # classification results
+        self.object_class_maps = {}  # classification results
 
     def __iter__(self):
         """Expose ROI meta information via the Pandas.DataFrame API"""
@@ -273,7 +273,7 @@ class RoiSet(object):
             se.loc[idx[i]] = object_class
 
         self.add_df_col('classify_by_' + name, se)
-        self.object_class_map[name] = InMemoryDataAccessor(om)
+        self.object_class_maps[name] = InMemoryDataAccessor(om)
 
     def run_exports(self, where, channel, prefix, params: RoiSetExportParams):
         if not self.count:
@@ -304,9 +304,10 @@ class RoiSet(object):
             if k == 'patch_masks':
                 self.export_patch_masks(subdir, prefix=pr, **kp)
             if k == 'annotated_zstacks':
-                self.export_annotated_zstack(prefix=pr, **kp)
+                self.export_annotated_zstack(subdir, prefix=pr, **kp)
             if k == 'object_classes':
-                write_accessor_data_to_file(subdir / (pr + '.tif'), self.object_class_map)
+                for k, acc in self.object_class_maps.items():
+                    write_accessor_data_to_file(subdir / k / (pr + '.tif'), acc)
             if k == 'dataframe':
                 dfpa = subdir / (pr + '.csv')
                 dfpa.parent.mkdir(parents=True, exist_ok=True)