# retinaface/image_processing_retinaface.py

from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
from transformers.image_utils import ImageInput
import numpy as np
from typing import List

class RetinaFaceImageProcessor(BaseImageProcessor):
    model_input_names = ["pixel_values"]
    image_processor_type = "RetinaFaceImageProcessor" # クラス名と同じにするのが一般的

    def __init__(self, mean: List[int] = None, std: int = 1, do_normalize: bool = True, **kwargs):
        super().__init__(**kwargs)
        # 元の実装の正規化: img -= (104, 117, 123)
        self.mean = mean if mean is not None else [104, 117, 123]
        self.std = std
        self.do_normalize = do_normalize

    def preprocess(self, images: ImageInput, return_tensors: str = "pt", **kwargs):
        """
        画像を前処理します。
        """
        if not isinstance(images, list):
            images = [images]

        processed_images = []
        for image in images:
            # Numpy配列に変換
            if not isinstance(image, np.ndarray):
                image = self.to_numpy_array(image)
            
            # BGR -> BGR (元のモデルはOpenCVのBGRを想定)
            
            # 正規化
            if self.do_normalize:
                image = self.normalize(image, mean=self.mean, std=self.std)

            # HWC to CHW
            image = image.transpose(2, 0, 1)
            processed_images.append(image)

        return BatchFeature(data={"pixel_values": processed_images}, return_tensors=return_tensors)

    def normalize(self, image, mean, std, **kwargs):
        """画像からmeanを引き、stdで割ります。"""
        image = image.astype(np.float32)
        image -= np.array(mean).astype(image.dtype)
        image /= std
        return image