(self, root, ann_file, is_train=True, return_location=False, logits_file=None)
| 55 | |
| 56 | class INAT_ImageDataset(data.Dataset): |
| 57 | def __init__(self, root, ann_file, is_train=True, return_location=False, logits_file=None): |
| 58 | |
| 59 | # load annotations |
| 60 | print('Loading annotations from: ' + os.path.basename(ann_file)) |
| 61 | with open(ann_file) as data_file: |
| 62 | ann_data = json.load(data_file) |
| 63 | |
| 64 | # set up the filenames and annotations |
| 65 | self.imgs = [aa['file_name'] for aa in ann_data['images']] |
| 66 | self.ids = [aa['id'] for aa in ann_data['images']] |
| 67 | |
| 68 | # if we dont have class labels set them to '0' |
| 69 | if 'annotations' in ann_data.keys(): |
| 70 | self.classes = [aa['category_id'] for aa in ann_data['annotations']] |
| 71 | else: |
| 72 | self.classes = [0]*len(self.imgs) |
| 73 | |
| 74 | # load taxonomy |
| 75 | self.tax_levels = ['id', 'genus', 'family', 'order', 'class', 'phylum', 'kingdom'] |
| 76 | #8142, 4412, 1120, 273, 57, 25, 6 |
| 77 | self.taxonomy, self.classes_taxonomic = load_taxonomy(ann_data, self.tax_levels, self.classes) |
| 78 | |
| 79 | # print out some stats |
| 80 | print('\t' + str(len(self.imgs)) + ' images') |
| 81 | print('\t' + str(len(set(self.classes))) + ' classes') |
| 82 | |
| 83 | self.root = root |
| 84 | self.is_train = is_train |
| 85 | self.loader = default_loader |
| 86 | |
| 87 | # augmentation params |
| 88 | self.im_size = [299, 299] # can change this to train on higher res |
| 89 | self.mu_data = [0.485, 0.456, 0.406] |
| 90 | self.std_data = [0.229, 0.224, 0.225] |
| 91 | self.brightness = 0.4 |
| 92 | self.contrast = 0.4 |
| 93 | self.saturation = 0.4 |
| 94 | self.hue = 0.25 |
| 95 | |
| 96 | # augmentations |
| 97 | self.center_crop = transforms.CenterCrop((self.im_size[0], self.im_size[1])) |
| 98 | self.scale_aug = transforms.RandomResizedCrop(size=self.im_size[0]) |
| 99 | self.flip_aug = transforms.RandomHorizontalFlip() |
| 100 | self.color_aug = transforms.ColorJitter(self.brightness, self.contrast, self.saturation, self.hue) |
| 101 | self.tensor_aug = transforms.ToTensor() |
| 102 | self.norm_aug = transforms.Normalize(mean=self.mu_data, std=self.std_data) |
| 103 | |
| 104 | # load logits |
| 105 | self.logits_file = logits_file |
| 106 | if self.logits_file is not None: |
| 107 | self.logits = np.load(self.logits_file, mmap_mode='r+') |
| 108 | |
| 109 | self.return_location = return_location |
| 110 | if return_location: |
| 111 | # location data |
| 112 | path, ext = os.path.splitext(ann_file) |
| 113 | locations, classes, users, dates, keep_indxs, data_ids = load_inat_location_data(root, os.path.basename(path) + "_locations" + ext, os.path.basename(ann_file)) |
| 114 | print(f"dropping {len(self.ids) - len(data_ids)} of {len(self.ids)} samples due to lacking location") |
no test coverage detected