Added and tested accessor method to run argmax on z-axis

model_server/base/accessors.py

@@ -67,10 +67,12 @@ class GenericImageDataAccessor(ABC):
         """
         return self.apply(lambda x: x.max(axis=self._ga('Z'), keepdims=True))
 
-
     def get_mono(self, channel: int, mip: bool = False):
         return self.get_channels([channel], mip=mip)
 
+    def get_z_argmax(self):
+        return self.apply(lambda x: x.argmax(axis=self.get_axis('Z')))
+
     @property
     def data_xy(self) -> np.ndarray:
         if not self.chroma == 1 and self.nz == 1:

tests/base/test_accessors.py

@@ -61,6 +61,17 @@ class TestCziImageFileAccess(unittest.TestCase):
         sc = cf.get_mono(c, mip=True)
         self.assertEqual(sc.shape, (h, w, 1, 1))
 
+    def test_get_single_channel_argmax_from_zstack(self):
+        w = 256
+        h = 512
+        nc = 4
+        nz = 11
+        c = 3
+        cf = InMemoryDataAccessor(np.random.rand(h, w, nc, nz))
+        am = cf.get_mono(c).get_z_argmax()
+        self.assertEqual(am.shape, (h, w, 1, 1))
+        self.assertTrue(np.all(am.unique()[0] == range(0, nz)))
+
     def test_get_zi(self):
         w = 256
         h = 512