WuWenc/tiny_coco

1import os2 3import datasets4from pycocotools.coco import COCO5 6_DESCRIPTION = 'A tiny coco2017 dataset example.'7 8_URLS = {9    'train': 'train2017.zip',10    'train_meta': 'annotations/instances_train2017.json',11    'val': 'val2017.zip',12    'val_meta': 'annotations/instances_val2017.json',13}14 15_CLASSES = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',16            'train', 'truck', 'boat', 'traffic light', 'fire hydrant',17            'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',18            'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',19            'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',20            'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',21            'baseball glove', 'skateboard', 'surfboard', 'tennis racket',22            'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',23            'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',24            'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',25            'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',26            'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',27            'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',28            'scissors', 'teddy bear', 'hair drier', 'toothbrush')29 30 31class TinyCoco(datasets.GeneratorBasedBuilder):32    """TODO: Short description of my dataset."""33 34    VERSION = datasets.Version('0.1.0')35 36    BUILDER_CONFIGS = [37        datasets.BuilderConfig(38            name='train', version=VERSION, description='Training set'),39        datasets.BuilderConfig(40            name='val', version=VERSION, description='Validation set'),41    ]42    # It's not mandatory to have a default configuration.43    # Just use one if it make sense.44    DEFAULT_CONFIG_NAME = 'train'45 46    def _info(self):47        return datasets.DatasetInfo(48            # This is the description that will appear on the datasets page.49            description=_DESCRIPTION+f'\nCLASSES: ({",".join(_CLASSES)})'50        )51 52    def _split_generators(self, dl_manager):53        data_dir = dl_manager.download_and_extract(_URLS[self.config.name])54        meta = dl_manager.download(_URLS[self.config.name + '_meta'])55        return [56            datasets.SplitGenerator(57                name=self.config.name,58                # These kwargs will be passed to _generate_examples59                gen_kwargs={60                    'img_prefix': data_dir,61                    'ann_file': meta62                })63        ]64 65    def _generate_examples(self, img_prefix, ann_file):66        """Parser coco format annotation file."""67        coco = COCO(ann_file)68        cat_ids = coco.getCatIds(_CLASSES)69        cat2label = {cat_id: i for i, cat_id in enumerate(cat_ids)}70        img_ids = coco.getImgIds()71        index = 072        for i in img_ids:73            sample = dict()74            info = coco.loadImgs([i])[0]75            sample['filename'] = os.path.join(img_prefix, info['file_name'])76            sample['height'] = info['height']77            sample['width'] = info['width']78            ann_ids = coco.getAnnIds([i])79            ann_info = coco.loadAnns(ann_ids)80            gt_bboxes = []81            gt_labels = []82            gt_bboxes_ignore = []83            gt_label_ignore = []84            gt_masks_ann = []85            for i, ann in enumerate(ann_info):86                if ann.get('ignore', False):87                    continue88                x1, y1, w, h = ann['bbox']89                inter_w = max(0, min(x1 + w, sample['width']) - max(x1, 0))90                inter_h = max(0, min(y1 + h, sample['height']) - max(y1, 0))91                if inter_w * inter_h == 0:92                    continue93                if ann['area'] <= 0 or w < 1 or h < 1:94                    continue95                if ann['category_id'] not in cat_ids:96                    continue97                bbox = [x1, y1, x1 + w, y1 + h]98                if ann.get('iscrowd', False):99                    gt_bboxes_ignore.append(bbox)100                    gt_label_ignore.append(cat2label[ann['category_id']])101                else:102                    gt_bboxes.append(bbox)103                    gt_labels.append(cat2label[ann['category_id']])104                    gt_masks_ann.append(ann.get('segmentation', None))105 106            sample['ann'] = dict(107                bboxes=gt_bboxes,108                labels=gt_labels,109                bboxes_ignore=gt_bboxes_ignore,110                label_ignore=gt_label_ignore)111            yield index, sample112            index += 1113