| import json |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
| from transformers.feature_extraction_utils import BatchFeature |
| from transformers.utils import PushToHubMixin |
|
|
|
|
| class VisionCondenserProcessor(PushToHubMixin): |
| """Resize images to a fixed square tensor batch.""" |
|
|
| _auto_class = "AutoProcessor" |
| model_input_names = ["pixel_values"] |
|
|
| def __init__( |
| self, |
| image_size: int = 224, |
| **kwargs, |
| ) -> None: |
| kwargs.pop("patch_size", None) |
| kwargs.pop("pool_size", None) |
| kwargs.pop("max_patches", None) |
| self.image_size = int(image_size) |
| self.extra_config = dict(kwargs) |
| if self.image_size <= 0: |
| raise ValueError(f"image_size must be positive, got {image_size}.") |
|
|
| @classmethod |
| def register_for_auto_class(cls, auto_class: str = "AutoProcessor"): |
| cls._auto_class = auto_class |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): |
| config_path = Path(pretrained_model_name_or_path) / "preprocessor_config.json" |
| with config_path.open("r", encoding="utf-8") as f: |
| config = json.load(f) |
|
|
| config.update(kwargs) |
| config.pop("processor_class", None) |
| config.pop("auto_map", None) |
| return cls(**config) |
|
|
| def save_pretrained(self, save_directory, **kwargs): |
| del kwargs |
| save_path = Path(save_directory) |
| save_path.mkdir(parents=True, exist_ok=True) |
|
|
| config = { |
| "processor_class": self.__class__.__name__, |
| "image_size": self.image_size, |
| **self.extra_config, |
| "auto_map": { |
| "AutoProcessor": "processing_vision_condenser.VisionCondenserProcessor" |
| }, |
| } |
|
|
| config_path = save_path / "preprocessor_config.json" |
| with config_path.open("w", encoding="utf-8") as f: |
| json.dump(config, f, indent=2, sort_keys=True) |
| f.write("\n") |
|
|
| source_file = Path(__file__) |
| target_file = save_path / source_file.name |
| if not target_file.exists(): |
| target_file.write_bytes(source_file.read_bytes()) |
|
|
| return (str(config_path),) |
|
|
| def _ensure_pil_image(self, image: Any) -> Image.Image: |
| if isinstance(image, Image.Image): |
| return image |
| return Image.fromarray(np.asarray(image)) |
|
|
| def _normalize_images(self, images: Any) -> list[Image.Image]: |
| if isinstance(images, Image.Image): |
| return [images] |
| if not isinstance(images, (list, tuple)) or len(images) == 0: |
| raise ValueError( |
| "images must be a PIL image or a non-empty list of PIL images." |
| ) |
| return [self._ensure_pil_image(image) for image in images] |
|
|
| def _image_to_tensor(self, image: Image.Image) -> torch.Tensor: |
| image = image.convert("RGB") |
| image = image.resize( |
| (self.image_size, self.image_size), |
| Image.Resampling.BICUBIC, |
| ) |
| pixels = torch.from_numpy(np.array(image)).permute(2, 0, 1).contiguous() |
| return pixels.float().div(255.0) |
|
|
| def __call__( |
| self, |
| images, |
| return_tensors: str | None = None, |
| **kwargs, |
| ) -> BatchFeature: |
| del kwargs |
| if return_tensors not in (None, "pt"): |
| raise ValueError("Only return_tensors='pt' is supported.") |
|
|
| pil_images = self._normalize_images(images) |
| pixel_values = torch.stack( |
| [self._image_to_tensor(image) for image in pil_images], |
| dim=0, |
| ).contiguous() |
|
|
| data = {"pixel_values": pixel_values} |
| tensor_type = "pt" if return_tensors == "pt" else None |
| return BatchFeature(data=data, tensor_type=tensor_type) |
|
|
