Can export both 2d and 3d patch masks

model_server/base/roiset.py

@@ -864,7 +864,8 @@ class RoiSet(object):
             if self.is_3d and make_3d:
                 return patch
             elif self.is_3d:
-                return patch[:, :, :, roi.zi: roi.zi + 1]
+                rzi = roi.zi - roi.z0
+                return patch[:, :, rzi: rzi + 1]
             else:
                 return np.expand_dims(patch, 2)
 

tests/base/test_roiset.py

@@ -302,7 +302,7 @@ class TestRoiSet3dProducts(unittest.TestCase):
                     'is_patch_mask': True,
                 },
                 '2d_masks': {
-                    'make_3d': True,
+                    'make_3d': False,
                     'is_patch_mask': True,
                 },
             },
@@ -322,7 +322,7 @@ class TestRoiSet3dProducts(unittest.TestCase):
 
         self.assertTrue(any([_get_nz_from_file(f) > 1 for f in res['patches_3d']]))
         self.assertTrue(any([_get_nz_from_file(f) > 1 for f in res['patches_3d_masks']]))
-        self.assertTrue(any([_get_nz_from_file(f) == 1 for f in res['patches_2d_masks']]))
+        self.assertTrue(all([_get_nz_from_file(f) == 1 for f in res['patches_2d_masks']]))
 
 
     def test_run_export_mono_3d_labels_overlay(self):