From 6c3ffc9d8247461653093e2bf5db69b99330c6c5 Mon Sep 17 00:00:00 2001 From: Rabah Khalek Date: Thu, 8 Aug 2024 16:23:32 +0200 Subject: [PATCH] mode switch in hf models --- giskard_vision/image_classification/models/base.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/giskard_vision/image_classification/models/base.py b/giskard_vision/image_classification/models/base.py index 92c7bbfe..57cd82f6 100644 --- 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)]