test_pipelines_image_segmentation.py 21.5 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# 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 hashlib
import unittest
17
from typing import Dict
18

19
import datasets
20
import numpy as np
21
from datasets import load_dataset
22

23
24
from transformers import (
    MODEL_FOR_IMAGE_SEGMENTATION_MAPPING,
25
    MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING,
26
    MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING,
27
28
    AutoFeatureExtractor,
    AutoModelForImageSegmentation,
29
    AutoModelForInstanceSegmentation,
30
    DetrForSegmentation,
31
    ImageSegmentationPipeline,
32
    MaskFormerForInstanceSegmentation,
33
34
35
    is_vision_available,
    pipeline,
)
36
from transformers.testing_utils import nested_simplify, require_tf, require_timm, require_torch, require_vision, slow
37
38
39
40
41
42
43
44
45
46
47
48
49
50

from .test_pipelines_common import ANY, PipelineTestCaseMeta


if is_vision_available():
    from PIL import Image
else:

    class Image:
        @staticmethod
        def open(*args, **kwargs):
            pass


51
52
def hashimage(image: Image) -> str:
    m = hashlib.md5(image.tobytes())
53
54
55
56
57
58
59
60
    return m.hexdigest()[:10]


def mask_to_test_readable(mask: Image) -> Dict:
    npimg = np.array(mask)
    white_pixels = (npimg == 255).sum()
    shape = npimg.shape
    return {"hash": hashimage(mask), "white_pixels": white_pixels, "shape": shape}
61
62


63
64
65
66
@require_vision
@require_timm
@require_torch
class ImageSegmentationPipelineTests(unittest.TestCase, metaclass=PipelineTestCaseMeta):
67
68
69
70
71
72
    model_mapping = {
        k: v
        for k, v in (
            list(MODEL_FOR_IMAGE_SEGMENTATION_MAPPING.items()) if MODEL_FOR_IMAGE_SEGMENTATION_MAPPING else []
        )
        + (MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING.items() if MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING else [])
73
        + (MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING.items() if MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING else [])
74
    }
75

76
    def get_test_pipeline(self, model, tokenizer, feature_extractor):
77
        image_segmenter = ImageSegmentationPipeline(model=model, feature_extractor=feature_extractor)
78
79
80
81
82
83
        return image_segmenter, [
            "./tests/fixtures/tests_samples/COCO/000000039769.png",
            "./tests/fixtures/tests_samples/COCO/000000039769.png",
        ]

    def run_pipeline_test(self, image_segmenter, examples):
84
        outputs = image_segmenter("./tests/fixtures/tests_samples/COCO/000000039769.png", threshold=0.0)
85
86
        self.assertIsInstance(outputs, list)
        n = len(outputs)
87
88
89
90
91
92
        if isinstance(image_segmenter.model, (MaskFormerForInstanceSegmentation)):
            # Instance segmentation (maskformer) have a slot for null class
            # and can output nothing even with a low threshold
            self.assertGreaterEqual(n, 0)
        else:
            self.assertGreaterEqual(n, 1)
93
94
95
        # XXX: PIL.Image implements __eq__ which bypasses ANY, so we inverse the comparison
        # to make it work
        self.assertEqual([{"score": ANY(float, type(None)), "label": ANY(str), "mask": ANY(Image.Image)}] * n, outputs)
96

97
        dataset = datasets.load_dataset("hf-internal-testing/fixtures_image_utils", "image", split="test")
98

99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
        # RGBA
        outputs = image_segmenter(dataset[0]["file"])
        m = len(outputs)
        self.assertEqual([{"score": ANY(float, type(None)), "label": ANY(str), "mask": ANY(Image.Image)}] * m, outputs)
        # LA
        outputs = image_segmenter(dataset[1]["file"])
        m = len(outputs)
        self.assertEqual([{"score": ANY(float, type(None)), "label": ANY(str), "mask": ANY(Image.Image)}] * m, outputs)
        # L
        outputs = image_segmenter(dataset[2]["file"])
        m = len(outputs)
        self.assertEqual([{"score": ANY(float, type(None)), "label": ANY(str), "mask": ANY(Image.Image)}] * m, outputs)

        if isinstance(image_segmenter.model, DetrForSegmentation):
            # We need to test batch_size with images with the same size.
            # Detr doesn't normalize the size of the images, meaning we can have
            # 800x800 or 800x1200, meaning we cannot batch simply.
            # We simply bail on this
            batch_size = 1
        else:
            batch_size = 2

        # 5 times the same image so the output shape is predictable
