Resolve "ilastik models do not validate dimensionality of input data"

Closes #1 (closed)

model_server/extensions/ilastik/tests/test_ilastik.py

@@ -11,6 +11,9 @@ from model_server.base.roiset import _get_label_ids, RoiSet, RoiSetMetaParams
 from model_server.base.workflows import classify_pixels
 from tests.test_api import TestServerBaseClass
 
+def _random_int(*args):
+    return np.random.randint(0, 2 ** 8, size=args, dtype='uint8')
+
 class TestIlastikPixelClassification(unittest.TestCase):
     def setUp(self) -> None:
         self.cf = CziImageFileAccessor(czifile['path'])
@@ -83,6 +86,40 @@ class TestIlastikPixelClassification(unittest.TestCase):
         self.mono_image = mono_image
         self.mask = mask
 
+    def test_pixel_classifier_enforces_input_shape(self):
+        model = ilm.IlastikPixelClassifierModel(
+            {'project_file': ilastik_classifiers['px']}
+        )
+        self.assertEqual(model.model_chroma, 1)
+        self.assertEqual(model.model_3d, False)
+
+        # correct data
+        self.assertIsInstance(
+            model.label_pixel_class(
+                InMemoryDataAccessor(
+                    _random_int(512, 256, 1, 1)
+                )
+            ),
+            InMemoryDataAccessor
+        )
+
+        # raise except with input of multiple channels
+        with self.assertRaises(ilm.IlastikInputShapeError):
+            mask = model.label_pixel_class(
+                InMemoryDataAccessor(
+                    _random_int(512, 256, 3, 1)
+                )
+            )
+
+        # raise except with input of multiple channels
+        with self.assertRaises(ilm.IlastikInputShapeError):
+            mask = model.label_pixel_class(
+                InMemoryDataAccessor(
+                    _random_int(512, 256, 1, 15)
+                )
+            )
+
+
     def test_run_object_classifier_from_pixel_predictions(self):
         self.test_run_pixel_classifier()
         fp = czifile['path']
@@ -97,7 +134,7 @@ class TestIlastikPixelClassification(unittest.TestCase):
                 objmap,
             )
         )
-        self.assertEqual(objmap.data.max(), 3)
+        self.assertEqual(objmap.data.max(), 2)
 
     def test_run_object_classifier_from_segmentation(self):
         self.test_run_pixel_classifier()
@@ -113,7 +150,7 @@ class TestIlastikPixelClassification(unittest.TestCase):
                 objmap,
             )
         )
-        self.assertEqual(objmap.data.max(), 3)
+        self.assertEqual(objmap.data.max(), 2)
 
     def test_ilastik_pixel_classification_as_workflow(self):
         result = classify_pixels(
@@ -269,16 +306,18 @@ class TestIlastikObjectClassification(unittest.TestCase):
             )
         )
 
-        self.object_classifier = ilm.PatchStackObjectClassifier(
+        self.object_classifier = ilm.IlastikObjectClassifierFromSegmentationModel(
             params={'project_file': ilastik_classifiers['seg_to_obj']}
         )
 
+
     def test_classify_patches(self):
         raw_patches = self.roiset.get_raw_patches()
         patch_masks = self.roiset.get_patch_masks()
-        res_patches, _ = self.object_classifier.infer(raw_patches, patch_masks)
+        res_patches = self.object_classifier.label_instance_class(raw_patches, patch_masks)
         self.assertEqual(res_patches.count, self.roiset.count)
+        res_patches.export_pyxcz(output_path / 'res_patches.tif')
         for pi in range(0, res_patches.count):  # assert that there is only one nonzero label per patch
-            unique = np.unique(res_patches.iat(pi).data)
-            self.assertEqual(len(unique), 2)
-            self.assertEqual(unique[0], 0)
+            la, ct = np.unique(res_patches.iat(pi).data, return_counts=True)
+            self.assertEqual(np.sum(ct > 1), 2)  # exclude single-pixel anomaly
+            self.assertEqual(la[0], 0)