dofa-base-patch16-224/image_processing_dofa.py

# Copyright 2024 The DOFA Authors and The HuggingFace Inc. team.
"""Image processor for DOFA."""

from typing import Optional, Union

import numpy as np

from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from transformers.image_transforms import resize, to_channel_dimension_format
from transformers.image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    infer_channel_dimension_format,
    make_flat_list_of_images,
    to_numpy_array,
    valid_images,
    validate_preprocess_arguments,
)
from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging


logger = logging.get_logger(__name__)


def _resize_multispectral(image: np.ndarray, size: dict[str, int], input_data_format: ChannelDimension) -> np.ndarray:
    target_height, target_width = size["height"], size["width"]

    if input_data_format == ChannelDimension.FIRST:
        image = np.transpose(image, (1, 2, 0))

    height, width, _ = image.shape
    if height == target_height and width == target_width:
        resized = image
    else:
        try:
            import cv2
        except ImportError as exc:
            raise ImportError(
                "Multispectral resize requires OpenCV (`opencv-python`) when input has more than 4 channels."
            ) from exc
        resized = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR)

    if input_data_format == ChannelDimension.FIRST:
        return np.transpose(resized, (2, 0, 1))
    return resized


class DOFAImageProcessor(BaseImageProcessor):
    model_input_names = ["pixel_values", "wavelengths"]

    def __init__(
        self,
        do_resize: bool = False,
        size: Optional[dict[str, int]] = None,
        resample: PILImageResampling = PILImageResampling.BILINEAR,
        do_rescale: bool = True,
        rescale_factor: float = 1 / 255.0,
        do_normalize: bool = True,
        image_mean: Optional[Union[float, list[float]]] = None,
        image_std: Optional[Union[float, list[float]]] = None,
        do_convert_rgb: bool = False,
        default_wavelengths: Optional[list[float]] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        size = size if size is not None else {"height": 224, "width": 224}
        self.do_resize = do_resize
        self.size = size
        self.resample = resample
        self.do_rescale = do_rescale
        self.rescale_factor = rescale_factor
        self.do_normalize = do_normalize
        self.image_mean = image_mean
        self.image_std = image_std
        self.do_convert_rgb = do_convert_rgb
        self.default_wavelengths = default_wavelengths

    @filter_out_non_signature_kwargs()
    def preprocess(
        self,
        images: ImageInput,
        do_resize: Optional[bool] = None,
        size: Optional[dict[str, int]] = None,
        resample: Optional[PILImageResampling] = None,
        do_rescale: Optional[bool] = None,
        rescale_factor: Optional[float] = None,
        do_normalize: Optional[bool] = None,
        image_mean: Optional[Union[float, list[float]]] = None,
        image_std: Optional[Union[float, list[float]]] = None,
        wavelengths: Optional[list[float]] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        do_convert_rgb: Optional[bool] = None,
    ):
        do_resize = do_resize if do_resize is not None else self.do_resize
        size = size if size is not None else self.size
        size = get_size_dict(size, default_to_square=True)
        resample = resample if resample is not None else self.resample
        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
        image_mean = image_mean if image_mean is not None else self.image_mean
        image_std = image_std if image_std is not None else self.image_std
        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
        wavelengths = wavelengths if wavelengths is not None else self.default_wavelengths

        if wavelengths is None:
            raise ValueError(
                "DOFA requires `wavelengths` (one value per channel, in micrometers). "
                "Pass them to the processor call or set `default_wavelengths` in the processor config."
            )

        if do_normalize and (image_mean is None or image_std is None):
            raise ValueError(
                "Normalization requires `image_mean` and `image_std` with one value per channel. "
                "Provide modality-specific statistics for your input bands."
            )

        if do_normalize and (
            len(image_mean) != len(wavelengths) or len(image_std) != len(wavelengths)
        ):
            raise ValueError(
                f"`image_mean` and `image_std` must match the number of wavelengths/channels "
                f"({len(wavelengths)}), got mean={len(image_mean)} std={len(image_std)}."
            )

        images = make_flat_list_of_images(images)
        if not valid_images(images):
            raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")

        validate_preprocess_arguments(
            do_rescale=do_rescale,
            rescale_factor=rescale_factor,
            do_normalize=do_normalize,
            image_mean=image_mean,
            image_std=image_std,
            do_resize=do_resize,
            size=size,
            resample=resample,
        )

        processed_images = []
        for image in images:
            image = to_numpy_array(image)
            if do_convert_rgb:
                image = self._convert_image_to_rgb(image)

            if input_data_format is None:
                try:
                    input_data_format = infer_channel_dimension_format(image)
                except ValueError:
                    input_data_format = ChannelDimension.LAST

            num_channels = image.shape[-1] if input_data_format == ChannelDimension.LAST else image.shape[0]

            if do_resize:
                if num_channels > 4:
                    image = _resize_multispectral(image, size=size, input_data_format=input_data_format)
                else:
                    image = resize(image, size=size, resample=resample, input_data_format=input_data_format)

            if do_rescale:
                image = image * rescale_factor

            if do_normalize:
                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)

            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
            processed_images.append(image)

        data = {
            "pixel_values": processed_images,
            "wavelengths": np.asarray(wavelengths, dtype=np.float32),
        }
        return BatchFeature(data=data, tensor_type=return_tensors)


__all__ = ["DOFAImageProcessor"]