Spaces:
Runtime error
Runtime error
File size: 2,165 Bytes
1c3eb47 |
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 |
from typing import List
from mmpl.registry import DATASETS
from mmdet.datasets.coco import CocoDataset
@DATASETS.register_module()
class NWPUInsSegDataset(CocoDataset):
"""Dataset for Cityscapes."""
METAINFO = {
'classes': ['airplane', 'ship', 'storage_tank', 'baseball_diamond',
'tennis_court', 'basketball_court', 'ground_track_field',
'harbor', 'bridge', 'vehicle'],
'palette': [(220, 20, 60), (255, 0, 0), (0, 0, 142), (0, 0, 70),
(0, 60, 100), (0, 80, 100), (0, 0, 230),
(119, 11, 32), (0, 255, 0), (0, 0, 255)]
}
def filter_data(self) -> List[dict]:
"""Filter annotations according to filter_cfg.
Returns:
List[dict]: Filtered results.
"""
if self.test_mode:
return self.data_list
if self.filter_cfg is None:
return self.data_list
filter_empty_gt = self.filter_cfg.get('filter_empty_gt', False)
min_size = self.filter_cfg.get('min_size', 0)
# obtain images that contain annotation
ids_with_ann = set(data_info['img_id'] for data_info in self.data_list)
# obtain images that contain annotations of the required categories
ids_in_cat = set()
for i, class_id in enumerate(self.cat_ids):
ids_in_cat |= set(self.cat_img_map[class_id])
# merge the image id sets of the two conditions and use the merged set
# to filter out images if self.filter_empty_gt=True
ids_in_cat &= ids_with_ann
valid_data_infos = []
for i, data_info in enumerate(self.data_list):
img_id = data_info['img_id']
width = data_info['width']
height = data_info['height']
all_is_crowd = all([
instance['ignore_flag'] == 1
for instance in data_info['instances']
])
if filter_empty_gt and (img_id not in ids_in_cat or all_is_crowd):
continue
if min(width, height) >= min_size:
valid_data_infos.append(data_info)
return valid_data_infos
|