|
- import os
- import numpy as np
- from PIL import Image
- from scipy.io import loadmat
-
- from torch.utils import data
- from .utils import download_url
- from shutil import move
-
- class StanfordDogs(data.Dataset):
- """Dataset for Stanford Dogs
- """
- urls = {"images.tar": "http://vision.stanford.edu/aditya86/ImageNetDogs/images.tar",
- "annotation.tar": "http://vision.stanford.edu/aditya86/ImageNetDogs/annotation.tar",
- "lists.tar": "http://vision.stanford.edu/aditya86/ImageNetDogs/lists.tar"}
-
- def __init__(self, root, split='train', download=False, transform=None, target_transform=None):
- self.root = os.path.abspath( os.path.expanduser(root) )
- self.split = split
- self.transform = transform
- self.target_transform = target_transform
- if download:
- self.download()
- list_file = os.path.join(self.root, self.split+'_list.mat')
- mat_file = loadmat(list_file)
- size = len(mat_file['file_list'])
- self.files = [str(mat_file['file_list'][i][0][0]) for i in range(size)]
- self.labels = np.array(
- [mat_file['labels'][i][0]-1 for i in range(size)])
- categories = os.listdir(os.path.join(self.root, 'Images'))
- categories.sort()
- self.object_categories = [c[10:] for c in categories]
- print('Stanford Dogs, Split: %s, Size: %d' %
- (self.split, self.__len__()))
-
- def __len__(self):
- return len(self.files)
-
- def __getitem__(self, idx):
- img = Image.open(os.path.join(self.root, 'Images',
- self.files[idx])).convert("RGB")
- lbl = self.labels[idx]
- if self.transform is not None:
- img = self.transform(img)
- if self.target_transform is not None:
- lbl = self.target_transform( lbl )
- return img, lbl
-
- def download(self):
- import tarfile
- os.makedirs(self.root, exist_ok=True)
- for fname, url in self.urls.items():
- if not os.path.isfile(os.path.join(self.root, fname)):
- download_url(url, self.root, fname)
- # extract file
- print("Extracting %s..." % fname)
- with tarfile.open(os.path.join(self.root, fname), "r") as tar:
- tar.extractall(path=self.root)
|
