test_image_processing_vit.py 7.06 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# coding=utf-8
# Copyright 2021 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
22
from transformers.utils import is_torch_available, is_vision_available
23

24
from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs
25
26
27
28
29
30
31
32


if is_torch_available():
    import torch

if is_vision_available():
    from PIL import Image

33
    from transformers import ViTImageProcessor
34
35


36
class ViTImageProcessingTester(unittest.TestCase):
37
38
39
40
41
42
43
44
45
    def __init__(
        self,
        parent,
        batch_size=7,
        num_channels=3,
        image_size=18,
        min_resolution=30,
        max_resolution=400,
        do_resize=True,
amyeroberts's avatar
amyeroberts committed
46
        size=None,
NielsRogge's avatar
NielsRogge committed
47
48
49
        do_normalize=True,
        image_mean=[0.5, 0.5, 0.5],
        image_std=[0.5, 0.5, 0.5],
50
    ):
amyeroberts's avatar
amyeroberts committed
51
        size = size if size is not None else {"height": 18, "width": 18}
52
53
54
55
56
57
58
59
        self.parent = parent
        self.batch_size = batch_size
        self.num_channels = num_channels
        self.image_size = image_size
        self.min_resolution = min_resolution
        self.max_resolution = max_resolution
        self.do_resize = do_resize
        self.size = size
NielsRogge's avatar
NielsRogge committed
60
61
62
        self.do_normalize = do_normalize
        self.image_mean = image_mean
        self.image_std = image_std
63

64
    def prepare_image_processor_dict(self):
65
66
67
68
69
70
71
72
73
74
75
        return {
            "image_mean": self.image_mean,
            "image_std": self.image_std,
            "do_normalize": self.do_normalize,
            "do_resize": self.do_resize,
            "size": self.size,
        }


@require_torch
@require_vision
76
77
class ViTImageProcessingTest(ImageProcessingSavingTestMixin, unittest.TestCase):
    image_processing_class = ViTImageProcessor if is_vision_available() else None
78
79

    def setUp(self):
80
        self.image_processor_tester = ViTImageProcessingTester(self)
81
82

    @property
83
84
    def image_processor_dict(self):
        return self.image_processor_tester.prepare_image_processor_dict()
85

86
87
88
89
90
91
92
    def test_image_processor_properties(self):
        image_processing = self.image_processing_class(**self.image_processor_dict)
        self.assertTrue(hasattr(image_processing, "image_mean"))
        self.assertTrue(hasattr(image_processing, "image_std"))
        self.assertTrue(hasattr(image_processing, "do_normalize"))
        self.assertTrue(hasattr(image_processing, "do_resize"))
        self.assertTrue(hasattr(image_processing, "size"))
93

94
95
96
    def test_image_processor_from_dict_with_kwargs(self):
        image_processor = self.image_processing_class.from_dict(self.image_processor_dict)
        self.assertEqual(image_processor.size, {"height": 18, "width": 18})
97

98
99
        image_processor = self.image_processing_class.from_dict(self.image_processor_dict, size=42)
        self.assertEqual(image_processor.size, {"height": 42, "width": 42})
100

101
102
103
104
    def test_batch_feature(self):
        pass

    def test_call_pil(self):
105
106
        # Initialize image_processing
        image_processing = self.image_processing_class(**self.image_processor_dict)
107
        # create random PIL images
108
        image_inputs = prepare_image_inputs(self.image_processor_tester, equal_resolution=False)
109
110
111
112
        for image in image_inputs:
            self.assertIsInstance(image, Image.Image)

        # Test not batched input
113
        encoded_images = image_processing(image_inputs[0], return_tensors="pt").pixel_values
114
115
116
117
        self.assertEqual(
            encoded_images.shape,
            (
                1,
118
119
120
                self.image_processor_tester.num_channels,
                self.image_processor_tester.size["height"],
                self.image_processor_tester.size["width"],
121
122
123
124
            ),
        )

        # Test batched
125
        encoded_images = image_processing(image_inputs, return_tensors="pt").pixel_values
126
127
128
        self.assertEqual(
            encoded_images.shape,
            (
129
130
131
132
                self.image_processor_tester.batch_size,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.size["height"],
                self.image_processor_tester.size["width"],
133
134
135
136
            ),
        )

    def test_call_numpy(self):
137
138
        # Initialize image_processing
        image_processing = self.image_processing_class(**self.image_processor_dict)
139
        # create random numpy tensors
140
        image_inputs = prepare_image_inputs(self.image_processor_tester, equal_resolution=False, numpify=True)
141
142
143
144
        for image in image_inputs:
            self.assertIsInstance(image, np.ndarray)

        # Test not batched input
145
        encoded_images = image_processing(image_inputs[0], return_tensors="pt").pixel_values
146
147
148
149
        self.assertEqual(
            encoded_images.shape,
            (
                1,
150
151
152
                self.image_processor_tester.num_channels,
                self.image_processor_tester.size["height"],
                self.image_processor_tester.size["width"],
153
154
155
156
            ),
        )

        # Test batched
157
        encoded_images = image_processing(image_inputs, return_tensors="pt").pixel_values
158
159
160
        self.assertEqual(
            encoded_images.shape,
            (
161
162
163
164
                self.image_processor_tester.batch_size,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.size["height"],
                self.image_processor_tester.size["width"],
165
166
167
168
            ),
        )

    def test_call_pytorch(self):
169
170
        # Initialize image_processing
        image_processing = self.image_processing_class(**self.image_processor_dict)
171
        # create random PyTorch tensors
172
        image_inputs = prepare_image_inputs(self.image_processor_tester, equal_resolution=False, torchify=True)
173
174
175
176
        for image in image_inputs:
            self.assertIsInstance(image, torch.Tensor)

        # Test not batched input
177
        encoded_images = image_processing(image_inputs[0], return_tensors="pt").pixel_values
178
179
180
181
        self.assertEqual(
            encoded_images.shape,
            (
                1,
182
183
184
                self.image_processor_tester.num_channels,
                self.image_processor_tester.size["height"],
                self.image_processor_tester.size["width"],
185
186
187
188
            ),
        )

        # Test batched
189
        encoded_images = image_processing(image_inputs, return_tensors="pt").pixel_values
190
191
192
        self.assertEqual(
            encoded_images.shape,
            (
193
194
195
196
                self.image_processor_tester.batch_size,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.size["height"],
                self.image_processor_tester.size["width"],
197
198
            ),
        )