mode switch in hf models

Giskard-AI · Aug 8, 2024 · 6c3ffc9 · 6c3ffc9
1 parent 7ba35b6
commit 6c3ffc9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/giskard_vision/image_classification/models/base.py b/giskard_vision/image_classification/models/base.py
@@ -87,5 +87,7 @@ def predict_image(self, image, mode=None) -> Types.label:
  image (np.ndarray): input image
  mode (str): mode of the image
  """
+ if len(image.shape) == 2:
+ mode = "L" # Grayscale - Needed to run the color robustness detector
  probas = self.predict_probas(image, mode=mode)
  return self.classification_labels[np.argmax(probas)]