Function for converting segmentation polygons in MS-COCO
object detection dataset to segmentation masks. The seg-
mentation masks are stored with a colour palette that's
randomly assigned based on class if specified. Change
the seed if you want to change colours.
Parameters:
| Name |
Type |
Description |
Default |
cann |
dict
|
COCO annotation object |
required
|
output_dir |
str
|
Directory to store segmentation masks |
required
|
palette |
bool
|
Use palette. Defaults to True. |
True
|
Raises:
| Type |
Description |
AssertionError
|
Incorrect shape for colour map array |
AssertionError
|
Muliple images have same image id |
Source code in coco_assistant/utils/det2seg.py
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 | def det2seg(cann, output_dir, palette=True):
"""
Function for converting segmentation polygons in MS-COCO
object detection dataset to segmentation masks. The seg-
mentation masks are stored with a colour palette that's
randomly assigned based on class if specified. Change
the seed if you want to change colours.
Args:
cann (dict): COCO annotation object
output_dir (str): Directory to store segmentation masks
palette (bool, optional): Use palette. Defaults to True.
Raises:
AssertionError: Incorrect shape for colour map array
AssertionError: Muliple images have same image id
"""
output_dir = Path(output_dir)
if not output_dir.is_dir():
output_dir.mkdir(parents=True, exist_ok=True)
imids = cann.getImgIds()
cats = cann.loadCats(cann.getCatIds())
cat_colours = {0: (0, 0, 0)}
# Set seed for palette colour
np.random.seed(121)
# Create category colourmap
for c in cats:
cat_colours[c["id"]] = (
np.random.randint(0, 256),
np.random.randint(0, 256),
np.random.randint(0, 256),
)
colour_map = np.array(list(cat_colours.values()))
if colour_map.shape != (len(cats) + 1, 3):
raise AssertionError("Incorrect shape of color map array")
for imid in tqdm(imids):
img = cann.loadImgs(imid)
if len(img) > 1:
raise AssertionError("Multiple images with same id")
h, w = img[0]["height"], img[0]["width"]
name = img[0]["file_name"]
name = Path(name)
if name.suffix.lower() != ".png":
name = name.stem + ".png"
im = np.zeros((h, w), dtype=np.uint8)
annids = cann.getAnnIds(imgIds=[imid])
if not annids:
# No annotations
res = Image.fromarray(im)
else:
anns = cann.loadAnns(annids)
areas = [i["area"] for i in anns]
area_ids = [i for i in range(1, len(areas) + 1)][::-1]
area_id_map = dict(zip(sorted(areas), area_ids))
area_cat_map = {}
# Assumption: area of objects are unique
for ann in anns:
aid = area_id_map[ann["area"]]
bMask = cann.annToMask(ann)
aMask = bMask * aid
im = np.maximum(im, aMask)
area_cat_map[aid] = ann["category_id"]
# Ref: https://stackoverflow.com/questions/55949809/efficiently-replace-elements-in-array-based-on-dictionary-numpy-python/55950051#55950051
k = np.array(list(area_cat_map.keys()))
v = np.array(list(area_cat_map.values()))
mapping_ar = np.zeros(k.max() + 1, dtype=np.uint8)
mapping_ar[k] = v
res = mapping_ar[im]
res = Image.fromarray(res)
if palette:
res.putpalette(colour_map.astype(np.uint8))
res.save(output_dir / f"{name}")
|