def get_coco(root, image_set, transforms, mode='instances'):
anno_file_template = "{}_{}2017.json"
PATHS = {
"train": ("train2017",
os.path.join("annotations",
anno_file_template.format(mode, "train"))),
"val": ("val2017",
os.path.join("annotations",
anno_file_template.format(mode, "val"))),
# "train": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val")))
}
t = [ConvertCocoPolysToMask()]
if transforms is not None:
t.append(transforms)
transforms = transforms_utils.Compose(t)
img_folder, ann_file = PATHS[image_set]
img_folder = os.path.join(root, img_folder)
ann_file = os.path.join(root, ann_file)
dataset = CocoDetection(img_folder, ann_file, transforms=transforms)
if image_set == "train":
dataset = _coco_remove_images_without_annotations(dataset)
# dataset = torch.utils.data.Subset(dataset, [i for i in range(500)])
return dataset
def get_transforms():
transforms_list = [transforms_utils.ToTensor()]
return transforms_utils.Compose(transforms_list)
target["iscrowd"] = iscrowd
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
def __len__(self):
return len(self.imgs_list)
if __name__ == '__main__':
import aracle.transforms_utils as transforms_utils
import aracle.data
transforms = transforms_utils.Compose([transforms_utils.ToTensor()])
#data_dir = os.path.dirname(aracle.data.__file__)
data_dir = '/nobackup/jpark45'
imgs_dir = os.path.join(data_dir, 'minidata',
'X_images_uncropped_circle_res256')
masks_dir = os.path.join(data_dir, 'minidata',
'Y_masks_uncropped_circle_res256')
hmi_dataset = HMIDataset(transforms,
imgs_dir=imgs_dir,
masks_dir=masks_dir)
n_data = len(hmi_dataset)
print(n_data)
print(hmi_dataset[0][0].shape)
sample_dict = hmi_dataset[0][1]
print(sample_dict['masks'].shape)