datasets.py

"""prepare CIFAR and SVHN
"""

from __future__ import print_function

import sys
import os
import glob
import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import Dataset
from PIL import Image
import numpy as np


crop_size = 32
padding = 4


class ImageNet_For_Search(Dataset):
    def __init__(self, root, transforms_=None, num_class=100):
        self.transform = transforms_
        self.files = []
        self.labels = []

        # dir_list = np.random.choice(os.listdir(os.path.join(root, mode)), num_class, replace=False)
        dir_list = os.listdir(root)[:num_class]
        for label, dirname in enumerate(dir_list):
            for fname in os.listdir(os.path.join(root, dirname)):
                assert 'JPEG' in fname or 'jpg' in fname or 'png' in fname
                self.files.append(os.path.join(root, dirname, fname))
                self.labels.append(label)

        data_list = list(zip(self.files, self.labels))
        np.random.shuffle(data_list)
        self.files, self.labels = zip(*data_list)

        self.files = list(self.files)
        self.labels = list(self.labels)


    def __getitem__(self, index):
        img = self.transform(Image.open(self.files[index % len(self.files)]).convert('RGB'))
        label = self.labels[index % len(self.files)]
        
        return img, label

    def __len__(self):
        # return max(len(self.files), len(self.files_B))
        return len(self.files)


def prepare_train_data_for_search(dataset='imagenet', datadir='/home/yf22/dataset', num_class=100):
    if 'imagenet' in dataset:
        train_dataset = ImageNet_For_Search(
            datadir,
            transforms.Compose([
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
            ]), 
            num_class=num_class)

    else:
        train_dataset = None
        
    return train_dataset


def prepare_test_data_for_search(dataset='imagenet', datadir='/home/yf22/dataset', num_class=100):
    if 'imagenet' in dataset:
        train_dataset = ImageNet_For_Search(
            datadir,
            transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                         std=[0.229, 0.224, 0.225])
            ]), 
            num_class=num_class)

    else:
        train_dataset = None
        
    return train_dataset


def prepare_train_data(dataset='imagenet', datadir='/home/yf22/dataset'):
    if 'imagenet' in dataset:
        train_dataset = torchvision.datasets.ImageFolder(
            datadir,
            transforms.Compose([
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
            ]))

    else:
        train_dataset = None
        
    return train_dataset


def prepare_test_data(dataset='imagenet', datadir='/home/yf22/dataset'):

    if 'imagenet' in dataset:
        test_dataset = torchvision.datasets.ImageFolder(datadir, transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
        ]))

    else:
        test_dataset = None

    return test_dataset