| 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 |
|
|
| 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 |
| |
| 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(): |
| 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 |
|
|
