Converge to a single PatchStack accessor class

model_server/base/roiset.py

@@ -17,7 +17,7 @@ from model_server.base.accessors import GenericImageDataAccessor, InMemoryDataAc
 from model_server.base.models import InstanceSegmentationModel
 from model_server.base.process import pad, rescale, resample_to_8bit, make_rgb
 from base.annotators import draw_box_on_patch, draw_contours_on_patch, draw_boxes_on_3d_image
-from model_server.extensions.chaeo.accessors import write_patch_to_file, MonoPatchStack, Multichannel3dPatchStack
+from model_server.extensions.chaeo.accessors import write_patch_to_file, MonoPatchStack, PatchStack
 from base.process import mask_largest_object
 
 
@@ -234,7 +234,7 @@ class RoiSet(object):
         if channel is not None or self.acc_raw.chroma == 1:
             return MonoPatchStack(patches)
         else:
-            return Multichannel3dPatchStack(patches)
+            return PatchStack(patches)
 
     def export_annotated_zstack(self, where, prefix='zstack', **kwargs):
         annotated = InMemoryDataAccessor(draw_boxes_on_3d_image(self, **kwargs))

model_server/extensions/chaeo/accessors.py

@@ -80,30 +80,33 @@ class MonoPatchStackFromFile(MonoPatchStack):
     def fpath(self):
         return self.file_acc.fpath
 
-# TODO: unify this into one accessor
-class Multichannel3dPatchStack(InMemoryDataAccessor):
+
+class PatchStack(InMemoryDataAccessor):
 
     def __init__(self, data):
         """
         A sequence of n (generally) color 3D images of the same size
-        :param data: a list of np.ndarrays of size YXCZ
+        :param data: either a list of np.ndarrays of size YXCZ, or np.ndarray of size PYXCZ
         """
 
         if isinstance(data, list):  # list of YXCZ patches
-            nda = np.zeros((len(data), *np.array([e.shape for e in data]).max(axis=0)), dtype=data[0].dtype)
+            n = len(data)
+            yxcz_shape = np.array([e.shape for e in data]).max(axis=0)
+            nda = np.zeros(
+                (n, *yxcz_shape), dtype=data[0].dtype
+            )
             for i in range(0, len(data)):
-                nzi = data[i].shape[-1]
-                nda[i, :, :, :, 0:nzi] = data[i]
-            assert nda.ndim == 5
-            # self._data = np.moveaxis( # pos-YXCZ
-            #         nda,
-            #         [0, 1, 2, 0, 3],
-            #         [0, 1, 2, 3]
-            # )
-            self._data = nda
+                s = tuple([slice(0, c) for c in data[i].shape])
+                nda[i][s] = data[i]
+
+        elif isinstance(data, np.ndarray) and len(data.shape) == 5:  # interpret as PYXCZ
+            nda = data
         else:
             raise InvalidDataForPatchStackError(f'Cannot create accessor from {type(data)}')
 
+        assert nda.ndim == 5
+        self._data = nda
+
     def iat(self, i):
         return self.data[i, :, :, :, :]
 
@@ -115,11 +118,15 @@ class Multichannel3dPatchStack(InMemoryDataAccessor):
         return self.shape_dict['P']
 
     @property
-    def data(self):
-        """
-        Return data as 5d with axes in order of pos, Y, X, C, Z
-        :return: np.ndarray
-        """
+    def shape_dict(self):
+        return dict(zip(('P', 'Y', 'X', 'C', 'Z'), self.data.shape))
+
+    def get_list(self):
+        n = self.nz
+        return [self.data[:, :, 0, zi] for zi in range(0, n)]
+
+    @property
+    def pyxcz(self):
         return self._data
 
     @property
@@ -141,7 +148,7 @@ def write_patch_to_file(where, fname, yxcz):
         assert yxcz.shape[3] == 1, f'Cannot export z-stacks as PNGs'
         if yxcz.shape[2] == 1:
             outdata = yxcz[:, :, 0, 0]
-        elif yxcz.shape[2] == 2: # add a blank blue channel
+        elif yxcz.shape[2] == 2:  # add a blank blue channel
             outdata = make_rgb(yxcz)
         else: # preserve RGB order
             outdata = yxcz[:, :, :, 0]

