""" ALOE image processor — standard ViT-style preprocessing with B-cos channel expansion. B-cos models expect 6-channel input ``[x, 1-x]`` (on the **normalised** 3-channel tensor ``x``), matching ``HFTransform`` in ``src/data/transforms/transforms.py`` when ``is_bcos=True``. Using ``-x`` for the second half was incorrect for training parity and broke native-Hub k-NN / linear-probe eval. This processor wraps resize → normalise → to-tensor and appends the second half, so the pipeline matches ModelFactory + dataloader code paths. The processor is serialisable, loadable with ``AutoImageProcessor.from_pretrained``, and can be pushed to the HF Hub alongside the model. Typical usage ------------- :: from transformers import AutoImageProcessor proc = AutoImageProcessor.from_pretrained("your-org/aloe-siglip2-base", trust_remote_code=True) inputs = proc(images=pil_image, return_tensors="pt") # inputs.pixel_values.shape == (1, 6, H, W) out = model(**inputs) Building from a backbone processor (during export) --------------------------------------------------- :: >>> from src.models.hf.aloe import AloeImageProcessor >>> backbone_proc = AutoImageProcessor.from_pretrained("google/siglip2-base-patch16-224") >>> aloe_proc = AloeImageProcessor.from_backbone_processor(backbone_proc) """ from __future__ import annotations from typing import Any, Dict, List, Optional, Union import numpy as np from transformers import ViTImageProcessor from transformers.image_processing_utils import BatchFeature _IMAGENET_MEAN = [0.485, 0.456, 0.406] _IMAGENET_STD = [0.229, 0.224, 0.225] class AloeImageProcessor(ViTImageProcessor): """ Image processor for all ALOE vision backbones. Extends :class:`~transformers.ViTImageProcessor` with an optional B-cos 6-channel expansion: the normalised 3-channel tensor ``x`` is expanded to ``torch.cat([x, 1.0 - x], dim=channel)``, same as ``HFTransform`` for B-cos training. Hub checkpoints produced by :func:`src.models.hf.aloe_models.export_native_model_to_hub` patch ``preprocessor_config.json`` with ``image_processor_type`` and ``auto_map["AutoImageProcessor"]`` so :class:`transformers.AutoImageProcessor` resolves the class when using ``trust_remote_code=True``. Parameters ---------- expand_to_6ch : bool If ``True`` (default), concatenate ``[pixel_values, 1 - pixel_values]`` along the channel dimension after preprocessing (on the normalised tensor). Set to ``False`` for standard 3-channel output (e.g. visualisation). All other parameters are forwarded to :class:`~transformers.ViTImageProcessor` (``do_resize``, ``size``, ``do_normalize``, ``image_mean``, ``image_std``, …). """ model_input_names = ["pixel_values"] def to_dict(self) -> dict[str, Any]: """Serialise with explicit type and B-cos flag (see :meth:`~transformers.image_processing_utils.ImageProcessingMixin.to_dict`).""" out = super().to_dict() out["image_processor_type"] = self.__class__.__name__ out["expand_to_6ch"] = self.expand_to_6ch return out def __init__( self, expand_to_6ch: bool = True, do_resize: bool = True, 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: bool = True, **kwargs: Any, ) -> None: self.expand_to_6ch = expand_to_6ch super().__init__( do_resize=do_resize, size=size or {"height": 224, "width": 224}, do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean if image_mean is not None else _IMAGENET_MEAN, image_std=image_std if image_std is not None else _IMAGENET_STD, do_convert_rgb=do_convert_rgb, **kwargs, ) def preprocess( self, images: Union[Any, List[Any]], *, return_tensors: Optional[str] = None, **kwargs: Any, ) -> BatchFeature: """ Preprocess one or more images. Delegates to :class:`~transformers.ViTImageProcessor` for all standard steps, then — if ``expand_to_6ch=True`` — appends ``(1 - pv)`` (same convention as ``HFTransform`` with ``is_bcos=True``). """ result = super().preprocess(images, return_tensors=return_tensors, **kwargs) if self.expand_to_6ch and "pixel_values" in result: pv = result["pixel_values"] if isinstance(pv, np.ndarray): result["pixel_values"] = np.concatenate([pv, 1.0 - pv], axis=1) else: import torch result["pixel_values"] = torch.cat([pv, 1.0 - pv], dim=1) return result # ------------------------------------------------------------------ # Factory helpers # ------------------------------------------------------------------ @classmethod def from_backbone_processor( cls, backbone_processor: Any, *, expand_to_6ch: bool = True, ) -> "AloeImageProcessor": """ Build an :class:`AloeImageProcessor` by copying preprocessing parameters from an existing backbone processor (e.g. the one loaded by ``AutoImageProcessor.from_pretrained("google/siglip2-base-…")``). Parameters ---------- backbone_processor : Any HF image processor with ``image_mean``, ``image_std``, and (optionally) a ``size`` or ``crop_size`` attribute. expand_to_6ch : Whether to append the 6-channel expansion step (default ``True``). """ mean = getattr(backbone_processor, "image_mean", _IMAGENET_MEAN) std = getattr(backbone_processor, "image_std", _IMAGENET_STD) # Resolve image size from various naming conventions used across HF processors. size: Dict[str, int] = {"height": 224, "width": 224} for attr in ("size", "crop_size"): raw = getattr(backbone_processor, attr, None) if isinstance(raw, dict) and "height" in raw and "width" in raw: size = {"height": raw["height"], "width": raw["width"]} break if isinstance(raw, int): size = {"height": raw, "width": raw} break do_rescale = getattr(backbone_processor, "do_rescale", True) rescale_factor = getattr(backbone_processor, "rescale_factor", 1 / 255) do_normalize = getattr(backbone_processor, "do_normalize", True) do_resize = getattr(backbone_processor, "do_resize", True) do_convert_rgb = getattr(backbone_processor, "do_convert_rgb", True) return cls( expand_to_6ch=expand_to_6ch, do_resize=do_resize, size=size, do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=list(mean), image_std=list(std), do_convert_rgb=do_convert_rgb, )