Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
plot_repurposing_annotations.py 2.41 KB
Newer Older
Allen Goodman's avatar
`masks_to_bounding_boxes` op (#4290)  
Allen Goodman committed
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 
"""
=======================
Repurposing annotations
=======================

The following example illustrates the operations available in the torchvision.ops module for repurposing object
localization annotations for different tasks (e.g. transforming masks used by instance and panoptic segmentation
methods into bounding boxes used by object detection methods).
"""
import os.path

import PIL.Image
import matplotlib.patches
import matplotlib.pyplot
import numpy
import torch
from torchvision.ops import masks_to_boxes

ASSETS_DIRECTORY = "../test/assets"

matplotlib.pyplot.rcParams["savefig.bbox"] = "tight"

####################################
# Masks
# -----
# In tasks like instance and panoptic segmentation, masks are commonly defined, and are defined by this package,
# as a multi-dimensional array (e.g. a NumPy array or a PyTorch tensor) with the following shape:
#
#       (objects, height, width)
#
# Where objects is the number of annotated objects in the image. Each (height, width) object corresponds to exactly
# one object. For example, if your input image has the dimensions 224 x 224 and has four annotated objects the shape
# of your masks annotation has the following shape:
#
#       (4, 224, 224).
#
# A nice property of masks is that they can be easily repurposed to be used in methods to solve a variety of object
# localization tasks.
#
# Masks to bounding boxes
# ----------------------------------------
# For example, the masks to bounding_boxes operation can be used to transform masks into bounding boxes that can be
# used in methods like Faster RCNN and YOLO.

with PIL.Image.open(os.path.join(ASSETS_DIRECTORY, "masks.tiff")) as image:
    masks = torch.zeros((image.n_frames, image.height, image.width), dtype=torch.int)

    for index in range(image.n_frames):
        image.seek(index)

        frame = numpy.array(image)

        masks[index] = torch.tensor(frame)

bounding_boxes = masks_to_boxes(masks)

figure = matplotlib.pyplot.figure()

a = figure.add_subplot(121)
b = figure.add_subplot(122)

labeled_image = torch.sum(masks, 0)

a.imshow(labeled_image)
b.imshow(labeled_image)

for bounding_box in bounding_boxes:
    x0, y0, x1, y1 = bounding_box

    rectangle = matplotlib.patches.Rectangle((x0, y0), x1 - x0, y1 - y0, linewidth=1, edgecolor="r", facecolor="none")

    b.add_patch(rectangle)

a.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
b.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])