supporting gray scale

giskard_vision/core/models/base.py

@@ -18,15 +18,37 @@ class ModelBase(ABC):
     prediction_result_cls = TypesBase.prediction_result
 
     @abstractmethod
-    def predict_image(self, image: np.ndarray) -> Any:
-        """abstract method that takes one image as input and outputs the prediction
+    def predict_rgb_image(self, image: np.ndarray) -> Any:
+        """abstract method that takes one RGB image as input and outputs the prediction
 
         Args:
             image (np.ndarray): input image
         """
 
         ...
 
+    def predict_gray_image(self, image: np.ndarray) -> Any:
+        """abstract method that takes one gray image as input and outputs the prediction
+
+        Args:
+            image (np.ndarray): input image
+        """
+
+        raise NotImplementedError("predict_gray_image method is not implemented")
+
+    def predict_image(self, image: np.ndarray) -> Any:
+        """abstract method that takes one image as input and outputs the prediction
+
+        Args:
+            image (np.ndarray): input image
+        """
+        if image.shape[-1] == 3:
+            return self.predict_rgb_image(image)
+        elif image.shape[-1] == 1 or len(image.shape) == 2:
+            return self.predict_gray_image(image)
+        else:
+            raise ValueError("predict_image: image shape not supported.")
+
     def predict_batch(self, idx: int, images: List[np.ndarray]) -> np.ndarray:
         """method that should be implemented if the passed dataloader has batch_size != 1
 

giskard_vision/image_classification/models/base.py

@@ -80,14 +80,10 @@ def predict_probas(self, image: np.ndarray, mode=None) -> np.ndarray:
 
         return np.array([_prediction[label] for label in self.classification_labels])
 
-    def predict_image(self, image, mode=None) -> Types.label:
-        """method that takes one image as input and outputs one class label
+    def predict_rgb_image(self, image: np.ndarray) -> Types.label:
+        probas = self.predict_probas(image, mode=None)
+        return self.classification_labels[np.argmax(probas)]
 
-        Args:
-            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)
+    def predict_gray_image(self, image: np.ndarray) -> Types.label:
+        probas = self.predict_probas(image, mode="L")
         return self.classification_labels[np.argmax(probas)]

giskard_vision/landmark_detection/models/wrappers.py

@@ -29,7 +29,7 @@ def __init__(self, model):
         super().__init__(n_landmarks=68, n_dimensions=2, name="FaceAlignment")
         self.model = model
 
-    def predict_image(self, image):
+    def predict_rgb_image(self, image):
         """
         Predict facial landmarks for a given image using the wrapped face alignment model.
 
@@ -100,7 +100,7 @@ def __init__(self):
         self.landmark_detector = cv2.face.createFacemarkLBF()
         self.landmark_detector.loadModel(LBFmodel)
 
-    def predict_image(self, image):
+    def predict_rgb_image(self, image):
         """
         Predict facial landmarks for a given image using the wrapped OpenCV face landmarks model.
 

giskard_vision/object_detection/models/wrappers.py

@@ -96,7 +96,7 @@ def shape_rescale(self, image, boxes):
     def positive_constraint(self, boxes):
         return np.clip(boxes, 0, None)
 
-    def predict_image(self, image: np.ndarray):
+    def predict_rgb_image(self, image: np.ndarray):
         try:
             from keras.applications.mobilenet import preprocess_input
         except ImportError: