Automatically convert images to RGB color space

Inference only works in RGB, so RGBA and grayscale images need to be
transformed first. Output is always RGB, regardless of input format.

Author: mdraw

deface/centerface.py

@@ -9,6 +9,15 @@
 default_onnx_path = f'{os.path.dirname(__file__)}/centerface.onnx'
 
 
+def ensure_rgb(img: np.ndarray) -> np.ndarray:
+    """Convert input image to RGB if it is in RGBA or L format"""
+    if img.ndim == 2:  # 1-channel grayscale -> RGB
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+    elif img.shape[2] == 4:  # 4-channel RGBA -> RGB
+        img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
+    return img
+
+
 class CenterFace:
     def __init__(self, onnx_path=None, in_shape=None, backend='auto'):
         self.in_shape = in_shape
@@ -71,6 +80,7 @@ def dynamicize_shapes(static_model):
         return dyn_model
 
     def __call__(self, img, threshold=0.5):
+        img = ensure_rgb(img)
         self.orig_shape = img.shape[:2]
         if self.in_shape is None:
             self.in_shape = self.orig_shape[::-1]