Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

update bdd2coco for detection and tracking #86

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open

update bdd2coco for detection and tracking #86

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
fc9cc22
update bdd2coco for detection and tracking
haofengac May 27, 2020
8b7574b
update detection
haofengac Jun 8, 2020
a0f7258
reformat code
xinw1012 Sep 21, 2020
4d5fb2b
Merge branch 'master' into bdd2coco
xinw1012 Sep 21, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
234 changes: 164 additions & 70 deletions bdd100k/bdd2coco.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,30 +2,57 @@
import json
import argparse
from tqdm import tqdm
from collections import defaultdict


def parse_arguments():
parser = argparse.ArgumentParser(description='BDD100K to COCO format')
parser.add_argument(
"-l", "--label_dir",
default="/path/to/bdd/label/",
help="root directory of BDD label Json files",
"-i", "--in_path",
default="/input/path/",
help="path to detection JSON file or tracking base folder",
)
parser.add_argument(
"-s", "--save_path",
default="/save/path",
help="path to save coco formatted label file",
"-o", "--out_path",
default="/output/path",
help="path to save coco formatted label file",
)
parser.add_argument(
"-m", "--mode",
default="det",
choices=["det", "track"],
help="conversion mode: detection or tracking",
)
return parser.parse_args()


def bdd2coco_detection(id_dict, labeled_images, fn):
def bdd2coco_det(labels):

images = list()
annotations = list()
# naming replacement for legacy release
naming_replacement_dict = {
'person': 'pedestrian',
'motor': 'motorcycle',
'bike': 'bicycle'
}

coco = defaultdict(list)
coco["categories"] = [
{"supercategory": "none", "id": 1, "name": "pedestrian"},
{"supercategory": "none", "id": 2, "name": "rider"},
{"supercategory": "none", "id": 3, "name": "car"},
{"supercategory": "none", "id": 4, "name": "truck"},
{"supercategory": "none", "id": 5, "name": "bus"},
{"supercategory": "none", "id": 6, "name": "train"},
{"supercategory": "none", "id": 7, "name": "motorcycle"},
{"supercategory": "none", "id": 8, "name": "bicycle"},
{"supercategory": "none", "id": 9, "name": "traffic light"},
{"supercategory": "none", "id": 10, "name": "traffic sign"},
{"supercategory": "none", "id": 11, "name": "ignored"}
]
attr_id_dict = {i['name']: i['id'] for i in coco['categories']}

counter = 0
for i in tqdm(labeled_images):
for i in tqdm(labels):
counter += 1
image = dict()
image['file_name'] = i['name']
Expand All @@ -34,79 +61,146 @@ def bdd2coco_detection(id_dict, labeled_images, fn):

image['id'] = counter

empty_image = True

for label in i['labels']:
annotation = dict()
if label['category'] in id_dict.keys():
empty_image = False
annotation["iscrowd"] = 0
if i['labels']:
for l in i['labels']:
# skip for drivable area and lane marking
if not 'box2d' in l:
continue
annotation = dict()
annotation["iscrowd"] = int(
l['attributes']['crowd']) if 'crowd' in l['attributes'] else 0
annotation["image_id"] = image['id']
x1 = label['box2d']['x1']
y1 = label['box2d']['y1']
x2 = label['box2d']['x2']
y2 = label['box2d']['y2']

x1 = l['box2d']['x1']
y1 = l['box2d']['y1']
x2 = l['box2d']['x2']
y2 = l['box2d']['y2']

annotation['bbox'] = [x1, y1, x2-x1, y2-y1]
annotation['area'] = float((x2 - x1) * (y2 - y1))
annotation['category_id'] = id_dict[label['category']]
annotation['ignore'] = 0
annotation['id'] = label['id']
# fix legacy naming
if l['category'] in naming_replacement_dict:
l['category'] = naming_replacement_dict[l['category']]
category_ignored = l['category'] not in attr_id_dict

annotation['category_id'] = attr_id_dict['ignored'] if category_ignored else attr_id_dict[l['category']]
annotation['ignore'] = int(category_ignored)
annotation['id'] = l['id']
annotation['segmentation'] = [[x1, y1, x1, y2, x2, y2, x2, y1]]
annotations.append(annotation)

if empty_image:
coco['annotations'].append(annotation)
else:
continue

images.append(image)

attr_dict["images"] = images
attr_dict["annotations"] = annotations
attr_dict["type"] = "instances"
coco['images'].append(image)
coco['type'] = 'instances'

print('saving...')
json_string = json.dumps(attr_dict)
with open(fn, "w") as file:
file.write(json_string)
return coco


if __name__ == '__main__':

args = parse_arguments()
def bdd2coco_track(labels):

