File size: 1,876 Bytes
f918a65 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | # 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)
|