| |
| import json |
| import os.path as osp |
| from typing import List, Optional |
|
|
| from mmengine.fileio import get_local_path |
|
|
| from mmdet.registry import DATASETS |
| from .base_det_dataset import BaseDetDataset |
|
|
|
|
| @DATASETS.register_module() |
| class ODVGDataset(BaseDetDataset): |
| """object detection and visual grounding dataset.""" |
|
|
| def __init__(self, |
| *args, |
| data_root: str = '', |
| label_map_file: Optional[str] = None, |
| need_text: bool = True, |
| **kwargs) -> None: |
| self.dataset_mode = 'VG' |
| self.need_text = need_text |
| if label_map_file: |
| label_map_file = osp.join(data_root, label_map_file) |
| with open(label_map_file, 'r') as file: |
| self.label_map = json.load(file) |
| self.dataset_mode = 'OD' |
| super().__init__(*args, data_root=data_root, **kwargs) |
| assert self.return_classes is True |
|
|
| def load_data_list(self) -> List[dict]: |
| with get_local_path( |
| self.ann_file, backend_args=self.backend_args) as local_path: |
| with open(local_path, 'r') as f: |
| data_list = [json.loads(line) for line in f] |
|
|
| out_data_list = [] |
| for data in data_list: |
| data_info = {} |
| img_path = osp.join(self.data_prefix['img'], data['filename']) |
| data_info['img_path'] = img_path |
| data_info['height'] = data['height'] |
| data_info['width'] = data['width'] |
| if self.dataset_mode == 'OD': |
| if self.need_text: |
| data_info['text'] = self.label_map |
| anno = data.get('detection', {}) |
| instances = [obj for obj in anno.get('instances', [])] |
| bboxes = [obj['bbox'] for obj in instances] |
| bbox_labels = [str(obj['label']) for obj in instances] |
|
|
| instances = [] |
| for bbox, label in zip(bboxes, bbox_labels): |
| instance = {} |
| x1, y1, x2, y2 = bbox |
| inter_w = max(0, min(x2, data['width']) - max(x1, 0)) |
| inter_h = max(0, min(y2, data['height']) - max(y1, 0)) |
| if inter_w * inter_h == 0: |
| continue |
| if (x2 - x1) < 1 or (y2 - y1) < 1: |
| continue |
| instance['ignore_flag'] = 0 |
| instance['bbox'] = bbox |
| instance['bbox_label'] = int(label) |
| instances.append(instance) |
| data_info['instances'] = instances |
| data_info['dataset_mode'] = self.dataset_mode |
| out_data_list.append(data_info) |
| else: |
| anno = data['grounding'] |
| data_info['text'] = anno['caption'] |
| regions = anno['regions'] |
|
|
| instances = [] |
| phrases = {} |
| for i, region in enumerate(regions): |
| bbox = region['bbox'] |
| phrase = region['phrase'] |
| tokens_positive = region['tokens_positive'] |
| if not isinstance(bbox[0], list): |
| bbox = [bbox] |
| for box in bbox: |
| instance = {} |
| x1, y1, x2, y2 = box |
| inter_w = max(0, min(x2, data['width']) - max(x1, 0)) |
| inter_h = max(0, min(y2, data['height']) - max(y1, 0)) |
| if inter_w * inter_h == 0: |
| continue |
| if (x2 - x1) < 1 or (y2 - y1) < 1: |
| continue |
| instance['ignore_flag'] = 0 |
| instance['bbox'] = box |
| instance['bbox_label'] = i |
| phrases[i] = { |
| 'phrase': phrase, |
| 'tokens_positive': tokens_positive |
| } |
| instances.append(instance) |
| data_info['instances'] = instances |
| data_info['phrases'] = phrases |
| data_info['dataset_mode'] = self.dataset_mode |
| out_data_list.append(data_info) |
|
|
| del data_list |
| return out_data_list |
|
|
