PatchStack no longer loses slice information when doing channel separation

model_server/base/accessors.py

@@ -486,6 +486,10 @@ class PatchStack(InMemoryDataAccessor):
     def shape_dict(self):
         return dict(zip(('P', 'Y', 'X', 'C', 'Z'), self.data.shape))
 
+    def get_channels(self, channels: list, mip: bool = False):
+        nda = super().get_channels(channels, mip).data
+        return self._derived_accessor([nda[i][self._slices[i]] for i in range(0, self.count)])
+
     def get_list(self, crop=True):
         if crop:
             return [self._data[i][self._slices[i]] for i in range(0, self.count)]

tests/base/test_accessors.py

@@ -272,7 +272,7 @@ class TestPatchStackAccessor(unittest.TestCase):
     def test_make_3d_patch_stack_from_nonuniform_list(self):
         w = 256
         h = 512
-        c = 1
+        c = 3
         nz = 5
         n = 4
 
@@ -296,6 +296,11 @@ class TestPatchStackAccessor(unittest.TestCase):
         ps_list_uncropped = acc.get_list(crop=False)
         self.assertTrue(all([p.shape == acc.shape[1:] for p in ps_list_uncropped]))
 
+        # test that this persists after channel selection
+        for i in range(0, acc.count):
+            self.assertEqual(patches[i].shape[0:2], acc.get_channels([0]).iat(i, crop=True).shape[0:2])
+            self.assertEqual(patches[i].shape[3], acc.get_channels([0]).iat(i, crop=True).shape[3])
+
     def test_make_3d_patch_stack_from_list_force_long_dim(self):
         def _r(h, w):
             return np.random.randint(0, 2 ** 8, size=(h, w, 1, 1), dtype='uint8')