122
        batch = [
123
124
125
126
127
            "./tests/fixtures/tests_samples/COCO/000000039769.png",
            "./tests/fixtures/tests_samples/COCO/000000039769.png",
            "./tests/fixtures/tests_samples/COCO/000000039769.png",
            "./tests/fixtures/tests_samples/COCO/000000039769.png",
            "./tests/fixtures/tests_samples/COCO/000000039769.png",
128
        ]
129
        outputs = image_segmenter(batch, threshold=0.0, batch_size=batch_size)
130
        self.assertEqual(len(batch), len(outputs))
131
        self.assertEqual(len(outputs[0]), n)
132
133
        self.assertEqual(
            [
134
135
136
137
138
                [{"score": ANY(float, type(None)), "label": ANY(str), "mask": ANY(Image.Image)}] * n,
                [{"score": ANY(float, type(None)), "label": ANY(str), "mask": ANY(Image.Image)}] * n,
                [{"score": ANY(float, type(None)), "label": ANY(str), "mask": ANY(Image.Image)}] * n,
                [{"score": ANY(float, type(None)), "label": ANY(str), "mask": ANY(Image.Image)}] * n,
                [{"score": ANY(float, type(None)), "label": ANY(str), "mask": ANY(Image.Image)}] * n,
139
            ],
140
141
            outputs,
            f"Expected [{n}, {n}, {n}, {n}, {n}], got {[len(item) for item in outputs]}",
142
143
144
145
146
147
148
149
150
        )

    @require_tf
    @unittest.skip("Image segmentation not implemented in TF")
    def test_small_model_tf(self):
        pass

    @require_torch
    def test_small_model_pt(self):
151
        model_id = "hf-internal-testing/tiny-detr-mobilenetsv3-panoptic"
152
153
154

        model = AutoModelForImageSegmentation.from_pretrained(model_id)
        feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
155
156
157
158
159
160
161
        image_segmenter = ImageSegmentationPipeline(
            model=model,
            feature_extractor=feature_extractor,
            task="semantic",
            threshold=0.0,
            overlap_mask_area_threshold=0.0,
        )
162

163
164
165
166
167
        outputs = image_segmenter(
            "http://images.cocodataset.org/val2017/000000039769.jpg",
        )

        # Shortening by hashing
168
        for o in outputs:
169
            o["mask"] = mask_to_test_readable(o["mask"])
170

171
        # This is extremely brittle, and those values are made specific for the CI.
172
173
174
175
        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
                {
176
177
178
179
180
181
182
183
                    "label": "LABEL_88",
                    "mask": {"hash": "7f0bf661a4", "shape": (480, 640), "white_pixels": 3},
                    "score": None,
                },
                {
                    "label": "LABEL_101",
                    "mask": {"hash": "10ab738dc9", "shape": (480, 640), "white_pixels": 8948},
                    "score": None,
184
185
                },
                {
186
                    "label": "LABEL_215",
187
188
                    "mask": {"hash": "b431e0946c", "shape": (480, 640), "white_pixels": 298249},
                    "score": None,
189
                },
190
191
192
193
194
195
196
197
198
199
200
201
202
203
            ]
            # Temporary: Keeping around the old values as they might provide useful later
            # [
            #     {
            #         "score": 0.004,
            #         "label": "LABEL_215",
            #         "mask": {"hash": "34eecd16bb", "shape": (480, 640), "white_pixels": 0},
            #     },
            #     {
            #         "score": 0.004,
            #         "label": "LABEL_215",
            #         "mask": {"hash": "34eecd16bb", "shape": (480, 640), "white_pixels": 0},
            #     },
            # ],
