src/inference/preprocessing.py

"""Frame preprocessing shared by every prediction entry point."""
import cv2
import numpy as np
from PIL import Image
from tensorflow.keras.applications.efficientnet_v2 import preprocess_input
from tensorflow.keras.utils import img_to_array

from src.configs import settings


def to_uint8_array(frame) -> np.ndarray:
    if isinstance(frame, Image.Image):
        frame = np.array(frame)
    if not isinstance(frame, np.ndarray):
        frame = np.array(frame)
    if frame.dtype != np.uint8:
        frame = frame.astype(np.uint8)
    return frame


def preprocess_frame(frame) -> np.ndarray:
    frame = to_uint8_array(frame)
    resized = cv2.resize(frame, settings.IMAGE_SIZE)
    array = img_to_array(resized)
    array = preprocess_input(array)
    return np.expand_dims(array, axis=0)