attr_dict = dict()
attr_dict["categories"] = [
{"supercategory": "none", "id": 1, "name": "person"},
coco = defaultdict(list)
coco["categories"] = [
{"supercategory": "none", "id": 1, "name": "pedestrian"},
{"supercategory": "none", "id": 2, "name": "rider"},
{"supercategory": "none", "id": 3, "name": "car"},
{"supercategory": "none", "id": 4, "name": "bus"},
{"supercategory": "none", "id": 5, "name": "truck"},
{"supercategory": "none", "id": 6, "name": "bike"},
{"supercategory": "none", "id": 7, "name": "motor"},
{"supercategory": "none", "id": 8, "name": "traffic light"},
{"supercategory": "none", "id": 9, "name": "traffic sign"},
{"supercategory": "none", "id": 10, "name": "train"}
{"supercategory": "none", "id": 4, "name": "truck"},
{"supercategory": "none", "id": 5, "name": "bus"},
{"supercategory": "none", "id": 6, "name": "train"},
{"supercategory": "none", "id": 7, "name": "motorcycle"},
{"supercategory": "none", "id": 8, "name": "bicycle"},
{"supercategory": "none", "id": 9, "name": "ignored"}
]
attr_id_dict = {i['name']: i['id'] for i in coco['categories']}

video_id, image_id, ann_id, global_instance_id = 0, 0, 0, 0
no_ann = 0

for video_index, video_anns in enumerate(tqdm(labels)):
instance_id_maps = dict()

# videos
video = dict(id=video_id, name=video_anns[0]['video_name'])
coco['videos'].append(video)

# images
for image_anns in video_anns:
video_name = '-'.join(image_anns['name'].split('-')[:-1])
image = dict(
file_name=os.path.join(video_name, image_anns['name']),
height=720,
width=1280,
id=image_id,
video_id=video_id,
index=image_anns['index'])
coco['images'].append(image)

# annotations
for lbl in image_anns['labels']:
category_ignored = False
if lbl['category'] not in attr_id_dict.keys():
lbl['category'] = 'ignored'
category_ignored = True

bdd_id = lbl['id']
if bdd_id in instance_id_maps.keys():
instance_id = instance_id_maps[bdd_id]
else:
instance_id = global_instance_id
global_instance_id += 1
instance_id_maps[bdd_id] = instance_id

x1 = lbl['box2d']['x1']
x2 = lbl['box2d']['x2']
y1 = lbl['box2d']['y1']
y2 = lbl['box2d']['y2']
area = float((x2 - x1) * (y2 - y1))
ann = dict(
id=ann_id,
image_id=image_id,
category_id=attr_id_dict[lbl['category']],
instance_id=instance_id,
is_occluded=lbl['attributes']['Occluded'],
is_truncated=lbl['attributes']['Truncated'],
bbox=[x1, y1, x2 - x1, y2 - y1],
area=area,
iscrowd=int(lbl['attributes']['Crowd']),
ignore=int(category_ignored),
segmentation=[[x1, y1, x1, y2, x2, y2, x2, y1]])

coco['annotations'].append(ann)
ann_id += 1
if len(image_anns['labels']) == 0:
no_ann += 1

image_id += 1
video_id += 1

return coco

attr_id_dict = {i['name']: i['id'] for i in attr_dict['categories']}

if __name__ == '__main__':

args = parse_arguments()

# create BDD training set detections in COCO format
print('Loading training set...')
with open(os.path.join(args.label_dir,
'bdd100k_labels_images_train.json')) as f:
train_labels = json.load(f)
print('Converting training set...')

out_fn = os.path.join(args.save_path,
'bdd100k_labels_images_det_coco_train.json')
bdd2coco_detection(attr_id_dict, train_labels, out_fn)

print('Loading validation set...')
# create BDD validation set detections in COCO format
with open(os.path.join(args.label_dir,
'bdd100k_labels_images_val.json')) as f:
val_labels = json.load(f)
print('Converting validation set...')

out_fn = os.path.join(args.save_path,
'bdd100k_labels_images_det_coco_val.json')
bdd2coco_detection(attr_id_dict, val_labels, out_fn)
print('Loading...')
if os.path.isdir(args.in_path):
labels = []
for p in sorted(os.listdir(args.in_path)):
with open(os.path.join(args.in_path, p)) as f:
labels.append(json.load(f))
else:
with open(args.in_path) as f:
labels = json.load(f)
print('Converting...')
out_fn = os.path.join(args.out_path)

if args.mode == 'det':
coco = bdd2coco_det(labels)
elif args.mode == 'track':
coco = bdd2coco_track(labels)

print('Saving...')
with open(out_fn, 'w') as f:
json.dump(coco, f)