Implement labels property for object classification models

model_server/extensions/ilastik/models.py

@@ -85,6 +85,14 @@ class IlastikModel(Model):
                 dd[ci] = 1
         return dd
 
+    @property
+    def labels(self):
+        return []
+
+    @property
+    def info(self):
+        return {**self._info, 'labels': self.labels}
+
     @property
     def model_chroma(self):
         return self.model_shape_dict['C']
@@ -183,6 +191,11 @@ class IlastikObjectClassifierFromMaskSegmentationModel(IlastikModel, InstanceMas
         from ilastik.workflows.objectClassification.objectClassificationWorkflow import ObjectClassificationWorkflowBinary
         return ObjectClassificationWorkflowBinary
 
+    @property
+    def labels(self):
+        h5 = self.shell.projectManager.currentProjectFile
+        return [None] + [l.decode() for l in h5['ObjectClassification/LabelNames'][()]]
+
     def infer(self, input_img: GenericImageDataAccessor, segmentation_img: GenericImageDataAccessor) -> (np.ndarray, dict):
         if self.model_chroma != input_img.chroma:
             raise IlastikInputShapeError(
@@ -248,6 +261,11 @@ class IlastikObjectClassifierFromPixelPredictionsModel(IlastikModel, ImageToImag
         from ilastik.workflows.objectClassification.objectClassificationWorkflow import ObjectClassificationWorkflowPrediction
         return ObjectClassificationWorkflowPrediction
 
+    @property
+    def labels(self):
+        h5 = self.shell.projectManager.currentProjectFile
+        return [None] + [l.decode() for l in h5['ObjectClassification/LabelNames'][()]]
+
     def infer(self, input_img: GenericImageDataAccessor, pxmap_img: GenericImageDataAccessor) -> (np.ndarray, dict):
         if self.model_chroma != input_img.chroma:
             raise IlastikInputShapeError(