IlastikPixelClassifierModel can optionally rescale pxmap output to range of its input

model_server/extensions/ilastik/models.py

@@ -152,9 +152,12 @@ class IlastikPixelClassifierModel(IlastikModel, SemanticSegmentationModel):
         )
         return InMemoryDataAccessor(data=yxcz)
 
-    def infer_patch_stack(self, img: PatchStack, crop=True, **kwargs) -> (np.ndarray, dict):
+    def infer_patch_stack(self, img: PatchStack, crop=True, normalize=False, **kwargs) -> (np.ndarray, dict):
         """
         Iterative over a patch stack, call inference separately on each cropped patch
+        :param img: patch stack of input data
+        :param crop: pass list of cropped (generally non-uniform) patches to classifier if True
+        :param normalize: scale the inference result (generally 0.0 to 1.0) to the range of img if True
         """
         dsi = [
             {
@@ -165,7 +168,12 @@ class IlastikPixelClassifierModel(IlastikModel, SemanticSegmentationModel):
         ]
         pxmaps = self.shell.workflow.batchProcessingApplet.run_export(dsi, export_to_array=True)  # [z x h x w x n]
         yxcz = [np.moveaxis(pm, [1, 2, 3, 0], [0, 1, 2, 3]) for pm in pxmaps]
-        return PatchStack(yxcz)
+        if normalize:
+            return PatchStack(yxcz).apply(
+                lambda x: (x * img.dtype_max).astype(img.dtype)
+            )
+        else:
+            return PatchStack(yxcz)
 
     def label_pixel_class(self, img: GenericImageDataAccessor, **kwargs):
         pxmap = self.infer(img)

tests/test_ilastik/test_ilastik.py

@@ -150,6 +150,9 @@ class TestIlastikPixelClassification(unittest.TestCase):
         self.assertEqual(pxmap.nz, acc.nz)
         self.assertEqual(pxmap.count, acc.count)
 
+        norm_pxmap = self.model.infer_patch_stack(acc, normalize=True)
+        self.assertEqual(norm_pxmap.dtype, 'uint8')
+
     def test_run_object_classifier_from_pixel_predictions(self):
         self.test_run_pixel_classifier()
         fp = czifile['path']