| 50 | |
| 51 | |
| 52 | def get_tinyImageNet(path): |
| 53 | # 100000 train 10000 test |
| 54 | raw_tr = datasets.ImageFolder(path + '/tinyImageNet/tiny-imagenet-200/train') |
| 55 | raw_te = datasets.ImageFolder(path + '/tinyImageNet/tiny-imagenet-200/val') |
| 56 | f = open(path + '/tinyImageNet/tiny-imagenet-200/val/val_annotations.txt') |
| 57 | |
| 58 | val_dict = {} |
| 59 | for line in f.readlines(): |
| 60 | val_dict[line.split()[0]] = raw_tr.class_to_idx[line.split()[1]] |
| 61 | X_tr,Y_tr,X_te, Y_te = [],[],[],[] |
| 62 | |
| 63 | div_list = [len(raw_tr)*(x+1)//10 for x in range(10)] # can not load at once, memory limitation |
| 64 | i=0 |
| 65 | for count in div_list: |
| 66 | loop = count - i |
| 67 | for j in range(loop): |
| 68 | image,target = raw_tr[i] |
| 69 | X_tr.append(np.array(image)) |
| 70 | Y_tr.append(target) |
| 71 | i += 1 |
| 72 | |
| 73 | for i in range(len(raw_te)): |
| 74 | img, label = raw_te[i] |
| 75 | img_pth = raw_te.imgs[i][0].split('/')[-1] |
| 76 | X_te.append(np.array(img)) |
| 77 | Y_te.append(val_dict[img_pth]) |
| 78 | |
| 79 | return X_tr,Y_tr,X_te, Y_te |
| 80 | # torch.tensor(X_tr), torch.tensor(Y_tr), torch.tensor(X_te), torch.tensor(Y_te) |
| 81 | |
| 82 | def get_MNIST(path): |
| 83 | raw_tr = datasets.MNIST(path + '/mnist', train=True, download=True) |