204
205
206
207
208
209
210
211
212
213
214
        )

        outputs = image_segmenter(
            [
                "http://images.cocodataset.org/val2017/000000039769.jpg",
                "http://images.cocodataset.org/val2017/000000039769.jpg",
            ],
            threshold=0.0,
        )
        for output in outputs:
            for o in output:
215
                o["mask"] = mask_to_test_readable(o["mask"])
216
217
218
219
220
221

        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
                [
                    {
222
223
224
225
226
227
228
229
                        "label": "LABEL_88",
                        "mask": {"hash": "7f0bf661a4", "shape": (480, 640), "white_pixels": 3},
                        "score": None,
                    },
                    {
                        "label": "LABEL_101",
                        "mask": {"hash": "10ab738dc9", "shape": (480, 640), "white_pixels": 8948},
                        "score": None,
230
231
                    },
                    {
232
                        "label": "LABEL_215",
233
234
                        "mask": {"hash": "b431e0946c", "shape": (480, 640), "white_pixels": 298249},
                        "score": None,
235
236
237
238
                    },
                ],
                [
                    {
239
240
241
242
243
244
245
246
                        "label": "LABEL_88",
                        "mask": {"hash": "7f0bf661a4", "shape": (480, 640), "white_pixels": 3},
                        "score": None,
                    },
                    {
                        "label": "LABEL_101",
                        "mask": {"hash": "10ab738dc9", "shape": (480, 640), "white_pixels": 8948},
                        "score": None,
247
248
                    },
                    {
249
                        "label": "LABEL_215",
250
251
                        "mask": {"hash": "b431e0946c", "shape": (480, 640), "white_pixels": 298249},
                        "score": None,
252
                    },
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
                ]
                # [
                #     {
                #         "score": 0.004,
                #         "label": "LABEL_215",
                #         "mask": {"hash": "34eecd16bb", "shape": (480, 640), "white_pixels": 0},
                #     },
                #     {
                #         "score": 0.004,
                #         "label": "LABEL_215",
                #         "mask": {"hash": "34eecd16bb", "shape": (480, 640), "white_pixels": 0},
                #     },
                # ],
                # [
                #     {
                #         "score": 0.004,
                #         "label": "LABEL_215",
                #         "mask": {"hash": "34eecd16bb", "shape": (480, 640), "white_pixels": 0},
                #     },
                #     {
                #         "score": 0.004,
                #         "label": "LABEL_215",
                #         "mask": {"hash": "34eecd16bb", "shape": (480, 640), "white_pixels": 0},
                #     },
                # ],
278
279
280
            ],
        )

281
282
283
284
285
286
287
    @require_torch
    def test_small_model_pt_semantic(self):
        model_id = "hf-internal-testing/tiny-random-beit-pipeline"
        image_segmenter = pipeline(model=model_id)
        outputs = image_segmenter("http://images.cocodataset.org/val2017/000000039769.jpg")
        for o in outputs:
            # shortening by hashing
288
            o["mask"] = mask_to_test_readable(o["mask"])
289
290
291
292

        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
293
294
295
296
297
                {
                    "score": None,
                    "label": "LABEL_0",
                    "mask": {"hash": "42d0907228", "shape": (480, 640), "white_pixels": 10714},
                },
298
299
300
                {
                    "score": None,
                    "label": "LABEL_1",
301
                    "mask": {"hash": "46b8cc3976", "shape": (480, 640), "white_pixels": 296486},
302
303
304
305
                },
            ],
        )

306
307
308
309
310
311
    @require_torch
    @slow
    def test_integration_torch_image_segmentation(self):
        model_id = "facebook/detr-resnet-50-panoptic"
        image_segmenter = pipeline("image-segmentation", model=model_id)

312
313
314
315
316
317
318
319
        outputs = image_segmenter(
            "http://images.cocodataset.org/val2017/000000039769.jpg",
            task="panoptic",
            threshold=0,
            overlap_mask_area_threshold=0.0,
        )

        # Shortening by hashing
320
        for o in outputs:
321
            o["mask"] = mask_to_test_readable(o["mask"])
