Skip to content

Commit

Permalink
Automatically convert images to RGB color space
Browse files Browse the repository at this point in the history
Inference only works in RGB, so RGBA and grayscale images need to be
transformed first. Output is always RGB, regardless of input format.
  • Loading branch information
mdraw committed Jan 10, 2023
1 parent edfe190 commit 852286e
Showing 1 changed file with 10 additions and 0 deletions.
10 changes: 10 additions & 0 deletions deface/centerface.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand Down

0 comments on commit 852286e

Please sign in to comment.