Normalise image inputs and force RGB channels

app.py

@@ -17,6 +17,17 @@ logger = logging.getLogger(__name__)
 IMG_SIZE = 244
 
 
+def _ensure_three_channels(array: np.ndarray) -> np.ndarray:
+    if array.ndim == 2:
+        array = np.stack([array] * 3, axis=-1)
+    elif array.ndim == 3:
+        if array.shape[-1] == 1:
+            array = np.repeat(array, 3, axis=-1)
+        elif array.shape[-1] > 3:
+            array = array[..., :3]
+    return array
+
+
 class FaceShapeModel:
     def __init__(self, model_path: str, labels: List[str]):
         if not os.path.exists(model_path):
@@ -33,7 +44,9 @@ class FaceShapeModel:
             image = image.convert("RGB")
 
         resized = image.resize((IMG_SIZE, IMG_SIZE), Image.BILINEAR)
-        array = tf.keras.preprocessing.image.img_to_array(resized)
+        array = np.asarray(resized, dtype="float32")
+        array = _ensure_three_channels(array)
+        array /= 255.0
         array = np.expand_dims(array, axis=0)
         return array
 

inference.py

@@ -17,6 +17,7 @@ LABELS: List[str] = [
 
 TARGET_SIZE = 244
 
+
 def _load_image(image_bytes: bytes) -> Image.Image:
     image = Image.open(io.BytesIO(image_bytes))
     if image.mode != "RGB":
@@ -24,10 +25,22 @@ def _load_image(image_bytes: bytes) -> Image.Image:
     return image
 
 
+def _ensure_three_channels(array: np.ndarray) -> np.ndarray:
+    if array.ndim == 2:
+        array = np.stack([array] * 3, axis=-1)
+    elif array.ndim == 3:
+        if array.shape[-1] == 1:
+            array = np.repeat(array, 3, axis=-1)
+        elif array.shape[-1] > 3:
+            array = array[..., :3]
+    return array
+
+
 def _preprocess(image_bytes: bytes) -> np.ndarray:
     image = _load_image(image_bytes)
     resized = image.resize((TARGET_SIZE, TARGET_SIZE), Image.BILINEAR)
     array = np.asarray(resized, dtype="float32")
+    array = _ensure_three_channels(array)
     array /= 255.0
     return np.expand_dims(array, axis=0)
 
@@ -37,23 +50,16 @@ class PreTrainedModel:
         self.model = tf.keras.models.load_model(model_path)
 
     def predict(self, inputs: bytes) -> List[Dict[str, Any]]:
-        batch = _preprocess(inputs)
-        preds = self.model.predict(batch, verbose=0)
+        x = _preprocess(inputs)
+        preds = self.model.predict(x, verbose=0)
 
         if isinstance(preds, (list, tuple)):
             preds = preds[0]
 
-        scores = np.asarray(preds).squeeze()
-        if scores.ndim == 0:
-            scores = np.array([float(scores)])
-
-        scores_list = scores.tolist()
-        top_index = int(np.argmax(scores_list))
+        probs = np.asarray(preds).squeeze().tolist()
+        idx = int(np.argmax(probs))
         return [
-            {
-                "label": LABELS[top_index],
-                "score": float(scores_list[top_index]),
-            }
+            {"label": LABELS[idx], "score": float(probs[idx])},
         ]