coco_parser.py 文件源码

def __init__(self, annotations_path, class_names='all'):
        self.coco = COCO(annotations_path)
        self.class_names = class_names
        if self.class_names == 'all':
            class_data = self.coco.loadCats(self.coco.getCatIds())
            self.class_names = [class_['name'] for class_ in class_data]
            coco_ids = [class_['id'] for class_ in class_data]
            one_hot_ids = list(range(1, len(coco_ids) + 1))
            self.coco_id_to_class_arg = dict(zip(coco_ids, one_hot_ids))
            self.class_names = ['background'] + self.class_names
            self.num_classes = len(self.class_names)
        else:
            """
            https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoDemo.ipynb
            catIds = self.coco.getCatIds(catNms=self.class_names)
            imgIds = self.coco.getImgIds(catIds=catIds)
            imgIds = self.coco.getImgIds(imgIds=[324158])
            """