| | |
| | import os.path as osp |
| | from typing import List, Union |
| |
|
| | from mmdet.registry import DATASETS |
| | from .base_video_dataset import BaseVideoDataset |
| |
|
| |
|
| | @DATASETS.register_module() |
| | class MOTChallengeDataset(BaseVideoDataset): |
| | """Dataset for MOTChallenge. |
| | |
| | Args: |
| | visibility_thr (float, optional): The minimum visibility |
| | for the objects during training. Default to -1. |
| | """ |
| |
|
| | METAINFO = { |
| | 'classes': |
| | ('pedestrian', 'person_on_vehicle', 'car', 'bicycle', 'motorbike', |
| | 'non_mot_vehicle', 'static_person', 'distractor', 'occluder', |
| | 'occluder_on_ground', 'occluder_full', 'reflection', 'crowd') |
| | } |
| |
|
| | def __init__(self, visibility_thr: float = -1, *args, **kwargs): |
| | self.visibility_thr = visibility_thr |
| | super().__init__(*args, **kwargs) |
| |
|
| | def parse_data_info(self, raw_data_info: dict) -> Union[dict, List[dict]]: |
| | """Parse raw annotation to target format. The difference between this |
| | function and the one in ``BaseVideoDataset`` is that the parsing here |
| | adds ``visibility`` and ``mot_conf``. |
| | |
| | Args: |
| | raw_data_info (dict): Raw data information load from ``ann_file`` |
| | |
| | Returns: |
| | Union[dict, List[dict]]: Parsed annotation. |
| | """ |
| | img_info = raw_data_info['raw_img_info'] |
| | ann_info = raw_data_info['raw_ann_info'] |
| | data_info = {} |
| |
|
| | data_info.update(img_info) |
| | if self.data_prefix.get('img_path', None) is not None: |
| | img_path = osp.join(self.data_prefix['img_path'], |
| | img_info['file_name']) |
| | else: |
| | img_path = img_info['file_name'] |
| | data_info['img_path'] = img_path |
| |
|
| | instances = [] |
| | for i, ann in enumerate(ann_info): |
| | instance = {} |
| |
|
| | if (not self.test_mode) and (ann['visibility'] < |
| | self.visibility_thr): |
| | continue |
| | if ann.get('ignore', False): |
| | continue |
| | x1, y1, w, h = ann['bbox'] |
| | inter_w = max(0, min(x1 + w, img_info['width']) - max(x1, 0)) |
| | inter_h = max(0, min(y1 + h, img_info['height']) - max(y1, 0)) |
| | if inter_w * inter_h == 0: |
| | continue |
| | if ann['area'] <= 0 or w < 1 or h < 1: |
| | continue |
| | if ann['category_id'] not in self.cat_ids: |
| | continue |
| | bbox = [x1, y1, x1 + w, y1 + h] |
| |
|
| | if ann.get('iscrowd', False): |
| | instance['ignore_flag'] = 1 |
| | else: |
| | instance['ignore_flag'] = 0 |
| | instance['bbox'] = bbox |
| | instance['bbox_label'] = self.cat2label[ann['category_id']] |
| | instance['instance_id'] = ann['instance_id'] |
| | instance['category_id'] = ann['category_id'] |
| | instance['mot_conf'] = ann['mot_conf'] |
| | instance['visibility'] = ann['visibility'] |
| | if len(instance) > 0: |
| | instances.append(instance) |
| | if not self.test_mode: |
| | assert len(instances) > 0, f'No valid instances found in ' \ |
| | f'image {data_info["img_path"]}!' |
| | data_info['instances'] = instances |
| | return data_info |
| |
|
