diff --git a/model_server/base/roiset.py b/model_server/base/roiset.py
index 7a95a94774b92929c824cdde52c3fdd893289360..954e878dc2126bf33d5b166c9509cf2c55a4db6d 100644
--- a/model_server/base/roiset.py
+++ b/model_server/base/roiset.py
@@ -91,6 +91,7 @@ def get_label_ids(acc_seg_mask: GenericImageDataAccessor, allow_3d=False, connec
             la_3d[:, :, 0, zi] = la_2d
         return InMemoryDataAccessor(la_3d)
     else:
+        # TODO: call argmax z method
         return InMemoryDataAccessor(
             label(
                 acc_seg_mask.data_xyz.max(axis=-1),
diff --git a/tests/base/test_roiset.py b/tests/base/test_roiset.py
index 87a1d88901a98f98f3940a6cdba603f8dea0ce2c..9de818eea86645dc4168c800414f22aba3803d97 100644
--- a/tests/base/test_roiset.py
+++ b/tests/base/test_roiset.py
@@ -625,8 +625,6 @@ class TestRoiSetSerialization(unittest.TestCase):
             t_acc = generate_file_accessor(pt)
             self.assertTrue(np.all(r_acc.data == t_acc.data))
 
-        self.assertTrue(ref_roiset.contains_segmentation)
-        self.assertTrue(ref_roiset.contains_segmentation)
 
 class TestRoiSetObjectDetection(unittest.TestCase):
 
@@ -660,12 +658,14 @@ class TestRoiSetObjectDetection(unittest.TestCase):
         self.assertEqual(len(table), patches_bbox.count)
 
         # roiset w/ seg for comparison
-        roiset_seg = RoiSet.from_binary_mask(self.stack_ch_pa, mask)
+        roiset_seg = RoiSet.from_binary_mask(self.stack_ch_pa, mask, allow_3d=True)
         patches_seg = roiset_seg.get_patches_acc()
 
-        # TODO: test segments reside in bounding boxes
+        # test bounding box dimensions match those from RoiSet generated directly from segmentation
         self.assertEqual(roiset_seg.count, roiset_bbox.count)
-        self.assertTrue(False)
+        for i in range(0, roiset_seg.count):
+            self.assertEqual(patches_seg.iat(0, crop=True).shape, patches_bbox.iat(0, crop=True).shape)
+
 
 class TestRoiSetPolygons(BaseTestRoiSetMonoProducts, unittest.TestCase):