import os
import os.path
import numpy as np
from torchvision.datasets.utils import check_integrity, download_and_extract_archive, verify_str_arg
# reference: https://pytorch.org/vision/0.8/_modules/torchvision/datasets/stl10.html#STL10
[docs]
def load_file(data_file, labels_file=None):
labels = None
if labels_file:
with open(labels_file, 'rb') as f:
labels = np.fromfile(f, dtype=np.uint8) - 1 # 0-based
with open(data_file, 'rb') as f:
# read whole file in uint8 chunks
everything = np.fromfile(f, dtype=np.uint8)
images = np.reshape(everything, (-1, 3, 96, 96))
images = np.transpose(images, (0, 1, 3, 2))
return images, labels
def _check_integrity(root, train_list, test_list, base_folder):
for fentry in (train_list + test_list):
filename, md5 = fentry[0], fentry[1]
fpath = os.path.join(root, base_folder, filename)
if not check_integrity(fpath, md5):
return False
print('Files already downloaded and verified')
return True
[docs]
def STL10(root='./data/', split='train'):
base_folder = 'stl10_binary'
url = "http://ai.stanford.edu/~acoates/stl10/stl10_binary.tar.gz"
filename = "stl10_binary.tar.gz"
class_names_file = 'class_names.txt'
folds_list_file = 'fold_indices.txt'
train_list = [
['train_X.bin', '918c2871b30a85fa023e0c44e0bee87f'],
['train_y.bin', '5a34089d4802c674881badbb80307741'],
['unlabeled_X.bin', '5242ba1fed5e4be9e1e742405eb56ca4']
]
test_list = [
['test_X.bin', '7f263ba9f9e0b06b93213547f721ac82'],
['test_y.bin', '36f9794fa4beb8a2c72628de14fa638e']
]
splits = ('train', 'train+unlabeled', 'unlabeled', 'test')
# download dataset
fpath = os.path.join(root, base_folder, filename)
if not _check_integrity(root, train_list, test_list, base_folder):
download_and_extract_archive(url=url, download_root=root, filename=filename)
# choose which set to load
if split=='train':
path_to_data = os.path.join(root, base_folder, train_list[0][0])
path_to_labels = os.path.join(root, base_folder, train_list[1][0])
data, targets = load_file(path_to_data, path_to_labels)
elif split == 'unlabeled':
path_to_data = os.path.join(root, base_folder, train_list[2][0])
data, _ = load_file(path_to_data)
targets = np.asarray([-1] * data.shape[0])
elif split == 'test':
path_to_data = os.path.join(root, base_folder, test_list[0][0])
path_to_labels = os.path.join(root, base_folder, test_list[1][0])
data, targets = load_file(path_to_data, path_to_labels)
else:
raise ValueError('Out of option!')
return data, targets
if __name__ == '__main__':
data, targets = STL10(root='./data/', split='test') # return numpy array
print (data.shape, targets.shape)