322
323
324
325

        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
                {
                    "score": 0.9094,
                    "label": "blanket",
                    "mask": {"hash": "dcff19a97a", "shape": (480, 640), "white_pixels": 16617},
                },
                {
                    "score": 0.9941,
                    "label": "cat",
                    "mask": {"hash": "9c0af87bd0", "shape": (480, 640), "white_pixels": 59185},
                },
                {
                    "score": 0.9987,
                    "label": "remote",
                    "mask": {"hash": "c7870600d6", "shape": (480, 640), "white_pixels": 4182},
                },
                {
                    "score": 0.9995,
                    "label": "remote",
                    "mask": {"hash": "ef899a25fd", "shape": (480, 640), "white_pixels": 2275},
                },
                {
                    "score": 0.9722,
                    "label": "couch",
                    "mask": {"hash": "37b8446ac5", "shape": (480, 640), "white_pixels": 172380},
                },
                {
                    "score": 0.9994,
                    "label": "cat",
                    "mask": {"hash": "6a09d3655e", "shape": (480, 640), "white_pixels": 52561},
                },
356
357
358
359
360
361
362
363
            ],
        )

        outputs = image_segmenter(
            [
                "http://images.cocodataset.org/val2017/000000039769.jpg",
                "http://images.cocodataset.org/val2017/000000039769.jpg",
            ],
364
            task="panoptic",
365
            threshold=0.0,
366
            overlap_mask_area_threshold=0.0,
367
        )
368
369

        # Shortening by hashing
370
371
        for output in outputs:
            for o in output:
372
                o["mask"] = mask_to_test_readable(o["mask"])
373
374
375
376
377

        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
                [
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
                    {
                        "score": 0.9094,
                        "label": "blanket",
                        "mask": {"hash": "dcff19a97a", "shape": (480, 640), "white_pixels": 16617},
                    },
                    {
                        "score": 0.9941,
                        "label": "cat",
                        "mask": {"hash": "9c0af87bd0", "shape": (480, 640), "white_pixels": 59185},
                    },
                    {
                        "score": 0.9987,
                        "label": "remote",
                        "mask": {"hash": "c7870600d6", "shape": (480, 640), "white_pixels": 4182},
                    },
                    {
                        "score": 0.9995,
                        "label": "remote",
                        "mask": {"hash": "ef899a25fd", "shape": (480, 640), "white_pixels": 2275},
                    },
                    {
                        "score": 0.9722,
                        "label": "couch",
                        "mask": {"hash": "37b8446ac5", "shape": (480, 640), "white_pixels": 172380},
                    },
                    {
                        "score": 0.9994,
                        "label": "cat",
                        "mask": {"hash": "6a09d3655e", "shape": (480, 640), "white_pixels": 52561},
                    },
408
409
                ],
                [
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
                    {
                        "score": 0.9094,
                        "label": "blanket",
                        "mask": {"hash": "dcff19a97a", "shape": (480, 640), "white_pixels": 16617},
                    },
                    {
                        "score": 0.9941,
                        "label": "cat",
                        "mask": {"hash": "9c0af87bd0", "shape": (480, 640), "white_pixels": 59185},
                    },
                    {
                        "score": 0.9987,
                        "label": "remote",
                        "mask": {"hash": "c7870600d6", "shape": (480, 640), "white_pixels": 4182},
                    },
                    {
                        "score": 0.9995,
                        "label": "remote",
                        "mask": {"hash": "ef899a25fd", "shape": (480, 640), "white_pixels": 2275},
                    },
                    {
                        "score": 0.9722,
                        "label": "couch",
                        "mask": {"hash": "37b8446ac5", "shape": (480, 640), "white_pixels": 172380},
                    },
                    {
                        "score": 0.9994,
                        "label": "cat",
                        "mask": {"hash": "6a09d3655e", "shape": (480, 640), "white_pixels": 52561},
                    },
440
441
442
443
444
445
446
447
448
449
                ],
            ],
        )

    @require_torch
    @slow
    def test_threshold(self):
        model_id = "facebook/detr-resnet-50-panoptic"
        image_segmenter = pipeline("image-segmentation", model=model_id)

