Spaces:
Sleeping
Sleeping
File size: 2,930 Bytes
b34d1d6 |
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 |
import os.path as osp
from mmdet.registry import DATASETS
from mmdet.datasets.coco_panoptic import CocoPanopticDataset
@DATASETS.register_module()
class CocoPanopticSAMDataset(CocoPanopticDataset):
"""Coco SAM dataset, stuff is treated as thing.
"""
def parse_data_info(self, raw_data_info: dict) -> dict:
"""Parse raw annotation to target format.
Args:
raw_data_info (dict): Raw data information load from ``ann_file``.
Returns:
dict: Parsed annotation.
"""
img_info = raw_data_info['raw_img_info']
ann_info = raw_data_info['raw_ann_info']
# filter out unmatched annotations which have
# same segment_id but belong to other image
ann_info = [
ann for ann in ann_info if ann['image_id'] == img_info['img_id']
]
data_info = {}
img_path = osp.join(self.data_prefix['img'], img_info['file_name'])
if self.data_prefix.get('seg', None):
seg_map_path = osp.join(
self.data_prefix['seg'],
img_info['file_name'].replace('jpg', 'png'))
else:
seg_map_path = None
data_info['img_path'] = img_path
data_info['img_id'] = img_info['img_id']
data_info['seg_map_path'] = seg_map_path
data_info['height'] = img_info['height']
data_info['width'] = img_info['width']
if self.return_classes:
data_info['text'] = self.metainfo['thing_classes']
data_info['stuff_text'] = self.metainfo['stuff_classes']
data_info['custom_entities'] = True # no important
instances = []
segments_info = []
for ann in ann_info:
instance = {}
x1, y1, w, h = ann['bbox']
if ann['area'] <= 0 or w < 1 or h < 1:
continue
bbox = [x1, y1, x1 + w, y1 + h]
category_id = ann['category_id']
contiguous_cat_id = self.cat2label[category_id]
_ = self.coco.load_cats(ids=category_id)[0]['isthing']
# always set is_thing to true
is_thing = True
if is_thing:
is_crowd = ann.get('iscrowd', False)
instance['bbox'] = bbox
instance['bbox_label'] = contiguous_cat_id
if not is_crowd:
instance['ignore_flag'] = 0
else:
instance['ignore_flag'] = 1
is_thing = False
segment_info = {
'id': ann['id'],
'category': contiguous_cat_id,
'is_thing': is_thing
}
segments_info.append(segment_info)
if len(instance) > 0 and is_thing:
instances.append(instance)
data_info['instances'] = instances
data_info['segments_info'] = segments_info
return data_info
|