# DreamSim (HuggingFace format) — unofficial port.
# Copyright (c) 2026 bigshanedogg. Released under the MIT License (see LICENSE).
#
# Derivative of DreamSim (MIT, (c) 2023 Shobhita Sundaram, Netanel Tamir,
# Stephanie Fu, Richard Zhang — https://github.com/ssundaram21/dreamsim).
# Not an official DreamSim release.

"""HF image processor for DreamSim.

Reproduces the upstream ``dreamsim`` preprocess exactly: resize to
``img_size × img_size`` with BICUBIC and scale to ``[0, 1]`` — NO mean/std
normalization (the backbones consume [0,1] tensors; DreamSim's own
mean/L2-normalization happens inside the model on the output embedding).
"""

from typing import Any, List, Optional, Union

import numpy as np
import PIL.Image
import torch
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature


class DreamSimImageProcessor(BaseImageProcessor):
    model_input_names = ["pixel_values"]

    def __init__(self, img_size: int = 224, **kwargs):
        super().__init__(**kwargs)
        self.img_size = img_size

    def _to_tensor(self, image: PIL.Image.Image) -> torch.Tensor:
        # BICUBIC resize to (img_size, img_size), then HWC uint8 → CHW float [0,1].
        image = image.convert("RGB").resize((self.img_size, self.img_size), PIL.Image.BICUBIC)
        _array = np.asarray(image, dtype=np.float32) / 255.0
        return torch.from_numpy(_array).permute(2, 0, 1).contiguous()

    def preprocess(
        self,
        images: Union[PIL.Image.Image, List[PIL.Image.Image]],
        return_tensors: Optional[str] = "pt",
        **kwargs: Any,
    ) -> BatchFeature:
        if isinstance(images, PIL.Image.Image):
            images = [images]
        _pixel_values = torch.stack([self._to_tensor(_image) for _image in images], dim=0)
        return BatchFeature(data={"pixel_values": _pixel_values}, tensor_type=return_tensors)