coco_dataset.py 文件源码

def __init__(self, data_dir, image_set, year):
        self.name = 'coco_' + year + '_' + image_set
        # name, paths
        self.year = year
        self.image_set = image_set
        self.data_dir = data_dir

        # COCO API
        self.COCO = COCO(self._get_ann_file())
        cats = self.COCO.loadCats(self.COCO.getCatIds())
        self.classes = tuple([c['name'] for c in cats])
        self.num_classes = len(self.classes)
        # e.g: 'person' -> 0, 'toothbrush' -> 79
        self.class_to_ind = dict(zip(self.classes, xrange(self.num_classes)))
        # e.g: 'person' -> 1, 'toothbrush' -> 90
        self.class_to_coco_cat_id = dict(zip([c['name'] for c in cats], self.COCO.getCatIds()))
        # coco' cat_id (1,,90) -> label_id (0,,79)
        self.coco_cat_id_to_class_ind = dict([(self.class_to_coco_cat_id[cls],
                                          self.class_to_ind[cls])
                                         for cls in self.classes])

        self.image_index = self._load_image_set_index()
        self.num_images = len(self.image_index)

        coco_name = image_set + year  # e.g., "val2014"
        self.data_name = coco_name