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']