File size: 5,520 Bytes
96e3a14
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
# Copyright 2025 S5 authors and Hugging Face converters.
"""Image processor for S5 remote-sensing models (ImageNet mean/std)."""

from __future__ import annotations

from typing import Optional, Union

from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from transformers.image_transforms import convert_to_rgb, 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, logging

try:
    from transformers.utils import filter_out_non_signature_kwargs
except ImportError:  # older transformers
    def filter_out_non_signature_kwargs():
        def decorator(fn):
            return fn

        return decorator


try:
    from configuration_s5 import IMAGENET_MEAN, IMAGENET_STD
except ImportError:
    from .configuration_s5 import IMAGENET_MEAN, IMAGENET_STD


logger = logging.get_logger(__name__)


class S5ImageProcessor(BaseImageProcessor):
    """Resize / rescale / ImageNet-normalize RGB remote-sensing images for S5."""

    model_input_names = ["pixel_values"]

    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,
        do_normalize: bool = True,
        image_mean: Optional[Union[float, list[float]]] = None,
        image_std: Optional[Union[float, list[float]]] = None,
        do_convert_rgb: bool = True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.do_resize = do_resize
        self.size = size if size is not None else {"height": 512, "width": 512}
        self.resample = resample
        self.do_rescale = do_rescale
        self.rescale_factor = rescale_factor
        self.do_normalize = do_normalize
        self.image_mean = list(IMAGENET_MEAN) if image_mean is None else image_mean
        self.image_std = list(IMAGENET_STD) if image_std is None else image_std
        self.do_convert_rgb = do_convert_rgb

    @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,
        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 = self.do_resize if do_resize is None else do_resize
        size = get_size_dict(self.size if size is None else size, default_to_square=True)
        resample = self.resample if resample is None else resample
        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 = self.image_mean if image_mean is None else image_mean
        image_std = 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

        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 = []
        for image in images:
            if do_convert_rgb:
                try:
                    image = convert_to_rgb(image)
                except Exception:
                    image = to_numpy_array(image)
            image = to_numpy_array(image)
            if input_data_format is None:
                try:
                    inferred = infer_channel_dimension_format(image)
                except ValueError:
                    inferred = ChannelDimension.LAST
            else:
                inferred = input_data_format
            if do_resize:
                image = resize(
                    image,
                    size=(size["height"], size["width"]),
                    resample=resample,
                    input_data_format=inferred,
                )
            if do_rescale:
                image = image * rescale_factor
            if do_normalize:
                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=inferred)
            processed.append(to_channel_dimension_format(image, data_format, input_channel_dim=inferred))

        return BatchFeature(data={"pixel_values": processed}, tensor_type=return_tensors)


__all__ = ["S5ImageProcessor"]