def _create_annotations(self, image_ids, num_annotations_per_image):
captions = [str(idx) for idx in range(num_annotations_per_image)]
annotations = datasets_utils.combinations_grid(image_id=image_ids,
caption=captions)
for id, annotation in enumerate(annotations):
annotation["id"] = id
return annotations, dict(captions=captions)
def _create_annotations(self, image_ids, num_annotations_per_image):
annotations = datasets_utils.combinations_grid(
image_id=image_ids,
bbox=([1.0, 2.0, 3.0, 4.0], ) * num_annotations_per_image)
for id, annotation in enumerate(annotations):
annotation["id"] = id
return annotations, dict()
class UCF101TestCase(datasets_utils.VideoDatasetTestCase):
DATASET_CLASS = datasets.UCF101
CONFIGS = datasets_utils.combinations_grid(fold=(1, 2, 3),
train=(True, False))
def inject_fake_data(self, tmpdir, config):
tmpdir = pathlib.Path(tmpdir)
video_folder = tmpdir / "videos"
os.makedirs(video_folder)
video_files = self._create_videos(video_folder)
annotations_folder = annotations_folder = tmpdir / "annotations"
os.makedirs(annotations_folder)
num_examples = self._create_annotation_files(annotations_folder,
video_files,
config["fold"],
config["train"])
return (str(video_folder), str(annotations_folder)), num_examples
def _create_videos(self, root, num_examples_per_class=3):
def file_name_fn(cls, idx, clips_per_group=2):
return f"v_{cls}_g{(idx // clips_per_group) + 1:02d}_c{(idx % clips_per_group) + 1:02d}.avi"
video_files = [
datasets_utils.create_video_folder(
root, cls, lambda idx: file_name_fn(cls, idx),
num_examples_per_class) for cls in ("ApplyEyeMakeup", "YoYo")
]
return [
path.relative_to(root) for path in itertools.chain(*video_files)
]
def _create_annotation_files(self, root, video_files, fold, train):
current_videos = random.sample(
video_files, random.randrange(1,
len(video_files) - 1))
current_annotation = self._annotation_file_name(fold, train)
self._create_annotation_file(root, current_annotation, current_videos)
other_videos = set(video_files) - set(current_videos)
other_annotations = [
self._annotation_file_name(fold, train)
for fold, train in itertools.product((1, 2, 3), (True, False))
]
other_annotations.remove(current_annotation)
for name in other_annotations:
self._create_annotation_file(root, name, other_videos)
return len(current_videos)
def _annotation_file_name(self, fold, train):
return f"{'train' if train else 'test'}list{fold:02d}.txt"
def _create_annotation_file(self, root, name, video_files):
with open(pathlib.Path(root) / name, "w") as fh:
fh.writelines(f"{file}\n" for file in sorted(video_files))
class VOCSegmentationTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.VOCSegmentation
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image)
CONFIGS = (
*datasets_utils.combinations_grid(
year=[f"20{year:02d}" for year in range(7, 13)],
image_set=("train", "val", "trainval")),
dict(year="2007", image_set="test"),
dict(year="2007-test", image_set="test"),
)
def inject_fake_data(self, tmpdir, config):
year, is_test_set = (("2007", True) if config["year"] == "2007-test"
or config["image_set"] == "test" else
(config["year"], False))
image_set = config["image_set"]
base_dir = pathlib.Path(tmpdir)
if year == "2011":
base_dir /= "TrainVal"
base_dir = base_dir / "VOCdevkit" / f"VOC{year}"
os.makedirs(base_dir)
num_images, num_images_per_image_set = self._create_image_set_files(
base_dir, "ImageSets", is_test_set)
datasets_utils.create_image_folder(base_dir, "JPEGImages",
lambda idx: f"{idx:06d}.jpg",
num_images)
datasets_utils.create_image_folder(base_dir, "SegmentationClass",
lambda idx: f"{idx:06d}.png",
num_images)
annotation = self._create_annotation_files(base_dir, "Annotations",
num_images)
return dict(num_examples=num_images_per_image_set[image_set],
annotation=annotation)
def _create_image_set_files(self, root, name, is_test_set):
root = pathlib.Path(root) / name
src = pathlib.Path(root) / "Main"
os.makedirs(src, exist_ok=True)
idcs = dict(train=(0, 1, 2), val=(3, 4), test=(5, ))
idcs["trainval"] = (*idcs["train"], *idcs["val"])
for image_set in ("test", ) if is_test_set else ("train", "val",
"trainval"):
self._create_image_set_file(src, image_set, idcs[image_set])
shutil.copytree(src, root / "Segmentation")
num_images = max(itertools.chain(*idcs.values())) + 1
num_images_per_image_set = dict([(image_set, len(idcs_))
for image_set, idcs_ in idcs.items()])
return num_images, num_images_per_image_set
def _create_image_set_file(self, root, image_set, idcs):
with open(pathlib.Path(root) / f"{image_set}.txt", "w") as fh:
fh.writelines([f"{idx:06d}\n" for idx in idcs])
def _create_annotation_files(self, root, name, num_images):
root = pathlib.Path(root) / name
os.makedirs(root)
for idx in range(num_images):
annotation = self._create_annotation_file(root, f"{idx:06d}.xml")
return annotation
def _create_annotation_file(self, root, name):
def add_child(parent, name, text=None):
child = ET.SubElement(parent, name)
child.text = text
return child
def add_name(obj, name="dog"):
add_child(obj, "name", name)
return name
def add_bndbox(obj, bndbox=None):
if bndbox is None:
bndbox = {"xmin": "1", "xmax": "2", "ymin": "3", "ymax": "4"}
obj = add_child(obj, "bndbox")
for name, text in bndbox.items():
add_child(obj, name, text)
return bndbox
annotation = ET.Element("annotation")
obj = add_child(annotation, "object")
data = dict(name=add_name(obj), bndbox=add_bndbox(obj))
with open(pathlib.Path(root) / name, "wb") as fh:
fh.write(ET.tostring(annotation))
return data
class CelebATestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.CelebA
FEATURE_TYPES = (PIL.Image.Image, (torch.Tensor, int, tuple, type(None)))
CONFIGS = datasets_utils.combinations_grid(
split=("train", "valid", "test", "all"),
target_type=("attr", "identity", "bbox", "landmarks",
["attr", "identity"]),
)
REQUIRED_PACKAGES = ("pandas", )
_SPLIT_TO_IDX = dict(train=0, valid=1, test=2)
def inject_fake_data(self, tmpdir, config):
base_folder = pathlib.Path(tmpdir) / "celeba"
os.makedirs(base_folder)
num_images, num_images_per_split = self._create_split_txt(base_folder)
datasets_utils.create_image_folder(base_folder, "img_align_celeba",
lambda idx: f"{idx + 1:06d}.jpg",
num_images)
attr_names = self._create_attr_txt(base_folder, num_images)
self._create_identity_txt(base_folder, num_images)
self._create_bbox_txt(base_folder, num_images)
self._create_landmarks_txt(base_folder, num_images)
return dict(num_examples=num_images_per_split[config["split"]],
attr_names=attr_names)
def _create_split_txt(self, root):
num_images_per_split = dict(train=3, valid=2, test=1)
data = [[self._SPLIT_TO_IDX[split]]
for split, num_images in num_images_per_split.items()
for _ in range(num_images)]
self._create_txt(root, "list_eval_partition.txt", data)
num_images_per_split["all"] = num_images = sum(
num_images_per_split.values())
return num_images, num_images_per_split
def _create_attr_txt(self, root, num_images):
header = ("5_o_Clock_Shadow", "Young")
data = torch.rand(
(num_images, len(header))).ge(0.5).int().mul(2).sub(1).tolist()
self._create_txt(root,
"list_attr_celeba.txt",
data,
header=header,
add_num_examples=True)
return header
def _create_identity_txt(self, root, num_images):
data = torch.randint(1, 4, size=(num_images, 1)).tolist()
self._create_txt(root, "identity_CelebA.txt", data)
def _create_bbox_txt(self, root, num_images):
header = ("x_1", "y_1", "width", "height")
data = torch.randint(10, size=(num_images, len(header))).tolist()
self._create_txt(root,
"list_bbox_celeba.txt",
data,
header=header,
add_num_examples=True,
add_image_id_to_header=True)
def _create_landmarks_txt(self, root, num_images):
header = ("lefteye_x", "rightmouth_y")
data = torch.randint(10, size=(num_images, len(header))).tolist()
self._create_txt(root,
"list_landmarks_align_celeba.txt",
data,
header=header,
add_num_examples=True)
def _create_txt(self,
root,
name,
data,
header=None,
add_num_examples=False,
add_image_id_to_header=False):
with open(pathlib.Path(root) / name, "w") as fh:
if add_num_examples:
fh.write(f"{len(data)}\n")
if header:
if add_image_id_to_header:
header = ("image_id", *header)
fh.write(f"{' '.join(header)}\n")
for idx, line in enumerate(data, 1):
fh.write(
f"{' '.join((f'{idx:06d}.jpg', *[str(value) for value in line]))}\n"
)
def test_combined_targets(self):
target_types = ["attr", "identity", "bbox", "landmarks"]
individual_targets = []
for target_type in target_types:
with self.create_dataset(target_type=target_type) as (dataset, _):
_, target = dataset[0]
individual_targets.append(target)
with self.create_dataset(target_type=target_types) as (dataset, _):
_, combined_targets = dataset[0]
actual = len(individual_targets)
expected = len(combined_targets)
self.assertEqual(
actual,
expected,
f"The number of the returned combined targets does not match the the number targets if requested "
f"individually: {actual} != {expected}",
)
for target_type, combined_target, individual_target in zip(
target_types, combined_targets, individual_targets):
with self.subTest(target_type=target_type):
actual = type(combined_target)
expected = type(individual_target)
self.assertIs(
actual,
expected,
f"Type of the combined target does not match the type of the corresponding individual target: "
f"{actual} is not {expected}",
)
def test_no_target(self):
with self.create_dataset(target_type=[]) as (dataset, _):
_, target = dataset[0]
self.assertIsNone(target)
def test_attr_names(self):
with self.create_dataset() as (dataset, info):
self.assertEqual(tuple(dataset.attr_names), info["attr_names"])
class CIFAR10TestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.CIFAR10
CONFIGS = datasets_utils.combinations_grid(train=(True, False))
_VERSION_CONFIG = dict(
base_folder="cifar-10-batches-py",
train_files=tuple(f"data_batch_{idx}" for idx in range(1, 6)),
test_files=("test_batch", ),
labels_key="labels",
meta_file="batches.meta",
num_categories=10,
categories_key="label_names",
)
def inject_fake_data(self, tmpdir, config):
tmpdir = pathlib.Path(tmpdir) / self._VERSION_CONFIG["base_folder"]
os.makedirs(tmpdir)
num_images_per_file = 1
for name in itertools.chain(self._VERSION_CONFIG["train_files"],
self._VERSION_CONFIG["test_files"]):
self._create_batch_file(tmpdir, name, num_images_per_file)
categories = self._create_meta_file(tmpdir)
return dict(
num_examples=num_images_per_file *
len(self._VERSION_CONFIG["train_files"]
if config["train"] else self._VERSION_CONFIG["test_files"]),
categories=categories,
)
def _create_batch_file(self, root, name, num_images):
data = datasets_utils.create_image_or_video_tensor(
(num_images, 32 * 32 * 3))
labels = np.random.randint(0,
self._VERSION_CONFIG["num_categories"],
size=num_images).tolist()
self._create_binary_file(root, name, {
"data": data,
self._VERSION_CONFIG["labels_key"]: labels
})
def _create_meta_file(self, root):
categories = [
f"{idx:0{len(str(self._VERSION_CONFIG['num_categories'] - 1))}d}"
for idx in range(self._VERSION_CONFIG["num_categories"])
]
self._create_binary_file(
root, self._VERSION_CONFIG["meta_file"],
{self._VERSION_CONFIG["categories_key"]: categories})
return categories
def _create_binary_file(self, root, name, content):
with open(pathlib.Path(root) / name, "wb") as fh:
pickle.dump(content, fh)
def test_class_to_idx(self):
with self.create_dataset() as (dataset, info):
expected = {
category: label
for label, category in enumerate(info["categories"])
}
actual = dataset.class_to_idx
self.assertEqual(actual, expected)
class Caltech101TestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Caltech101
FEATURE_TYPES = (PIL.Image.Image, (int, np.ndarray, tuple))
CONFIGS = datasets_utils.combinations_grid(
target_type=("category", "annotation", ["category", "annotation"]))
REQUIRED_PACKAGES = ("scipy", )
def inject_fake_data(self, tmpdir, config):
root = pathlib.Path(tmpdir) / "caltech101"
images = root / "101_ObjectCategories"
annotations = root / "Annotations"
categories = (("Faces", "Faces_2"), ("helicopter", "helicopter"),
("ying_yang", "ying_yang"))
num_images_per_category = 2
for image_category, annotation_category in categories:
datasets_utils.create_image_folder(
root=images,
name=image_category,
file_name_fn=lambda idx: f"image_{idx + 1:04d}.jpg",
num_examples=num_images_per_category,
)
self._create_annotation_folder(
root=annotations,
name=annotation_category,
file_name_fn=lambda idx: f"annotation_{idx + 1:04d}.mat",
num_examples=num_images_per_category,
)
# This is included in the original archive, but is removed by the dataset. Thus, an empty directory suffices.
os.makedirs(images / "BACKGROUND_Google")
return num_images_per_category * len(categories)
def _create_annotation_folder(self, root, name, file_name_fn,
num_examples):
root = pathlib.Path(root) / name
os.makedirs(root)
for idx in range(num_examples):
self._create_annotation_file(root, file_name_fn(idx))
def _create_annotation_file(self, root, name):
mdict = dict(obj_contour=torch.rand((2, torch.randint(3, 6, size=())),
dtype=torch.float64).numpy())
datasets_utils.lazy_importer.scipy.io.savemat(
str(pathlib.Path(root) / name), mdict)
def test_combined_targets(self):
target_types = ["category", "annotation"]
individual_targets = []
for target_type in target_types:
with self.create_dataset(target_type=target_type) as (dataset, _):
_, target = dataset[0]
individual_targets.append(target)
with self.create_dataset(target_type=target_types) as (dataset, _):
_, combined_targets = dataset[0]
actual = len(individual_targets)
expected = len(combined_targets)
self.assertEqual(
actual,
expected,
f"The number of the returned combined targets does not match the the number targets if requested "
f"individually: {actual} != {expected}",
)
for target_type, combined_target, individual_target in zip(
target_types, combined_targets, individual_targets):
with self.subTest(target_type=target_type):
actual = type(combined_target)
expected = type(individual_target)
self.assertIs(
actual,
expected,
f"Type of the combined target does not match the type of the corresponding individual target: "
f"{actual} is not {expected}",
)
class LSUNTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.LSUN
REQUIRED_PACKAGES = ("lmdb", )
CONFIGS = datasets_utils.combinations_grid(
classes=("train", "test", "val",
["bedroom_train", "church_outdoor_train"]))
_CATEGORIES = (
"bedroom",
"bridge",
"church_outdoor",
"classroom",
"conference_room",
"dining_room",
"kitchen",
"living_room",
"restaurant",
"tower",
)
def inject_fake_data(self, tmpdir, config):
root = pathlib.Path(tmpdir)
num_images = 0
for cls in self._parse_classes(config["classes"]):
num_images += self._create_lmdb(root, cls)
return num_images
@contextlib.contextmanager
def create_dataset(self, *args, **kwargs):
with super().create_dataset(*args, **kwargs) as output:
yield output
# Currently datasets.LSUN caches the keys in the current directory rather than in the root directory. Thus,
# this creates a number of unique _cache_* files in the current directory that will not be removed together
# with the temporary directory
for file in os.listdir(os.getcwd()):
if file.startswith("_cache_"):
os.remove(file)
def _parse_classes(self, classes):
if not isinstance(classes, str):
return classes
split = classes
if split == "test":
return [split]
return [f"{category}_{split}" for category in self._CATEGORIES]
def _create_lmdb(self, root, cls):
lmdb = datasets_utils.lazy_importer.lmdb
hexdigits_lowercase = string.digits + string.ascii_lowercase[:6]
folder = f"{cls}_lmdb"
num_images = torch.randint(1, 4, size=()).item()
format = "webp"
files = datasets_utils.create_image_folder(
root, folder, lambda idx: f"{idx}.{format}", num_images)
with lmdb.open(str(root /
folder)) as env, env.begin(write=True) as txn:
for file in files:
key = "".join(
random.choice(hexdigits_lowercase)
for _ in range(40)).encode()
buffer = io.BytesIO()
Image.open(file).save(buffer, format)
buffer.seek(0)
value = buffer.read()
txn.put(key, value)
os.remove(file)
return num_images
