def __init__(self, root, image_sets, preproc=None, target_transform=None, dataset_name='COCO'): self.root = root self.cache_path = os.path.join(self.root, 'cache') self.image_set = image_sets self.preproc = preproc #target_transform = transforms.RandomResizedCrop(300) self.target_transform = target_transform self.name = dataset_name self.ids = list() self.annotations = list() self._view_map = { 'minival2014' : 'val2014', # 5k val2014 subset 'valminusminival2014' : 'val2014', # val2014 \setminus minival2014 'test-dev2015' : 'test2015', } # self.data_name = list() # self.data_len = list() for (year, image_set) in image_sets: self.image_set = image_set coco_name = image_set+year data_name = (self._view_map[coco_name] if coco_name in self._view_map else coco_name) annofile = self._get_ann_file(coco_name) print('annofile:',annofile) _COCO = COCO(annofile) self._COCO = _COCO self.coco_name = coco_name cats = _COCO.loadCats(_COCO.getCatIds()) print('cats ', cats) self._classes = tuple(['__background__'] + [c['name'] for c in cats]) self.num_classes = len(self._classes) print('class ', self._classes) self._class_to_ind = dict(zip(self._classes, range(self.num_classes))) self._class_to_coco_cat_id = dict(zip([c['name'] for c in cats], _COCO.getCatIds())) indexes = _COCO.getImgIds() self.image_indexes = indexes #print('index ', indexes) # seems it will reduce the speed during the training. # self.ids.extend(indexes) # self.data_len.append(len(indexes)) # self.data_name.append(data_name) self.ids.extend(self._load_coco_img_path(coco_name, _COCO.getImgFilenames())) if image_set.find('test') != -1: print('test set will not load annotations!') else: self.annotations.extend(self._load_coco_annotations(coco_name, indexes,_COCO))
def __init__(self, root, image_sets, preproc=None, target_transform=None, dataset_name='COCO'): self.root = None self.cache_path = os.path.join('/home/zhihuai/yolo', 'cache') self.image_set = image_sets self.preproc = preproc self.target_transform = target_transform self.name = dataset_name self.ids = list() self.annotations = list() self.OfficeDataset = [ 'head_difficult_train', 'head_office_train', # 5k val2014 subset 'head_squatting_train', 'head_steve_train', # val2014 \setminus minival2014 'head_v2_benchmark_0_train' ] self.CrowdHuman = ['head_CrowdHuman_train'] self.coco_crowdhuman = ['coco_crowdhuman_train'] self.failed_frames = ['failed_frames_train'] self.training_data_resources1 = ['training_data_resources1'] self.training_data_resources2 = ['training_data_resources2'] self.training_data_resources3 = ['training_data_resources3'] self.test_set = ['annotations'] self.num = 0 for image_set in image_sets: if image_set in self.OfficeDataset: self.root = os.path.join(root, 'OfficeDataset') coco_name = image_set self.data_name = 'OfficeDataset_' #in order to diff some anno has the same name elif image_set in self.CrowdHuman: self.root = os.path.join(root, 'CrowdHuman') coco_name = image_set self.data_name = 'CrowdHuman_' elif image_set in self.coco_crowdhuman: self.root = os.path.join(root, 'tracking_data/training_only', 'coco_crowdhuman') coco_name = 'annotations' self.data_name = 'coco_crowdhuman_' elif image_set in self.failed_frames: self.root = os.path.join(root, 'tracking_data/training_only', 'failed_frames') coco_name = 'annotations' self.data_name = 'failed_frames_' elif image_set in self.training_data_resources1: self.root = os.path.join(root, 'tracking_data/training_only', 'training_data_resources1') coco_name = 'annotations' self.data_name = 'training_data_resources1_' elif image_set in self.training_data_resources2: self.root = os.path.join(root, 'tracking_data/training_only', 'training_data_resources2') coco_name = 'annotations' self.data_name = 'training_data_resources2_' elif image_set in self.training_data_resources3: self.root = os.path.join(root, 'tracking_data/training_only', 'training_data_resources3') coco_name = 'annotations' self.data_name = 'training_data_resources3_' elif image_set in self.test_set: self.root = root coco_name = image_set self.data_name = 'test_' annofile = self._get_ann_file(coco_name) print('annofile', annofile) _COCO = COCO(annofile) self._COCO = _COCO self.coco_name = coco_name cats = _COCO.loadCats(_COCO.getCatIds()) self._classes = tuple(['__background__'] + [c['name'] for c in cats]) self.num_classes = len(self._classes) self._class_to_ind = dict( zip(self._classes, range(self.num_classes))) self._class_to_coco_cat_id = dict( zip([c['name'] for c in cats], _COCO.getCatIds())) indexes = _COCO.getImgIds() self.image_indexes = indexes # seems it will reduce the speed during the training. # self.ids.extend(indexes) # self.data_len.append(len(indexes)) # self.data_name.append(data_name) self.ids.extend(self._load_coco_img_path(coco_name, indexes)) if image_set.find('test') != -1: print('test set will not load annotations!') else: self.annotations.extend( self._load_coco_annotations(coco_name, indexes, _COCO)) self.num += len(self.image_indexes) print("image_num", self.num)