Dubhe/model_measuring/kamal/vision/datasets/stanford_cars.py

import os
import glob
from PIL import Image
import numpy as np
from scipy.io import loadmat

from torch.utils import data
from .utils import download_url, mkdir

from shutil import copyfile


class StanfordCars(data.Dataset):
    """Dataset for Stanford Cars
    """

    urls = {'cars_train.tgz':     'http://imagenet.stanford.edu/internal/car196/cars_train.tgz',
            'cars_test.tgz':       'http://imagenet.stanford.edu/internal/car196/cars_test.tgz',
            'car_devkit.tgz':      'https://ai.stanford.edu/~jkrause/cars/car_devkit.tgz',
            'cars_test_annos_withlabels.mat': 'http://imagenet.stanford.edu/internal/car196/cars_test_annos_withlabels.mat'}

    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()

        if self.split == 'train':
            annos = os.path.join(self.root, 'devkit', 'cars_train_annos.mat')
        else:
            annos = os.path.join(self.root, 'devkit',
                                 'cars_test_annos_withlabels.mat')

        annos = loadmat(annos)
        size = len(annos['annotations'][0])

        self.files = glob.glob(os.path.join(
            self.root, 'cars_'+self.split, '*.jpg'))
        self.files.sort()

        self.labels = np.array([int(l[4])-1 for l in annos['annotations'][0]])

        lbl_annos = loadmat(os.path.join(self.root, 'devkit', 'cars_meta.mat'))

        self.object_categories = [str(c[0])
                                  for c in lbl_annos['class_names'][0]]

        print('Stanford Cars, 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

        mkdir(self.root)
        for fname, url in self.urls.items():
            if not os.path.isfile(os.path.join(self.root, fname)):
                download_url(url, self.root, fname)
            if fname.endswith('tgz'):
                print("Extracting %s..." % fname)
                with tarfile.open(os.path.join(self.root, fname), "r:gz") as tar:
                    tar.extractall(path=self.root)

        copyfile(os.path.join(self.root, 'cars_test_annos_withlabels.mat'),
                 os.path.join(self.root, 'devkit', 'cars_test_annos_withlabels.mat'))