450
451
452
453
454
        outputs = image_segmenter(
            "http://images.cocodataset.org/val2017/000000039769.jpg", task="panoptic", threshold=0.999
        )
        # Shortening by hashing
        for o in outputs:
455
            o["mask"] = mask_to_test_readable(o["mask"])
456
457
458
459

        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
460
461
462
463
464
465
466
467
468
469
                {
                    "score": 0.9995,
                    "label": "remote",
                    "mask": {"hash": "d02404f578", "shape": (480, 640), "white_pixels": 2789},
                },
                {
                    "score": 0.9994,
                    "label": "cat",
                    "mask": {"hash": "eaa115b40c", "shape": (480, 640), "white_pixels": 304411},
                },
470
471
472
473
474
475
            ],
        )

        outputs = image_segmenter(
            "http://images.cocodataset.org/val2017/000000039769.jpg", task="panoptic", threshold=0.5
        )
476
477

        for o in outputs:
478
            o["mask"] = mask_to_test_readable(o["mask"])
479
480
481
482

        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
                {
                    "score": 0.9941,
                    "label": "cat",
                    "mask": {"hash": "9c0af87bd0", "shape": (480, 640), "white_pixels": 59185},
                },
                {
                    "score": 0.9987,
                    "label": "remote",
                    "mask": {"hash": "c7870600d6", "shape": (480, 640), "white_pixels": 4182},
                },
                {
                    "score": 0.9995,
                    "label": "remote",
                    "mask": {"hash": "ef899a25fd", "shape": (480, 640), "white_pixels": 2275},
                },
                {
                    "score": 0.9722,
                    "label": "couch",
                    "mask": {"hash": "37b8446ac5", "shape": (480, 640), "white_pixels": 172380},
                },
                {
                    "score": 0.9994,
                    "label": "cat",
                    "mask": {"hash": "6a09d3655e", "shape": (480, 640), "white_pixels": 52561},
                },
508
509
            ],
        )
510
511
512
513

    @require_torch
    @slow
    def test_maskformer(self):
514
        threshold = 0.8
515
516
        model_id = "facebook/maskformer-swin-base-ade"

517
518
        model = AutoModelForInstanceSegmentation.from_pretrained(model_id)
        feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
519
520
521
522

        image_segmenter = pipeline("image-segmentation", model=model, feature_extractor=feature_extractor)

        image = load_dataset("hf-internal-testing/fixtures_ade20k", split="test")
523
        file = image[0]["file"]
524
        outputs = image_segmenter(file, task="panoptic", threshold=threshold)
525

526
        # Shortening by hashing
527
        for o in outputs:
528
            o["mask"] = mask_to_test_readable(o["mask"])
529
530
531
532

        self.assertEqual(
            nested_simplify(outputs, decimals=4),
            [
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
                {
                    "score": 0.9974,
                    "label": "wall",
                    "mask": {"hash": "a547b7c062", "shape": (512, 683), "white_pixels": 14252},
                },
                {
                    "score": 0.949,
                    "label": "house",
                    "mask": {"hash": "0da9b7b38f", "shape": (512, 683), "white_pixels": 132177},
                },
                {
                    "score": 0.9995,
                    "label": "grass",
                    "mask": {"hash": "1d07ea0a26", "shape": (512, 683), "white_pixels": 53444},
                },
                {
                    "score": 0.9976,
                    "label": "tree",
                    "mask": {"hash": "6cdc97c7da", "shape": (512, 683), "white_pixels": 7944},
                },
                {
                    "score": 0.8239,
                    "label": "plant",
                    "mask": {"hash": "1ab4ce378f", "shape": (512, 683), "white_pixels": 4136},
                },
                {
                    "score": 0.9942,
                    "label": "road, route",
                    "mask": {"hash": "39c5d17be5", "shape": (512, 683), "white_pixels": 1941},
                },
                {
                    "score": 1.0,
                    "label": "sky",
                    "mask": {"hash": "a3756324a6", "shape": (512, 683), "white_pixels": 135802},
                },
568
569
            ],
        )