model_server/extensions/chaeo/tests/test_accessors.py

@@ -3,7 +3,7 @@ import unittest
 import numpy as np
 
 from model_server.conf.testing import monozstackmask
-from model_server.extensions.chaeo.accessors import MonoPatchStack, MonoPatchStackFromFile, Multichannel3dPatchStack
+from model_server.extensions.chaeo.accessors import MonoPatchStackFromFile, PatchStack
 
 
 
@@ -15,27 +15,21 @@ class TestMultipositionCziImageFileAccess(unittest.TestCase):
         w = 256
         h = 512
         n = 4
-        acc = MonoPatchStack(np.random.rand(h, w, n))
+        acc = PatchStack(np.random.rand(n, h, w, 1, 1))
         self.assertEqual(acc.count, n)
         self.assertEqual(acc.hw, (h, w))
-        self.assertEqual(acc.make_tczyx().shape, (n, 1, 1, h, w))
+        self.assertEqual(acc.pyxcz.shape, (n, h, w, 1, 1))
 
     def test_make_patch_stack_from_list(self):
         w = 256
         h = 512
         n = 4
-        acc = MonoPatchStack([np.random.rand(h, w) for _ in range(0, n)])
+        acc = PatchStack([np.random.rand(h, w, 1, 1) for _ in range(0, n)])
         self.assertEqual(acc.count, n)
         self.assertEqual(acc.hw, (h, w))
-        self.assertEqual(acc.make_tczyx().shape, (n, 1, 1, h, w))
+        self.assertEqual(acc.pyxcz.shape, (n, h, w, 1, 1))
         return acc
 
-    def test_make_patch_stack_clone(self):
-        w = 256
-        h = 512
-        n = 4
-        acc = MonoPatchStack([np.random.rand(h, w) for _ in range(0, n)])
-        self.assertEqual(MonoPatchStack(acc.data).shape, acc.shape)
 
     def test_make_patch_stack_from_file(self):
         h = monozstackmask['h']
@@ -53,30 +47,18 @@ class TestMultipositionCziImageFileAccess(unittest.TestCase):
         with self.assertRaises(FileNotFoundError):
             acc = MonoPatchStackFromFile('c:/fake/file/name.tif')
 
-    def test_patch_as_yxcz_array(self):
-        w = 256
-        h = 512
-        n = 4
-        acc = MonoPatchStack([np.random.rand(h, w) for _ in range(0, 4)])
-        self.assertEqual(acc.iat_yxcz(0).shape, (h, w, 1, 1))
-
-    def test_make_3d_patch_stack_from_list(self):
+    def test_make_3d_patch_stack_from_nonuniform_list(self):
         w = 256
         h = 512
         c = 1
         nz = 5
         n = 4
-        acc = Multichannel3dPatchStack([np.random.rand(h, w, c, nz) for _ in range(0, n)])
-        self.assertEqual(acc.count, n)
-        self.assertEqual(acc.hw, (h, w))
-        self.assertEqual(acc.chroma, c)
-        self.assertEqual(acc.iat(0).shape, (h, w, c, nz))
 
-    def test_3d_patch_as_yxcz_array(self):
-        w = 256
-        h = 512
-        nz = 5
-        c = 1
-        n = 4
-        acc = Multichannel3dPatchStack([np.random.rand(h, w, c, nz) for _ in range(0, n)])
-        self.assertEqual(acc.iat_yxcz(0).shape, (h, w, c, nz))
\ No newline at end of file
+        patches = [np.random.rand(h, w, c, nz) for _ in range(0, n)]
+        patches.append(np.random.rand(h, 2 * w, c, nz))
+        acc = PatchStack(patches)
+        self.assertEqual(acc.count, n + 1)
+        self.assertEqual(acc.hw, (h, 2 * w))
+        self.assertEqual(acc.chroma, c)
+        self.assertEqual(acc.iat(0).shape, (h, 2 * w, c, nz))
+        self.assertEqual(acc.iat_yxcz(0).shape, (h, 2 * w, c, nz))
\ No newline at end of file