51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
import sys
|
|
import os.path as osp
|
|
from torchvision.datasets import SVHN, CIFAR10, CIFAR100
|
|
|
|
from dassl.utils import mkdir_if_missing
|
|
|
|
|
|
def extract_and_save_image(dataset, save_dir):
|
|
if osp.exists(save_dir):
|
|
print('Folder "{}" already exists'.format(save_dir))
|
|
return
|
|
|
|
print('Extracting images to "{}" ...'.format(save_dir))
|
|
mkdir_if_missing(save_dir)
|
|
|
|
for i in range(len(dataset)):
|
|
img, label = dataset[i]
|
|
class_dir = osp.join(save_dir, str(label).zfill(3))
|
|
mkdir_if_missing(class_dir)
|
|
impath = osp.join(class_dir, str(i + 1).zfill(5) + ".jpg")
|
|
img.save(impath)
|
|
|
|
|
|
def download_and_prepare(name, root):
|
|
print("Dataset: {}".format(name))
|
|
print("Root: {}".format(root))
|
|
|
|
if name == "cifar10":
|
|
train = CIFAR10(root, train=True, download=True)
|
|
test = CIFAR10(root, train=False)
|
|
elif name == "cifar100":
|
|
train = CIFAR100(root, train=True, download=True)
|
|
test = CIFAR100(root, train=False)
|
|
elif name == "svhn":
|
|
train = SVHN(root, split="train", download=True)
|
|
test = SVHN(root, split="test", download=True)
|
|
else:
|
|
raise ValueError
|
|
|
|
train_dir = osp.join(root, name, "train")
|
|
test_dir = osp.join(root, name, "test")
|
|
|
|
extract_and_save_image(train, train_dir)
|
|
extract_and_save_image(test, test_dir)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
download_and_prepare("cifar10", sys.argv[1])
|
|
download_and_prepare("cifar100", sys.argv[1])
|
|
download_and_prepare("svhn", sys.argv[1])
|