diff --git a/model_server/base/roiset.py b/model_server/base/roiset.py
index 954e878dc2126bf33d5b166c9509cf2c55a4db6d..9b270de5be9c2b99a4e4d4388d4c5f81b5573617 100644
--- a/model_server/base/roiset.py
+++ b/model_server/base/roiset.py
@@ -91,10 +91,9 @@ 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),
+                acc_seg_mask.get_mip().data_xy,
                 connectivity=1,
             ).astype('uint16')
         )