File size: 3,106 Bytes
67bb36a |
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 |
import numpy as np
import json
from .AutoDriveDataset import AutoDriveDataset
from .convert import convert, id_dict, id_dict_single
from tqdm import tqdm
import os
single_cls = False # just detect vehicle
class HustDataset(AutoDriveDataset):
def __init__(self, cfg, is_train, inputsize, transform=None):
super().__init__(cfg, is_train, inputsize, transform)
self.db = self._get_db()
self.cfg = cfg
def _get_db(self):
"""
get database from the annotation file
Inputs:
Returns:
gt_db: (list)database [a,b,c,...]
a: (dictionary){'image':, 'information':, ......}
image: image path
mask: path of the segmetation label
label: [cls_id, center_x//256, center_y//256, w//256, h//256] 256=IMAGE_SIZE
"""
print('building database...')
gt_db = []
height, width = self.shapes
for mask in tqdm(list(self.mask_list)):
mask_path = str(mask)
label_path = self.label_root
# label_path = mask_path.replace(str(self.mask_root), str(self.label_root)).replace(".png", ".json")
image_path = mask_path.replace(str(self.mask_root), str(self.img_root)).replace(".png", ".jpg")
lane_path = mask_path.replace(str(self.mask_root), str(self.lane_root))
with open(label_path, 'r') as f:
label = json.load(f)
data = label[int(os.path.basename(image_path)[:-4])]["labels"]
data = self.filter_data(data)
gt = np.zeros((len(data), 5))
for idx, obj in enumerate(data):
category = obj['category']
if category == "traffic light":
color = obj['attributes']['Traffic Light Color'][0]
category = "tl_" + color
if category in id_dict.keys():
x1 = float(obj['box2d']['x1'])
y1 = float(obj['box2d']['y1'])
x2 = float(obj['box2d']['x2'])
y2 = float(obj['box2d']['y2'])
cls_id = id_dict[category]
if single_cls:
cls_id=0
gt[idx][0] = cls_id
box = convert((width, height), (x1, x2, y1, y2))
gt[idx][1:] = list(box)
rec = [{
'image': image_path,
'label': gt,
'mask': mask_path,
'lane': lane_path
}]
gt_db += rec
print('database build finish')
return gt_db
def filter_data(self, data):
remain = []
for obj in data:
if 'box2d' in obj.keys(): # obj.has_key('box2d'):
if single_cls:
if obj['category'] in id_dict_single.keys():
remain.append(obj)
else:
remain.append(obj)
return remain
def evaluate(self, cfg, preds, output_dir, *args, **kwargs):
"""
"""
pass
|