from __future__ import annotations import json from pathlib import Path from typing import Any, Dict, Iterable, List, Optional, Tuple, Union import numpy as np from PIL import Image, ImageOps import torch from transformers.image_processing_base import ImageProcessingMixin def _as_list(val): if isinstance(val, (list, tuple)): return list(val) return [val] def _to_numpy(image: Any) -> np.ndarray: if isinstance(image, np.ndarray): return image if torch.is_tensor(image): return image.detach().cpu().numpy() if isinstance(image, Image.Image): return np.array(image) raise TypeError(f"Unsupported image type: {type(image)}") def _infer_input_format(arr: np.ndarray) -> str: if arr.ndim == 3 and arr.shape[0] in (1, 3) and arr.shape[-1] not in (1, 3): return "channels_first" return "channels_last" def _to_channels_last(arr: np.ndarray, input_format: str) -> np.ndarray: if input_format == "channels_first": return np.transpose(arr, (1, 2, 0)) return arr def _to_channels_first(arr: np.ndarray, input_format: str) -> np.ndarray: if input_format == "channels_last": return np.transpose(arr, (2, 0, 1)) return arr def _compute_resize_size(image_size: Tuple[int, int], size: Dict[str, int]) -> Tuple[int, int]: height, width = image_size if "height" in size and "width" in size: return int(size["height"]), int(size["width"]) if "shortest_edge" in size: target = int(size["shortest_edge"]) if height <= width: new_h = target new_w = int(round(width * target / max(height, 1))) else: new_w = target new_h = int(round(height * target / max(width, 1))) return new_h, new_w raise ValueError(f"Unsupported size dict: {size}") def _resolve_resample(resample: Optional[int]) -> int: if resample is None: return Image.BICUBIC try: return Image.Resampling(resample) except Exception: return resample def _center_crop_pil(image: Image.Image, crop_size: Dict[str, int]) -> Image.Image: target_h = int(crop_size["height"]) target_w = int(crop_size["width"]) width, height = image.size if width < target_w or height < target_h: padded_w = max(width, target_w) padded_h = max(height, target_h) padded = Image.new(image.mode, (padded_w, padded_h), (0, 0, 0)) padded.paste(image, ((padded_w - width) // 2, (padded_h - height) // 2)) image = padded width, height = image.size left = int(round((width - target_w) / 2.0)) top = int(round((height - target_h) / 2.0)) return image.crop((left, top, left + target_w, top + target_h)) def _normalize_return_tensors(value: Optional[Union[str, Any]]) -> Optional[str]: if value is None: return None if isinstance(value, str): return value.lower() name = getattr(value, "name", None) if name: return name.lower() return str(value).lower() class Florence2ImageProcessorLite(ImageProcessingMixin): model_input_names = ["pixel_values"] def __init__( self, image_seq_length: int, do_resize: bool = True, size: Optional[Dict[str, int]] = None, resample: Optional[int] = None, do_center_crop: bool = False, crop_size: Optional[Dict[str, int]] = None, do_rescale: bool = True, rescale_factor: float = 1 / 255, do_normalize: bool = True, image_mean: Optional[List[float]] = None, image_std: Optional[List[float]] = None, do_convert_rgb: Optional[bool] = True, ) -> None: super().__init__() self.image_seq_length = int(image_seq_length) self.do_resize = bool(do_resize) self.size = size or {"height": 224, "width": 224} self.resample = resample self.do_center_crop = bool(do_center_crop) self.crop_size = crop_size or {"height": 224, "width": 224} self.do_rescale = bool(do_rescale) self.rescale_factor = float(rescale_factor) self.do_normalize = bool(do_normalize) self.image_mean = image_mean or [0.485, 0.456, 0.406] self.image_std = image_std or [0.229, 0.224, 0.225] self.do_convert_rgb = do_convert_rgb @classmethod def from_preprocessor_config(cls, model_dir: Union[str, Path]) -> "Florence2ImageProcessorLite": config_path = Path(model_dir) / "preprocessor_config.json" if not config_path.exists(): raise FileNotFoundError(f"Missing Florence2 preprocessor_config.json in {model_dir}") data = json.loads(config_path.read_text(encoding="utf-8")) return cls( image_seq_length=data.get("image_seq_length", 0), do_resize=data.get("do_resize", True), size=data.get("size") or data.get("crop_size") or {"height": 224, "width": 224}, resample=data.get("resample"), do_center_crop=data.get("do_center_crop", False), crop_size=data.get("crop_size") or data.get("size") or {"height": 224, "width": 224}, do_rescale=data.get("do_rescale", True), rescale_factor=data.get("rescale_factor", 1 / 255), do_normalize=data.get("do_normalize", True), image_mean=data.get("image_mean"), image_std=data.get("image_std"), do_convert_rgb=data.get("do_convert_rgb"), ) def __call__( self, images: Union[Image.Image, np.ndarray, torch.Tensor, List[Any]], do_resize: Optional[bool] = None, size: Optional[Dict[str, int]] = None, resample: Optional[int] = None, do_center_crop: Optional[bool] = None, crop_size: Optional[Dict[str, int]] = None, do_rescale: Optional[bool] = None, rescale_factor: Optional[float] = None, do_normalize: Optional[bool] = None, image_mean: Optional[Iterable[float]] = None, image_std: Optional[Iterable[float]] = None, do_convert_rgb: Optional[bool] = None, return_tensors: Optional[Union[str, Any]] = "pt", data_format: Optional[str] = "channels_first", input_data_format: Optional[str] = None, **kwargs, ) -> Dict[str, Any]: do_resize = self.do_resize if do_resize is None else do_resize size = self.size if size is None else size resample = self.resample if resample is None else resample do_center_crop = self.do_center_crop if do_center_crop is None else do_center_crop crop_size = self.crop_size if crop_size is None else crop_size do_rescale = self.do_rescale if do_rescale is None else do_rescale rescale_factor = self.rescale_factor if rescale_factor is None else rescale_factor do_normalize = self.do_normalize if do_normalize is None else do_normalize image_mean = list(self.image_mean if image_mean is None else image_mean) image_std = list(self.image_std if image_std is None else image_std) do_convert_rgb = self.do_convert_rgb if do_convert_rgb is None else do_convert_rgb resample = _resolve_resample(resample) want_torch = _normalize_return_tensors(return_tensors) in ("pt", "pytorch", "tensortype.pytorch") processed: List[np.ndarray] = [] for image in _as_list(images): if isinstance(image, Image.Image): img = image if do_convert_rgb: img = ImageOps.exif_transpose(img).convert("RGB") else: arr = _to_numpy(image) input_fmt = input_data_format or _infer_input_format(arr) arr = _to_channels_last(arr, input_fmt) img = Image.fromarray(arr.astype(np.uint8)) if do_convert_rgb: img = img.convert("RGB") if do_resize: out_h, out_w = _compute_resize_size((img.size[1], img.size[0]), size) img = img.resize((out_w, out_h), resample=resample) if do_center_crop: img = _center_crop_pil(img, crop_size) arr = np.array(img).astype(np.float32) if do_rescale: arr = arr * float(rescale_factor) if do_normalize: mean = np.array(image_mean, dtype=np.float32) std = np.array(image_std, dtype=np.float32) arr = (arr - mean) / std if data_format in ("channels_first", "first"): arr = _to_channels_first(arr, "channels_last") processed.append(arr) batch = np.stack(processed, axis=0) if want_torch: batch = torch.from_numpy(batch).float() return {"pixel_values": batch}