-
Notifications
You must be signed in to change notification settings - Fork 184
/
Copy pathpanoptic2detection_coco_format.py
153 lines (134 loc) · 6.27 KB
/
panoptic2detection_coco_format.py
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
145
146
147
148
149
150
151
152
153
#!/usr/bin/env python
'''
This script converts panoptic COCO format to detection COCO format. More
information about the formats can be found here:
http://cocodataset.org/#format-data. All segments will be stored in RLE format.
Additional option:
- using option '--things_only' the script can discard all stuff
segments, saving segments of things classes only.
'''
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import os, sys
import argparse
import numpy as np
import json
import time
import multiprocessing
import PIL.Image as Image
from panopticapi.utils import get_traceback, rgb2id, save_json
try:
# set up path for pycocotools
# sys.path.append('./cocoapi-master/PythonAPI/')
from pycocotools import mask as COCOmask
except Exception:
raise Exception("Please install pycocotools module from https://github.com/cocodataset/cocoapi")
@get_traceback
def convert_panoptic_to_detection_coco_format_single_core(
proc_id, annotations_set, categories, segmentations_folder, things_only
):
annotations_detection = []
for working_idx, annotation in enumerate(annotations_set):
if working_idx % 100 == 0:
print('Core: {}, {} from {} images processed'.format(proc_id,
working_idx,
len(annotations_set)))
file_name = '{}.png'.format(annotation['file_name'].rsplit('.')[0])
try:
pan_format = np.array(
Image.open(os.path.join(segmentations_folder, file_name)), dtype=np.uint32
)
except IOError:
raise KeyError('no prediction png file for id: {}'.format(annotation['image_id']))
pan = rgb2id(pan_format)
for segm_info in annotation['segments_info']:
if things_only and categories[segm_info['category_id']]['isthing'] != 1:
continue
mask = (pan == segm_info['id']).astype(np.uint8)
mask = np.expand_dims(mask, axis=2)
segm_info.pop('id')
segm_info['image_id'] = annotation['image_id']
rle = COCOmask.encode(np.asfortranarray(mask))[0]
rle['counts'] = rle['counts'].decode('utf8')
segm_info['segmentation'] = rle
annotations_detection.append(segm_info)
print('Core: {}, all {} images processed'.format(proc_id, len(annotations_set)))
return annotations_detection
def convert_panoptic_to_detection_coco_format(input_json_file,
segmentations_folder,
output_json_file,
categories_json_file,
things_only):
start_time = time.time()
if segmentations_folder is None:
segmentations_folder = input_json_file.rsplit('.', 1)[0]
print("CONVERTING...")
print("COCO panoptic format:")
print("\tSegmentation folder: {}".format(segmentations_folder))
print("\tJSON file: {}".format(input_json_file))
print("TO")
print("COCO detection format")
print("\tJSON file: {}".format(output_json_file))
if things_only:
print("Saving only segments of things classes.")
print('\n')
print("Reading annotation information from {}".format(input_json_file))
with open(input_json_file, 'r') as f:
d_coco = json.load(f)
annotations_panoptic = d_coco['annotations']
with open(categories_json_file, 'r') as f:
categories_list = json.load(f)
categories = {category['id']: category for category in categories_list}
cpu_num = multiprocessing.cpu_count()
annotations_split = np.array_split(annotations_panoptic, cpu_num)
print("Number of cores: {}, images per core: {}".format(cpu_num, len(annotations_split[0])))
workers = multiprocessing.Pool(processes=cpu_num)
processes = []
for proc_id, annotations_set in enumerate(annotations_split):
p = workers.apply_async(convert_panoptic_to_detection_coco_format_single_core,
(proc_id, annotations_set, categories, segmentations_folder, things_only))
processes.append(p)
annotations_coco_detection = []
for p in processes:
annotations_coco_detection.extend(p.get())
for idx, ann in enumerate(annotations_coco_detection):
ann['id'] = idx
d_coco['annotations'] = annotations_coco_detection
categories_coco_detection = []
for category in d_coco['categories']:
if things_only and category['isthing'] != 1:
continue
category.pop('isthing')
category.pop('color')
categories_coco_detection.append(category)
d_coco['categories'] = categories_coco_detection
save_json(d_coco, output_json_file)
t_delta = time.time() - start_time
print("Time elapsed: {:0.2f} seconds".format(t_delta))
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="The script converts panoptic COCO format to detection \
COCO format. See this file's head for more information."
)
parser.add_argument('--input_json_file', type=str,
help="JSON file with panoptic COCO format")
parser.add_argument(
'--segmentations_folder', type=str, default=None, help="Folder with \
panoptic COCO format segmentations. Default: X if input_json_file is \
X.json"
)
parser.add_argument('--output_json_file', type=str,
help="JSON file with detection COCO format")
parser.add_argument('--categories_json_file', type=str,
help="JSON file with Panoptic COCO categories information",
default='./panoptic_coco_categories.json')
parser.add_argument('--things_only', action='store_true',
help="discard stuff classes")
args = parser.parse_args()
convert_panoptic_to_detection_coco_format(args.input_json_file,
args.segmentations_folder,
args.output_json_file,
args.categories_json_file,
args.things_only)