import os import contextlib import tarfile import json import numpy as np import PIL import torch from common_utils import get_tmp_dir import pickle import random from itertools import cycle from torchvision.io.video import write_video import unittest.mock import hashlib from distutils import dir_util import re def mock_class_attribute(stack, target, new): mock = unittest.mock.patch(target, new_callable=unittest.mock.PropertyMock, return_value=new) stack.enter_context(mock) return mock def compute_md5(file): with open(file, "rb") as fh: return hashlib.md5(fh.read()).hexdigest() def make_tar(root, name, *files, compression=None): ext = ".tar" mode = "w" if compression is not None: ext = f"{ext}.{compression}" mode = f"{mode}:{compression}" name = os.path.splitext(name)[0] + ext archive = os.path.join(root, name) with tarfile.open(archive, mode) as fh: for file in files: fh.add(os.path.join(root, file), arcname=file) return name, compute_md5(archive) def clean_dir(root, *keep): pattern = re.compile(f"({f')|('.join(keep)})") for file_or_dir in os.listdir(root): if pattern.search(file_or_dir): continue file_or_dir = os.path.join(root, file_or_dir) if os.path.isfile(file_or_dir): os.remove(file_or_dir) else: dir_util.remove_tree(file_or_dir) @contextlib.contextmanager def mnist_root(num_images, cls_name): def _encode(v): return torch.tensor(v, dtype=torch.int32).numpy().tobytes()[::-1] def _make_image_file(filename, num_images): img = torch.randint(0, 256, size=(28 * 28 * num_images,), dtype=torch.uint8) with open(filename, "wb") as f: f.write(_encode(2051)) # magic header f.write(_encode(num_images)) f.write(_encode(28)) f.write(_encode(28)) f.write(img.numpy().tobytes()) def _make_label_file(filename, num_images): labels = torch.zeros((num_images,), dtype=torch.uint8) with open(filename, "wb") as f: f.write(_encode(2049)) # magic header f.write(_encode(num_images)) f.write(labels.numpy().tobytes()) with get_tmp_dir() as tmp_dir: raw_dir = os.path.join(tmp_dir, cls_name, "raw") os.makedirs(raw_dir) _make_image_file(os.path.join(raw_dir, "train-images-idx3-ubyte"), num_images) _make_label_file(os.path.join(raw_dir, "train-labels-idx1-ubyte"), num_images) _make_image_file(os.path.join(raw_dir, "t10k-images-idx3-ubyte"), num_images) _make_label_file(os.path.join(raw_dir, "t10k-labels-idx1-ubyte"), num_images) yield tmp_dir @contextlib.contextmanager def cifar_root(version): def _get_version_params(version): if version == 'CIFAR10': return { 'base_folder': 'cifar-10-batches-py', 'train_files': ['data_batch_{}'.format(batch) for batch in range(1, 6)], 'test_file': 'test_batch', 'target_key': 'labels', 'meta_file': 'batches.meta', 'classes_key': 'label_names', } elif version == 'CIFAR100': return { 'base_folder': 'cifar-100-python', 'train_files': ['train'], 'test_file': 'test', 'target_key': 'fine_labels', 'meta_file': 'meta', 'classes_key': 'fine_label_names', } else: raise ValueError def _make_pickled_file(obj, file): with open(file, 'wb') as fh: pickle.dump(obj, fh, 2) def _make_data_file(file, target_key): obj = { 'data': np.zeros((1, 32 * 32 * 3), dtype=np.uint8), target_key: [0] } _make_pickled_file(obj, file) def _make_meta_file(file, classes_key): obj = { classes_key: ['fakedata'], } _make_pickled_file(obj, file) params = _get_version_params(version) with get_tmp_dir() as root: base_folder = os.path.join(root, params['base_folder']) os.mkdir(base_folder) for file in list(params['train_files']) + [params['test_file']]: _make_data_file(os.path.join(base_folder, file), params['target_key']) _make_meta_file(os.path.join(base_folder, params['meta_file']), params['classes_key']) yield root @contextlib.contextmanager def imagenet_root(): import scipy.io as sio WNID = 'n01234567' CLS = 'fakedata' def _make_image(file): PIL.Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)).save(file) def _make_tar(archive, content, arcname=None, compress=False): mode = 'w:gz' if compress else 'w' if arcname is None: arcname = os.path.basename(content) with tarfile.open(archive, mode) as fh: fh.add(content, arcname=arcname) def _make_train_archive(root): with get_tmp_dir() as tmp: wnid_dir = os.path.join(tmp, WNID) os.mkdir(wnid_dir) _make_image(os.path.join(wnid_dir, WNID + '_1.JPEG')) wnid_archive = wnid_dir + '.tar' _make_tar(wnid_archive, wnid_dir) train_archive = os.path.join(root, 'ILSVRC2012_img_train.tar') _make_tar(train_archive, wnid_archive) def _make_val_archive(root): with get_tmp_dir() as tmp: val_image = os.path.join(tmp, 'ILSVRC2012_val_00000001.JPEG') _make_image(val_image) val_archive = os.path.join(root, 'ILSVRC2012_img_val.tar') _make_tar(val_archive, val_image) def _make_devkit_archive(root): with get_tmp_dir() as tmp: data_dir = os.path.join(tmp, 'data') os.mkdir(data_dir) meta_file = os.path.join(data_dir, 'meta.mat') synsets = np.core.records.fromarrays([ (0.0, 1.0), (WNID, ''), (CLS, ''), ('fakedata for the torchvision testsuite', ''), (0.0, 1.0), ], names=['ILSVRC2012_ID', 'WNID', 'words', 'gloss', 'num_children']) sio.savemat(meta_file, {'synsets': synsets}) groundtruth_file = os.path.join(data_dir, 'ILSVRC2012_validation_ground_truth.txt') with open(groundtruth_file, 'w') as fh: fh.write('0\n') devkit_name = 'ILSVRC2012_devkit_t12' devkit_archive = os.path.join(root, devkit_name + '.tar.gz') _make_tar(devkit_archive, tmp, arcname=devkit_name, compress=True) with get_tmp_dir() as root: _make_train_archive(root) _make_val_archive(root) _make_devkit_archive(root) yield root @contextlib.contextmanager def widerface_root(): """ Generates a dataset with the following folder structure and returns the path root: └── widerface ├── wider_face_split ├── WIDER_train ├── WIDER_val └── WIDER_test The dataset consist of 1 image for each dataset split (train, val, test) and annotation files for each split """ def _make_image(file): PIL.Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)).save(file) def _make_train_archive(root): extracted_dir = os.path.join(root, 'WIDER_train', 'images', '0--Parade') os.makedirs(extracted_dir) _make_image(os.path.join(extracted_dir, '0_Parade_marchingband_1_1.jpg')) def _make_val_archive(root): extracted_dir = os.path.join(root, 'WIDER_val', 'images', '0--Parade') os.makedirs(extracted_dir) _make_image(os.path.join(extracted_dir, '0_Parade_marchingband_1_2.jpg')) def _make_test_archive(root): extracted_dir = os.path.join(root, 'WIDER_test', 'images', '0--Parade') os.makedirs(extracted_dir) _make_image(os.path.join(extracted_dir, '0_Parade_marchingband_1_3.jpg')) def _make_annotations_archive(root): train_bbox_contents = '0--Parade/0_Parade_marchingband_1_1.jpg\n1\n449 330 122 149 0 0 0 0 0 0\n' val_bbox_contents = '0--Parade/0_Parade_marchingband_1_2.jpg\n1\n501 160 285 443 0 0 0 0 0 0\n' test_filelist_contents = '0--Parade/0_Parade_marchingband_1_3.jpg\n' extracted_dir = os.path.join(root, 'wider_face_split') os.mkdir(extracted_dir) # bbox training file bbox_file = os.path.join(extracted_dir, "wider_face_train_bbx_gt.txt") with open(bbox_file, "w") as txt_file: txt_file.write(train_bbox_contents) # bbox validation file bbox_file = os.path.join(extracted_dir, "wider_face_val_bbx_gt.txt") with open(bbox_file, "w") as txt_file: txt_file.write(val_bbox_contents) # test filelist file filelist_file = os.path.join(extracted_dir, "wider_face_test_filelist.txt") with open(filelist_file, "w") as txt_file: txt_file.write(test_filelist_contents) with get_tmp_dir() as root: root_base = os.path.join(root, "widerface") os.mkdir(root_base) _make_train_archive(root_base) _make_val_archive(root_base) _make_test_archive(root_base) _make_annotations_archive(root_base) yield root @contextlib.contextmanager def cityscapes_root(): def _make_image(file): PIL.Image.fromarray(np.zeros((1024, 2048, 3), dtype=np.uint8)).save(file) def _make_regular_target(file): PIL.Image.fromarray(np.zeros((1024, 2048), dtype=np.uint8)).save(file) def _make_color_target(file): PIL.Image.fromarray(np.zeros((1024, 2048, 4), dtype=np.uint8)).save(file) def _make_polygon_target(file): polygon_example = { 'imgHeight': 1024, 'imgWidth': 2048, 'objects': [{'label': 'sky', 'polygon': [[1241, 0], [1234, 156], [1478, 197], [1611, 172], [1606, 0]]}, {'label': 'road', 'polygon': [[0, 448], [1331, 274], [1473, 265], [2047, 605], [2047, 1023], [0, 1023]]}]} with open(file, 'w') as outfile: json.dump(polygon_example, outfile) with get_tmp_dir() as tmp_dir: for mode in ['Coarse', 'Fine']: gt_dir = os.path.join(tmp_dir, 'gt%s' % mode) os.makedirs(gt_dir) if mode == 'Coarse': splits = ['train', 'train_extra', 'val'] else: splits = ['train', 'test', 'val'] for split in splits: split_dir = os.path.join(gt_dir, split) os.makedirs(split_dir) for city in ['bochum', 'bremen']: city_dir = os.path.join(split_dir, city) os.makedirs(city_dir) _make_color_target(os.path.join(city_dir, '{city}_000000_000000_gt{mode}_color.png'.format( city=city, mode=mode))) _make_regular_target(os.path.join(city_dir, '{city}_000000_000000_gt{mode}_instanceIds.png'.format( city=city, mode=mode))) _make_regular_target(os.path.join(city_dir, '{city}_000000_000000_gt{mode}_labelIds.png'.format( city=city, mode=mode))) _make_polygon_target(os.path.join(city_dir, '{city}_000000_000000_gt{mode}_polygons.json'.format( city=city, mode=mode))) # leftImg8bit dataset leftimg_dir = os.path.join(tmp_dir, 'leftImg8bit') os.makedirs(leftimg_dir) for split in ['test', 'train_extra', 'train', 'val']: split_dir = os.path.join(leftimg_dir, split) os.makedirs(split_dir) for city in ['bochum', 'bremen']: city_dir = os.path.join(split_dir, city) os.makedirs(city_dir) _make_image(os.path.join(city_dir, '{city}_000000_000000_leftImg8bit.png'.format(city=city))) yield tmp_dir @contextlib.contextmanager def svhn_root(): import scipy.io as sio def _make_mat(file): images = np.zeros((32, 32, 3, 2), dtype=np.uint8) targets = np.zeros((2,), dtype=np.uint8) sio.savemat(file, {'X': images, 'y': targets}) with get_tmp_dir() as root: _make_mat(os.path.join(root, "train_32x32.mat")) _make_mat(os.path.join(root, "test_32x32.mat")) _make_mat(os.path.join(root, "extra_32x32.mat")) yield root @contextlib.contextmanager def ucf101_root(): with get_tmp_dir() as tmp_dir: ucf_dir = os.path.join(tmp_dir, 'UCF-101') video_dir = os.path.join(ucf_dir, 'video') annotations = os.path.join(ucf_dir, 'annotations') os.makedirs(ucf_dir) os.makedirs(video_dir) os.makedirs(annotations) fold_files = [] for split in {'train', 'test'}: for fold in range(1, 4): fold_file = '{:s}list{:02d}.txt'.format(split, fold) fold_files.append(os.path.join(annotations, fold_file)) file_handles = [open(x, 'w') for x in fold_files] file_iter = cycle(file_handles) for i in range(0, 2): current_class = 'class_{0}'.format(i + 1) class_dir = os.path.join(video_dir, current_class) os.makedirs(class_dir) for group in range(0, 3): for clip in range(0, 4): # Save sample file clip_name = 'v_{0}_g{1}_c{2}.avi'.format( current_class, group, clip) clip_path = os.path.join(class_dir, clip_name) length = random.randrange(10, 21) this_clip = torch.randint( 0, 256, (length * 25, 320, 240, 3), dtype=torch.uint8) write_video(clip_path, this_clip, 25) # Add to annotations ann_file = next(file_iter) ann_file.write('{0}\n'.format( os.path.join(current_class, clip_name))) # Close all file descriptors for f in file_handles: f.close() yield (video_dir, annotations) @contextlib.contextmanager def places365_root(split="train-standard", small=False): VARIANTS = { "train-standard": "standard", "train-challenge": "challenge", "val": "standard", } # {split: file} DEVKITS = { "train-standard": "filelist_places365-standard.tar", "train-challenge": "filelist_places365-challenge.tar", "val": "filelist_places365-standard.tar", } CATEGORIES = "categories_places365.txt" # {split: file} FILE_LISTS = { "train-standard": "places365_train_standard.txt", "train-challenge": "places365_train_challenge.txt", "val": "places365_train_standard.txt", } # {(split, small): (archive, folder_default, folder_renamed)} IMAGES = { ("train-standard", False): ("train_large_places365standard.tar", "data_large", "data_large_standard"), ("train-challenge", False): ("train_large_places365challenge.tar", "data_large", "data_large_challenge"), ("val", False): ("val_large.tar", "val_large", "val_large"), ("train-standard", True): ("train_256_places365standard.tar", "data_256", "data_256_standard"), ("train-challenge", True): ("train_256_places365challenge.tar", "data_256", "data_256_challenge"), ("val", True): ("val_256.tar", "val_256", "val_256"), } # (class, idx) CATEGORIES_CONTENT = (("/a/airfield", 0), ("/a/apartment_building/outdoor", 8), ("/b/badlands", 30)) # (file, idx) FILE_LIST_CONTENT = ( ("Places365_val_00000001.png", 0), *((f"{category}/Places365_train_00000001.png", idx) for category, idx in CATEGORIES_CONTENT), ) def mock_target(attr, partial="torchvision.datasets.places365.Places365"): return f"{partial}.{attr}" def make_txt(root, name, seq): file = os.path.join(root, name) with open(file, "w") as fh: for string, idx in seq: fh.write(f"{string} {idx}\n") return name, compute_md5(file) def make_categories_txt(root, name): return make_txt(root, name, CATEGORIES_CONTENT) def make_file_list_txt(root, name): return make_txt(root, name, FILE_LIST_CONTENT) def make_image(file, size): os.makedirs(os.path.dirname(file), exist_ok=True) PIL.Image.fromarray(np.zeros((*size, 3), dtype=np.uint8)).save(file) def make_devkit_archive(stack, root, split): archive = DEVKITS[split] files = [] meta = make_categories_txt(root, CATEGORIES) mock_class_attribute(stack, mock_target("_CATEGORIES_META"), meta) files.append(meta[0]) meta = {split: make_file_list_txt(root, FILE_LISTS[split])} mock_class_attribute(stack, mock_target("_FILE_LIST_META"), meta) files.extend([item[0] for item in meta.values()]) meta = {VARIANTS[split]: make_tar(root, archive, *files)} mock_class_attribute(stack, mock_target("_DEVKIT_META"), meta) def make_images_archive(stack, root, split, small): archive, folder_default, folder_renamed = IMAGES[(split, small)] image_size = (256, 256) if small else (512, random.randint(512, 1024)) files, idcs = zip(*FILE_LIST_CONTENT) images = [file.lstrip("/").replace("/", os.sep) for file in files] for image in images: make_image(os.path.join(root, folder_default, image), image_size) meta = {(split, small): make_tar(root, archive, folder_default)} mock_class_attribute(stack, mock_target("_IMAGES_META"), meta) return [(os.path.join(root, folder_renamed, image), idx) for image, idx in zip(images, idcs)] with contextlib.ExitStack() as stack, get_tmp_dir() as root: make_devkit_archive(stack, root, split) class_to_idx = dict(CATEGORIES_CONTENT) classes = list(class_to_idx.keys()) data = {"class_to_idx": class_to_idx, "classes": classes} data["imgs"] = make_images_archive(stack, root, split, small) clean_dir(root, ".tar$") yield root, data @contextlib.contextmanager def stl10_root(_extracted=False): CLASS_NAMES = ("airplane", "bird") ARCHIVE_NAME = "stl10_binary" NUM_FOLDS = 10 def mock_target(attr, partial="torchvision.datasets.stl10.STL10"): return f"{partial}.{attr}" def make_binary_file(num_elements, root, name): file = os.path.join(root, name) np.zeros(num_elements, dtype=np.uint8).tofile(file) return name, compute_md5(file) def make_image_file(num_images, root, name, num_channels=3, height=96, width=96): return make_binary_file(num_images * num_channels * height * width, root, name) def make_label_file(num_images, root, name): return make_binary_file(num_images, root, name) def make_class_names_file(root, name="class_names.txt"): with open(os.path.join(root, name), "w") as fh: for name in CLASS_NAMES: fh.write(f"{name}\n") def make_fold_indices_file(root): offset = 0 with open(os.path.join(root, "fold_indices.txt"), "w") as fh: for fold in range(NUM_FOLDS): line = " ".join([str(idx) for idx in range(offset, offset + fold + 1)]) fh.write(f"{line}\n") offset += fold + 1 return tuple(range(1, NUM_FOLDS + 1)) def make_train_files(stack, root, num_unlabeled_images=1): num_images_in_fold = make_fold_indices_file(root) num_train_images = sum(num_images_in_fold) train_list = [ list(make_image_file(num_train_images, root, "train_X.bin")), list(make_label_file(num_train_images, root, "train_y.bin")), list(make_image_file(1, root, "unlabeled_X.bin")) ] mock_class_attribute(stack, target=mock_target("train_list"), new=train_list) return num_images_in_fold, dict(train=num_train_images, unlabeled=num_unlabeled_images) def make_test_files(stack, root, num_images=2): test_list = [ list(make_image_file(num_images, root, "test_X.bin")), list(make_label_file(num_images, root, "test_y.bin")), ] mock_class_attribute(stack, target=mock_target("test_list"), new=test_list) return dict(test=num_images) def make_archive(stack, root, name): archive, md5 = make_tar(root, name, name, compression="gz") mock_class_attribute(stack, target=mock_target("tgz_md5"), new=md5) return archive with contextlib.ExitStack() as stack, get_tmp_dir() as root: archive_folder = os.path.join(root, ARCHIVE_NAME) os.mkdir(archive_folder) num_images_in_folds, num_images_in_split = make_train_files(stack, archive_folder) num_images_in_split.update(make_test_files(stack, archive_folder)) make_class_names_file(archive_folder) archive = make_archive(stack, root, ARCHIVE_NAME) dir_util.remove_tree(archive_folder) data = dict(num_images_in_folds=num_images_in_folds, num_images_in_split=num_images_in_split, archive=archive) yield root, data