Allow to set global mode for an HF ppl model for PIL conversion

giskard_vision/core/dataloaders/hf.py

@@ -31,7 +31,11 @@ class HFDataLoader(DataIteratorBase):
     """
 
     def __init__(
-        self, hf_id: str, hf_config: Optional[str] = None, hf_split: str = "test", name: Optional[str] = None
+        self,
+        hf_id: str,
+        hf_config: Optional[str] = None,
+        hf_split: str = "test",
+        name: Optional[str] = None,
     ) -> None:
         """
         Initializes the general HuggingFace Datasets instance.

giskard_vision/core/models/hf_pipeline.py

@@ -19,9 +19,9 @@ def __init__(
         """init method that accepts a model object, number of landmarks and dimensions
 
         Args:
-            model_id (str): Hugging Face model ID
-            name (Optional[str]): name of the model
-            pipeline_task (HFPipelineTask): HuggingFace pipeline task
+            model_id (str): Hugging Face model ID.
+            name (Optional[str]): name of the model.
+            pipeline_task (HFPipelineTask): HuggingFace pipeline task.
 
         Raises:
             GiskardImportError: If there are missing Hugging Face dependencies.

giskard_vision/image_classification/models/base.py

@@ -11,9 +11,10 @@ class ImageClassificationHFModel(HFPipelineModelBase):
     """Hugging Face pipeline wrapper class that serves as a template for image classification predictions
 
     Args:
-        model_id (str): Hugging Face model ID
-        name (Optional[str]): name of the model
-        device (str): device to run the model on
+        model_id (str): Hugging Face model ID.
+        name (Optional[str]): name of the model.
+        device (str): device to run the model on.
+        mode (str): The mode to convert the numpy image data to PIL image, defaulting to "RGB".
 
     Attributes:
         classification_labels: list of classification labels, where the position of the label corresponds to the class index
@@ -22,13 +23,14 @@ class ImageClassificationHFModel(HFPipelineModelBase):
     model_type = "image_classification"
     prediction_result_cls = Types.prediction_result
 
-    def __init__(self, model_id: str, name: Optional[str] = None, device: str = "cpu"):
+    def __init__(self, model_id: str, name: Optional[str] = None, device: str = "cpu", mode: str = "RGB"):
         """init method that accepts a model id, name and device
 
         Args:
-            model_id (str): Hugging Face model ID
-            name (Optional[str]): name of the model
-            device (str): device to run the model on
+            model_id (str): Hugging Face model ID.
+            name (Optional[str]): name of the model.
+            device (str): device to run the model on.
+            mode (str): The mode to convert the numpy image data to PIL image, defaulting to "RGB".
         """
 
         super().__init__(
@@ -39,6 +41,7 @@ def __init__(self, model_id: str, name: Optional[str] = None, device: str = "cpu
         )
 
         self._classification_labels = list(self.pipeline.model.config.id2label.values())
+        self._mode = mode
 
     @property
     def classification_labels(self):
@@ -58,14 +61,15 @@ class SingleLabelImageClassificationHFModelWrapper(ImageClassificationHFModel):
         classification_labels: list of classification labels, where the position of the label corresponds to the class index
     """
 
-    def predict_probas(self, image: np.ndarray, mode="RGB") -> np.ndarray:
+    def predict_probas(self, image: np.ndarray, mode=None) -> np.ndarray:
         """method that takes one image as input and outputs the prediction of probabilities for each class
 
         Args:
             image (np.ndarray): input image
             mode (str): mode of the image
         """
-        pil_image = Image.fromarray(image, mode=mode)
+        m = mode or self._mode
+        pil_image = Image.fromarray(image, mode=m)
 
         # Pipeline takes a PIL image as input
         _raw_prediction = self.pipeline(
@@ -76,7 +80,7 @@ def predict_probas(self, image: np.ndarray, mode="RGB") -> np.ndarray:
 
         return np.array([_prediction[label] for label in self.classification_labels])
 
-    def predict_image(self, image, mode="RGB") -> Types.label:
+    def predict_image(self, image, mode=None) -> Types.label:
         """method that takes one image as input and outputs one class label
 
         Args: