Skip to content

GitLab

Commit bf491463 authored by limm's avatar limm
Browse files

add v0.19.1 release

parent e17f5ea2
Hide whitespace changes
Inline Side-by-side
Showing with 3007 additions and 1700 deletions
test/fakedata_generation.py deleted 100644 → 0
View file @ e17f5ea2
import os
import contextlib
import hashlib
import pickle
import re
import tarfile
import unittest.mock
from distutils import dir_util
import numpy as np
import PIL
import torch
from common_utils import get_tmp_dir
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 widerface_root():
"""
Generates a dataset with the following folder structure and returns the path root:
<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
test/optests_failures_dict.json 0 → 100644
View file @ bf491463
{
"_description": "This is a dict containing failures for tests autogenerated by generate_opcheck_tests. For more details, please see https://docs.google.com/document/d/1Pj5HRZvdOq3xpFpbEjUZp2hBovhy7Wnxw14m6lF2154/edit",
"_version": 1,
"data": {}
}
test/preprocess-bench.py
View file @ bf491463
import argparse import argparse
import os import os
from timeit import default_timer as timer from timeit import default_timer as timer
from torch.utils.model_zoo import tqdm
import torch import torch
import torch.utils.data import torch.utils.data
import torchvision import torchvision
import torchvision.transforms as transforms
import torchvision.datasets as datasets import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.model_zoo import tqdm
parser = argparse.ArgumentParser(description='PyTorch ImageNet Training') parser = argparse.ArgumentParser(description="PyTorch ImageNet Training")
parser.add_argument('--data', metavar='PATH', required=True, parser.add_argument("--data", metavar="PATH", required=True, help="path to dataset")
help='path to dataset') parser.add_argument(
parser.add_argument('--nThreads', '-j', default=2, type=int, metavar='N', "--nThreads", "-j", default=2, type=int, metavar="N", help="number of data loading threads (default: 2)"
help='number of data loading threads (default: 2)') )
parser.add_argument('--batchSize', '-b', default=256, type=int, metavar='N', parser.add_argument(
help='mini-batch size (1 = pure stochastic) Default: 256') "--batchSize", "-b", default=256, type=int, metavar="N", help="mini-batch size (1 = pure stochastic) Default: 256"
parser.add_argument('--accimage', action='store_true', )
help='use accimage') parser.add_argument("--accimage", action="store_true", help="use accimage")
if __name__ == "__main__": if __name__ == "__main__":
args = parser.parse_args() args = parser.parse_args()
if args.accimage: if args.accimage:
torchvision.set_image_backend('accimage') torchvision.set_image_backend("accimage")
print('Using {}'.format(torchvision.get_image_backend())) print(f"Using {torchvision.get_image_backend()}")
# Data loading code # Data loading code
transform = transforms.Compose([ transform = transforms.Compose(
transforms.RandomSizedCrop(224), [
transforms.RandomHorizontalFlip(), transforms.RandomSizedCrop(224),
transforms.ToTensor(), transforms.RandomHorizontalFlip(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], transforms.PILToTensor(),
std=[0.229, 0.224, 0.225]), transforms.ConvertImageDtype(torch.float),
]) transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
traindir = os.path.join(args.data, 'train') )
valdir = os.path.join(args.data, 'val')
traindir = os.path.join(args.data, "train")
valdir = os.path.join(args.data, "val")
train = datasets.ImageFolder(traindir, transform) train = datasets.ImageFolder(traindir, transform)
val = datasets.ImageFolder(valdir, transform) val = datasets.ImageFolder(valdir, transform)
train_loader = torch.utils.data.DataLoader( train_loader = torch.utils.data.DataLoader(
train, batch_size=args.batchSize, shuffle=True, num_workers=args.nThreads) train, batch_size=args.batchSize, shuffle=True, num_workers=args.nThreads
)
train_iter = iter(train_loader) train_iter = iter(train_loader)
start_time = timer() start_time = timer()
...@@ -51,9 +55,12 @@ if __name__ == "__main__": ...@@ -51,9 +55,12 @@ if __name__ == "__main__":
pbar.update(1) pbar.update(1)
batch = next(train_iter) batch = next(train_iter)
end_time = timer() end_time = timer()
print("Performance: {dataset:.0f} minutes/dataset, {batch:.1f} ms/batch," print(
" {image:.2f} ms/image {rate:.0f} images/sec" "Performance: {dataset:.0f} minutes/dataset, {batch:.1f} ms/batch,"
.format(dataset=(end_time - start_time) * (float(len(train_loader)) / batch_count / 60.0), " {image:.2f} ms/image {rate:.0f} images/sec".format(
batch=(end_time - start_time) / float(batch_count) * 1.0e+3, dataset=(end_time - start_time) * (float(len(train_loader)) / batch_count / 60.0),
image=(end_time - start_time) / (batch_count * args.batchSize) * 1.0e+3, batch=(end_time - start_time) / float(batch_count) * 1.0e3,
rate=(batch_count * args.batchSize) / (end_time - start_time))) image=(end_time - start_time) / (batch_count * args.batchSize) * 1.0e3,
rate=(batch_count * args.batchSize) / (end_time - start_time),
)
)
test/sanity_checks.ipynb deleted 100644 → 0
View file @ e17f5ea2
This source diff could not be displayed because it is too large. You can view the blob instead.
test/smoke_test.py
View file @ bf491463
"""Run smoke tests"""
import sys
from pathlib import Path
import torch import torch
import torchvision import torchvision
import torchvision.datasets as dset from torchvision.io import decode_jpeg, read_file, read_image
import torchvision.transforms from torchvision.models import resnet50, ResNet50_Weights
SCRIPT_DIR = Path(__file__).parent
def smoke_test_torchvision() -> None:
print(
"Is torchvision usable?",
all(x is not None for x in [torch.ops.image.decode_png, torch.ops.torchvision.roi_align]),
)
def smoke_test_torchvision_read_decode() -> None:
img_jpg = read_image(str(SCRIPT_DIR / "assets" / "encode_jpeg" / "grace_hopper_517x606.jpg"))
if img_jpg.shape != (3, 606, 517):
raise RuntimeError(f"Unexpected shape of img_jpg: {img_jpg.shape}")
img_png = read_image(str(SCRIPT_DIR / "assets" / "interlaced_png" / "wizard_low.png"))
if img_png.shape != (4, 471, 354):
raise RuntimeError(f"Unexpected shape of img_png: {img_png.shape}")
def smoke_test_torchvision_decode_jpeg(device: str = "cpu"):
img_jpg_data = read_file(str(SCRIPT_DIR / "assets" / "encode_jpeg" / "grace_hopper_517x606.jpg"))
img_jpg = decode_jpeg(img_jpg_data, device=device)
if img_jpg.shape != (3, 606, 517):
raise RuntimeError(f"Unexpected shape of img_jpg: {img_jpg.shape}")
def smoke_test_compile() -> None:
try:
model = resnet50().cuda()
model = torch.compile(model)
x = torch.randn(1, 3, 224, 224, device="cuda")
out = model(x)
print(f"torch.compile model output: {out.shape}")
except RuntimeError:
if sys.platform == "win32":
print("Successfully caught torch.compile RuntimeError on win")
else:
raise
def smoke_test_torchvision_resnet50_classify(device: str = "cpu") -> None:
img = read_image(str(SCRIPT_DIR / ".." / "gallery" / "assets" / "dog2.jpg")).to(device)
# Step 1: Initialize model with the best available weights
weights = ResNet50_Weights.DEFAULT
model = resnet50(weights=weights, progress=False).to(device)
model.eval()
# Step 2: Initialize the inference transforms
preprocess = weights.transforms(antialias=(device != "mps")) # antialias not supported on MPS
# Step 3: Apply inference preprocessing transforms
batch = preprocess(img).unsqueeze(0)
# Step 4: Use the model and print the predicted category
prediction = model(batch).squeeze(0).softmax(0)
class_id = prediction.argmax().item()
score = prediction[class_id].item()
category_name = weights.meta["categories"][class_id]
expected_category = "German shepherd"
print(f"{category_name} ({device}): {100 * score:.1f}%")
if category_name != expected_category:
raise RuntimeError(f"Failed ResNet50 classify {category_name} Expected: {expected_category}")
def main() -> None:
print(f"torchvision: {torchvision.__version__}")
print(f"torch.cuda.is_available: {torch.cuda.is_available()}")
# Turn 1.11.0aHASH into 1.11 (major.minor only)
version = ".".join(torchvision.__version__.split(".")[:2])
if version >= "0.16":
print(f"{torch.ops.image._jpeg_version() = }")
assert torch.ops.image._is_compiled_against_turbo()
smoke_test_torchvision()
smoke_test_torchvision_read_decode()
smoke_test_torchvision_resnet50_classify()
smoke_test_torchvision_decode_jpeg()
if torch.cuda.is_available():
smoke_test_torchvision_decode_jpeg("cuda")
smoke_test_torchvision_resnet50_classify("cuda")
# TODO: remove once pytorch/pytorch#110436 is resolved
if sys.version_info < (3, 12, 0):
smoke_test_compile()
if torch.backends.mps.is_available():
smoke_test_torchvision_resnet50_classify("mps")
if __name__ == "__main__":
main()
test/test_architecture_ops.py 0 → 100644
View file @ bf491463
import unittest
import pytest
import torch
from torchvision.models.maxvit import SwapAxes, WindowDepartition, WindowPartition
class MaxvitTester(unittest.TestCase):
def test_maxvit_window_partition(self):
input_shape = (1, 3, 224, 224)
partition_size = 7
n_partitions = input_shape[3] // partition_size
x = torch.randn(input_shape)
partition = WindowPartition()
departition = WindowDepartition()
x_hat = partition(x, partition_size)
x_hat = departition(x_hat, partition_size, n_partitions, n_partitions)
torch.testing.assert_close(x, x_hat)
def test_maxvit_grid_partition(self):
input_shape = (1, 3, 224, 224)
partition_size = 7
n_partitions = input_shape[3] // partition_size
x = torch.randn(input_shape)
pre_swap = SwapAxes(-2, -3)
post_swap = SwapAxes(-2, -3)
partition = WindowPartition()
departition = WindowDepartition()
x_hat = partition(x, n_partitions)
x_hat = pre_swap(x_hat)
x_hat = post_swap(x_hat)
x_hat = departition(x_hat, n_partitions, partition_size, partition_size)
torch.testing.assert_close(x, x_hat)
if __name__ == "__main__":
pytest.main([__file__])
test/test_backbone_utils.py
View file @ bf491463
import unittest import random
from itertools import chain
from typing import Mapping, Sequence
import pytest
import torch import torch
from torchvision.models.detection.backbone_utils import resnet_fpn_backbone from common_utils import set_rng_seed
from torchvision import models
from torchvision.models._utils import IntermediateLayerGetter
class ResnetFPNBackboneTester(unittest.TestCase): from torchvision.models.detection.backbone_utils import BackboneWithFPN, mobilenet_backbone, resnet_fpn_backbone
@classmethod from torchvision.models.feature_extraction import create_feature_extractor, get_graph_node_names
def setUpClass(cls):
cls.dtype = torch.float32
@pytest.mark.parametrize("backbone_name", ("resnet18", "resnet50"))
def test_resnet18_fpn_backbone(self): def test_resnet_fpn_backbone(backbone_name):
device = torch.device('cpu') x = torch.rand(1, 3, 300, 300, dtype=torch.float32, device="cpu")
x = torch.rand(1, 3, 300, 300, dtype=self.dtype, device=device) model = resnet_fpn_backbone(backbone_name=backbone_name, weights=None)
resnet18_fpn = resnet_fpn_backbone(backbone_name='resnet18', pretrained=False) assert isinstance(model, BackboneWithFPN)
y = resnet18_fpn(x) y = model(x)
self.assertEqual(list(y.keys()), ['0', '1', '2', '3', 'pool']) assert list(y.keys()) == ["0", "1", "2", "3", "pool"]
def test_resnet50_fpn_backbone(self): with pytest.raises(ValueError, match=r"Trainable layers should be in the range"):
device = torch.device('cpu') resnet_fpn_backbone(backbone_name=backbone_name, weights=None, trainable_layers=6)
x = torch.rand(1, 3, 300, 300, dtype=self.dtype, device=device) with pytest.raises(ValueError, match=r"Each returned layer should be in the range"):
resnet50_fpn = resnet_fpn_backbone(backbone_name='resnet50', pretrained=False) resnet_fpn_backbone(backbone_name=backbone_name, weights=None, returned_layers=[0, 1, 2, 3])
y = resnet50_fpn(x) with pytest.raises(ValueError, match=r"Each returned layer should be in the range"):
self.assertEqual(list(y.keys()), ['0', '1', '2', '3', 'pool']) resnet_fpn_backbone(backbone_name=backbone_name, weights=None, returned_layers=[2, 3, 4, 5])
@pytest.mark.parametrize("backbone_name", ("mobilenet_v2", "mobilenet_v3_large", "mobilenet_v3_small"))
def test_mobilenet_backbone(backbone_name):
with pytest.raises(ValueError, match=r"Trainable layers should be in the range"):
mobilenet_backbone(backbone_name=backbone_name, weights=None, fpn=False, trainable_layers=-1)
with pytest.raises(ValueError, match=r"Each returned layer should be in the range"):
mobilenet_backbone(backbone_name=backbone_name, weights=None, fpn=True, returned_layers=[-1, 0, 1, 2])
with pytest.raises(ValueError, match=r"Each returned layer should be in the range"):
mobilenet_backbone(backbone_name=backbone_name, weights=None, fpn=True, returned_layers=[3, 4, 5, 6])
model_fpn = mobilenet_backbone(backbone_name=backbone_name, weights=None, fpn=True)
assert isinstance(model_fpn, BackboneWithFPN)
model = mobilenet_backbone(backbone_name=backbone_name, weights=None, fpn=False)
assert isinstance(model, torch.nn.Sequential)
# Needed by TestFxFeatureExtraction.test_leaf_module_and_function
def leaf_function(x):
return int(x)
# Needed by TestFXFeatureExtraction. Checking that node naming conventions
# are respected. Particularly the index postfix of repeated node names
class TestSubModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.relu = torch.nn.ReLU()
def forward(self, x):
x = x + 1
x = x + 1
x = self.relu(x)
x = self.relu(x)
return x
class TestModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.submodule = TestSubModule()
self.relu = torch.nn.ReLU()
def forward(self, x):
x = self.submodule(x)
x = x + 1
x = x + 1
x = self.relu(x)
x = self.relu(x)
return x
test_module_nodes = [
"x",
"submodule.add",
"submodule.add_1",
"submodule.relu",
"submodule.relu_1",
"add",
"add_1",
"relu",
"relu_1",
]
class TestFxFeatureExtraction:
inp = torch.rand(1, 3, 224, 224, dtype=torch.float32, device="cpu")
model_defaults = {"num_classes": 1}
leaf_modules = []
def _create_feature_extractor(self, *args, **kwargs):
"""
Apply leaf modules
"""
tracer_kwargs = {}
if "tracer_kwargs" not in kwargs:
tracer_kwargs = {"leaf_modules": self.leaf_modules}
else:
tracer_kwargs = kwargs.pop("tracer_kwargs")
return create_feature_extractor(*args, **kwargs, tracer_kwargs=tracer_kwargs, suppress_diff_warning=True)
def _get_return_nodes(self, model):
set_rng_seed(0)
exclude_nodes_filter = [
"getitem",
"floordiv",
"size",
"chunk",
"_assert",
"eq",
"dim",
"getattr",
]
train_nodes, eval_nodes = get_graph_node_names(
model, tracer_kwargs={"leaf_modules": self.leaf_modules}, suppress_diff_warning=True
)
# Get rid of any nodes that don't return tensors as they cause issues
# when testing backward pass.
train_nodes = [n for n in train_nodes if not any(x in n for x in exclude_nodes_filter)]
eval_nodes = [n for n in eval_nodes if not any(x in n for x in exclude_nodes_filter)]
return random.sample(train_nodes, 10), random.sample(eval_nodes, 10)
@pytest.mark.parametrize("model_name", models.list_models(models))
def test_build_fx_feature_extractor(self, model_name):
set_rng_seed(0)
model = models.get_model(model_name, **self.model_defaults).eval()
train_return_nodes, eval_return_nodes = self._get_return_nodes(model)
# Check that it works with both a list and dict for return nodes
self._create_feature_extractor(
model, train_return_nodes={v: v for v in train_return_nodes}, eval_return_nodes=eval_return_nodes
)
self._create_feature_extractor(
model, train_return_nodes=train_return_nodes, eval_return_nodes=eval_return_nodes
)
# Check must specify return nodes
with pytest.raises(ValueError):
self._create_feature_extractor(model)
# Check return_nodes and train_return_nodes / eval_return nodes
# mutual exclusivity
with pytest.raises(ValueError):
self._create_feature_extractor(
model, return_nodes=train_return_nodes, train_return_nodes=train_return_nodes
)
# Check train_return_nodes / eval_return nodes must both be specified
with pytest.raises(ValueError):
self._create_feature_extractor(model, train_return_nodes=train_return_nodes)
# Check invalid node name raises ValueError
with pytest.raises(ValueError):
# First just double check that this node really doesn't exist
if not any(n.startswith("l") or n.startswith("l.") for n in chain(train_return_nodes, eval_return_nodes)):
self._create_feature_extractor(model, train_return_nodes=["l"], eval_return_nodes=["l"])
else: # otherwise skip this check
raise ValueError
def test_node_name_conventions(self):
model = TestModule()
train_nodes, _ = get_graph_node_names(model)
assert all(a == b for a, b in zip(train_nodes, test_module_nodes))
@pytest.mark.parametrize("model_name", models.list_models(models))
def test_forward_backward(self, model_name):
model = models.get_model(model_name, **self.model_defaults).train()
train_return_nodes, eval_return_nodes = self._get_return_nodes(model)
model = self._create_feature_extractor(
model, train_return_nodes=train_return_nodes, eval_return_nodes=eval_return_nodes
)
out = model(self.inp)
out_agg = 0
for node_out in out.values():
if isinstance(node_out, Sequence):
out_agg += sum(o.float().mean() for o in node_out if o is not None)
elif isinstance(node_out, Mapping):
out_agg += sum(o.float().mean() for o in node_out.values() if o is not None)
else:
# Assume that the only other alternative at this point is a Tensor
out_agg += node_out.float().mean()
out_agg.backward()
def test_feature_extraction_methods_equivalence(self):
model = models.resnet18(**self.model_defaults).eval()
return_layers = {"layer1": "layer1", "layer2": "layer2", "layer3": "layer3", "layer4": "layer4"}
ilg_model = IntermediateLayerGetter(model, return_layers).eval()
fx_model = self._create_feature_extractor(model, return_layers)
# Check that we have same parameters
for (n1, p1), (n2, p2) in zip(ilg_model.named_parameters(), fx_model.named_parameters()):
assert n1 == n2
assert p1.equal(p2)
# And that outputs match
with torch.no_grad():
ilg_out = ilg_model(self.inp)
fgn_out = fx_model(self.inp)
assert all(k1 == k2 for k1, k2 in zip(ilg_out.keys(), fgn_out.keys()))
for k in ilg_out.keys():
assert ilg_out[k].equal(fgn_out[k])
@pytest.mark.parametrize("model_name", models.list_models(models))
def test_jit_forward_backward(self, model_name):
set_rng_seed(0)
model = models.get_model(model_name, **self.model_defaults).train()
train_return_nodes, eval_return_nodes = self._get_return_nodes(model)
model = self._create_feature_extractor(
model, train_return_nodes=train_return_nodes, eval_return_nodes=eval_return_nodes
)
model = torch.jit.script(model)
fgn_out = model(self.inp)
out_agg = 0
for node_out in fgn_out.values():
if isinstance(node_out, Sequence):
out_agg += sum(o.float().mean() for o in node_out if o is not None)
elif isinstance(node_out, Mapping):
out_agg += sum(o.float().mean() for o in node_out.values() if o is not None)
else:
# Assume that the only other alternative at this point is a Tensor
out_agg += node_out.float().mean()
out_agg.backward()
def test_train_eval(self):
class TestModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.dropout = torch.nn.Dropout(p=1.0)
def forward(self, x):
x = x.float().mean()
x = self.dropout(x) # dropout
if self.training:
x += 100 # add
else:
x *= 0 # mul
x -= 0 # sub
return x
model = TestModel()
train_return_nodes = ["dropout", "add", "sub"]
eval_return_nodes = ["dropout", "mul", "sub"]
def checks(model, mode):
with torch.no_grad():
out = model(torch.ones(10, 10))
if mode == "train":
# Check that dropout is respected
assert out["dropout"].item() == 0
# Check that control flow dependent on training_mode is respected
assert out["sub"].item() == 100
assert "add" in out
assert "mul" not in out
elif mode == "eval":
# Check that dropout is respected
assert out["dropout"].item() == 1
# Check that control flow dependent on training_mode is respected
assert out["sub"].item() == 0
assert "mul" in out
assert "add" not in out
# Starting from train mode
model.train()
fx_model = self._create_feature_extractor(
model, train_return_nodes=train_return_nodes, eval_return_nodes=eval_return_nodes
)
# Check that the models stay in their original training state
assert model.training
assert fx_model.training
# Check outputs
checks(fx_model, "train")
# Check outputs after switching to eval mode
fx_model.eval()
checks(fx_model, "eval")
# Starting from eval mode
model.eval()
fx_model = self._create_feature_extractor(
model, train_return_nodes=train_return_nodes, eval_return_nodes=eval_return_nodes
)
# Check that the models stay in their original training state
assert not model.training
assert not fx_model.training
# Check outputs
checks(fx_model, "eval")
# Check outputs after switching to train mode
fx_model.train()
checks(fx_model, "train")
def test_leaf_module_and_function(self):
class LeafModule(torch.nn.Module):
def forward(self, x):
# This would raise a TypeError if it were not in a leaf module
int(x.shape[0])
return torch.nn.functional.relu(x + 4)
class TestModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(3, 1, 3)
self.leaf_module = LeafModule()
def forward(self, x):
leaf_function(x.shape[0])
x = self.conv(x)
return self.leaf_module(x)
model = self._create_feature_extractor(
TestModule(),
return_nodes=["leaf_module"],
tracer_kwargs={"leaf_modules": [LeafModule], "autowrap_functions": [leaf_function]},
).train()
# Check that LeafModule is not in the list of nodes
assert "relu" not in [str(n) for n in model.graph.nodes]
assert "leaf_module" in [str(n) for n in model.graph.nodes]
# Check forward
out = model(self.inp)
# And backward
out["leaf_module"].float().mean().backward()
test/test_cpp_models.py deleted 100644 → 0
View file @ e17f5ea2
import torch
import os
import unittest
from torchvision import models, transforms
import sys
from PIL import Image
import torchvision.transforms.functional as F
try:
from torchvision import _C_tests
except ImportError:
_C_tests = None
def process_model(model, tensor, func, name):
model.eval()
traced_script_module = torch.jit.trace(model, tensor)
traced_script_module.save("model.pt")
py_output = model.forward(tensor)
cpp_output = func("model.pt", tensor)
assert torch.allclose(py_output, cpp_output), 'Output mismatch of ' + name + ' models'
def read_image1():
image_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'assets', 'encode_jpeg',
'grace_hopper_517x606.jpg')
image = Image.open(image_path)
image = image.resize((224, 224))
x = F.to_tensor(image)
return x.view(1, 3, 224, 224)
def read_image2():
image_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'assets', 'encode_jpeg',
'grace_hopper_517x606.jpg')
image = Image.open(image_path)
image = image.resize((299, 299))
x = F.to_tensor(image)
x = x.view(1, 3, 299, 299)
return torch.cat([x, x], 0)
@unittest.skipIf(
sys.platform == "darwin" or True,
"C++ models are broken on OS X at the moment, and there's a BC breakage on master; "
"see https://github.com/pytorch/vision/issues/1191")
class Tester(unittest.TestCase):
pretrained = False
image = read_image1()
def test_alexnet(self):
process_model(models.alexnet(self.pretrained), self.image, _C_tests.forward_alexnet, 'Alexnet')
def test_vgg11(self):
process_model(models.vgg11(self.pretrained), self.image, _C_tests.forward_vgg11, 'VGG11')
def test_vgg13(self):
process_model(models.vgg13(self.pretrained), self.image, _C_tests.forward_vgg13, 'VGG13')
def test_vgg16(self):
process_model(models.vgg16(self.pretrained), self.image, _C_tests.forward_vgg16, 'VGG16')
def test_vgg19(self):
process_model(models.vgg19(self.pretrained), self.image, _C_tests.forward_vgg19, 'VGG19')
def test_vgg11_bn(self):
process_model(models.vgg11_bn(self.pretrained), self.image, _C_tests.forward_vgg11bn, 'VGG11BN')
def test_vgg13_bn(self):
process_model(models.vgg13_bn(self.pretrained), self.image, _C_tests.forward_vgg13bn, 'VGG13BN')
def test_vgg16_bn(self):
process_model(models.vgg16_bn(self.pretrained), self.image, _C_tests.forward_vgg16bn, 'VGG16BN')
def test_vgg19_bn(self):
process_model(models.vgg19_bn(self.pretrained), self.image, _C_tests.forward_vgg19bn, 'VGG19BN')
def test_resnet18(self):
process_model(models.resnet18(self.pretrained), self.image, _C_tests.forward_resnet18, 'Resnet18')
def test_resnet34(self):
process_model(models.resnet34(self.pretrained), self.image, _C_tests.forward_resnet34, 'Resnet34')
def test_resnet50(self):
process_model(models.resnet50(self.pretrained), self.image, _C_tests.forward_resnet50, 'Resnet50')
def test_resnet101(self):
process_model(models.resnet101(self.pretrained), self.image, _C_tests.forward_resnet101, 'Resnet101')
def test_resnet152(self):
process_model(models.resnet152(self.pretrained), self.image, _C_tests.forward_resnet152, 'Resnet152')
def test_resnext50_32x4d(self):
process_model(models.resnext50_32x4d(), self.image, _C_tests.forward_resnext50_32x4d, 'ResNext50_32x4d')
def test_resnext101_32x8d(self):
process_model(models.resnext101_32x8d(), self.image, _C_tests.forward_resnext101_32x8d, 'ResNext101_32x8d')
def test_wide_resnet50_2(self):
process_model(models.wide_resnet50_2(), self.image, _C_tests.forward_wide_resnet50_2, 'WideResNet50_2')
def test_wide_resnet101_2(self):
process_model(models.wide_resnet101_2(), self.image, _C_tests.forward_wide_resnet101_2, 'WideResNet101_2')
def test_squeezenet1_0(self):
process_model(models.squeezenet1_0(self.pretrained), self.image,
_C_tests.forward_squeezenet1_0, 'Squeezenet1.0')
def test_squeezenet1_1(self):
process_model(models.squeezenet1_1(self.pretrained), self.image,
_C_tests.forward_squeezenet1_1, 'Squeezenet1.1')
def test_densenet121(self):
process_model(models.densenet121(self.pretrained), self.image, _C_tests.forward_densenet121, 'Densenet121')
def test_densenet169(self):
process_model(models.densenet169(self.pretrained), self.image, _C_tests.forward_densenet169, 'Densenet169')
def test_densenet201(self):
process_model(models.densenet201(self.pretrained), self.image, _C_tests.forward_densenet201, 'Densenet201')
def test_densenet161(self):
process_model(models.densenet161(self.pretrained), self.image, _C_tests.forward_densenet161, 'Densenet161')
def test_mobilenet_v2(self):
process_model(models.mobilenet_v2(self.pretrained), self.image, _C_tests.forward_mobilenetv2, 'MobileNet')
def test_googlenet(self):
process_model(models.googlenet(self.pretrained), self.image, _C_tests.forward_googlenet, 'GoogLeNet')
def test_mnasnet0_5(self):
process_model(models.mnasnet0_5(self.pretrained), self.image, _C_tests.forward_mnasnet0_5, 'MNASNet0_5')
def test_mnasnet0_75(self):
process_model(models.mnasnet0_75(self.pretrained), self.image, _C_tests.forward_mnasnet0_75, 'MNASNet0_75')
def test_mnasnet1_0(self):
process_model(models.mnasnet1_0(self.pretrained), self.image, _C_tests.forward_mnasnet1_0, 'MNASNet1_0')
def test_mnasnet1_3(self):
process_model(models.mnasnet1_3(self.pretrained), self.image, _C_tests.forward_mnasnet1_3, 'MNASNet1_3')
def test_inception_v3(self):
self.image = read_image2()
process_model(models.inception_v3(self.pretrained), self.image, _C_tests.forward_inceptionv3, 'Inceptionv3')
if __name__ == '__main__':
unittest.main()
test/test_datasets.py
View file @ bf491463
import bz2 import bz2
import contextlib import contextlib
import csv
import io import io
import itertools import itertools
import json
import os import os
import pathlib import pathlib
import pickle import pickle
import json
import random import random
import re
import shutil import shutil
import string import string
import unittest import unittest
import xml.etree.ElementTree as ET import xml.etree.ElementTree as ET
import zipfile import zipfile
from typing import Callable, Tuple, Union
import PIL
import datasets_utils import datasets_utils
import numpy as np import numpy as np
import PIL
import pytest
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
from common_utils import combinations_grid
from torchvision import datasets from torchvision import datasets
from torchvision.transforms import v2
class STL10TestCase(datasets_utils.ImageDatasetTestCase): class STL10TestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.STL10 DATASET_CLASS = datasets.STL10
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( ADDITIONAL_CONFIGS = combinations_grid(split=("train", "test", "unlabeled", "train+unlabeled"))
split=("train", "test", "unlabeled", "train+unlabeled"))
@staticmethod @staticmethod
def _make_binary_file(num_elements, root, name): def _make_binary_file(num_elements, root, name):
...@@ -88,20 +93,20 @@ class STL10TestCase(datasets_utils.ImageDatasetTestCase): ...@@ -88,20 +93,20 @@ class STL10TestCase(datasets_utils.ImageDatasetTestCase):
def test_folds(self): def test_folds(self):
for fold in range(10): for fold in range(10):
with self.create_dataset(split="train", folds=fold) as (dataset, _): with self.create_dataset(split="train", folds=fold) as (dataset, _):
self.assertEqual(len(dataset), fold + 1) assert len(dataset) == fold + 1
def test_unlabeled(self): def test_unlabeled(self):
with self.create_dataset(split="unlabeled") as (dataset, _): with self.create_dataset(split="unlabeled") as (dataset, _):
labels = [dataset[idx][1] for idx in range(len(dataset))] labels = [dataset[idx][1] for idx in range(len(dataset))]
self.assertTrue(all(label == -1 for label in labels)) assert all(label == -1 for label in labels)
def test_invalid_folds1(self): def test_invalid_folds1(self):
with self.assertRaises(ValueError): with pytest.raises(ValueError):
with self.create_dataset(folds=10): with self.create_dataset(folds=10):
pass pass
def test_invalid_folds2(self): def test_invalid_folds2(self):
with self.assertRaises(ValueError): with pytest.raises(ValueError):
with self.create_dataset(folds="0"): with self.create_dataset(folds="0"):
pass pass
...@@ -110,9 +115,7 @@ class Caltech101TestCase(datasets_utils.ImageDatasetTestCase): ...@@ -110,9 +115,7 @@ class Caltech101TestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Caltech101 DATASET_CLASS = datasets.Caltech101
FEATURE_TYPES = (PIL.Image.Image, (int, np.ndarray, tuple)) FEATURE_TYPES = (PIL.Image.Image, (int, np.ndarray, tuple))
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( ADDITIONAL_CONFIGS = combinations_grid(target_type=("category", "annotation", ["category", "annotation"]))
target_type=("category", "annotation", ["category", "annotation"])
)
REQUIRED_PACKAGES = ("scipy",) REQUIRED_PACKAGES = ("scipy",)
def inject_fake_data(self, tmpdir, config): def inject_fake_data(self, tmpdir, config):
...@@ -167,23 +170,24 @@ class Caltech101TestCase(datasets_utils.ImageDatasetTestCase): ...@@ -167,23 +170,24 @@ class Caltech101TestCase(datasets_utils.ImageDatasetTestCase):
actual = len(individual_targets) actual = len(individual_targets)
expected = len(combined_targets) expected = len(combined_targets)
self.assertEqual( assert (
actual, actual == expected
expected, ), "The number of the returned combined targets does not match the the number targets if requested "
f"The number of the returned combined targets does not match the the number targets if requested " f"individually: {actual} != {expected}",
f"individually: {actual} != {expected}",
)
for target_type, combined_target, individual_target in zip(target_types, combined_targets, individual_targets): for target_type, combined_target, individual_target in zip(target_types, combined_targets, individual_targets):
with self.subTest(target_type=target_type): with self.subTest(target_type=target_type):
actual = type(combined_target) actual = type(combined_target)
expected = type(individual_target) expected = type(individual_target)
self.assertIs( assert (
actual, actual is expected
expected, ), "Type of the combined target does not match the type of the corresponding individual target: "
f"Type of the combined target does not match the type of the corresponding individual target: " f"{actual} is not {expected}",
f"{actual} is not {expected}",
) def test_transforms_v2_wrapper_spawn(self):
expected_size = (123, 321)
with self.create_dataset(target_type="category", transform=v2.Resize(size=expected_size)) as (dataset, _):
datasets_utils.check_transforms_v2_wrapper_spawn(dataset, expected_size=expected_size)
class Caltech256TestCase(datasets_utils.ImageDatasetTestCase): class Caltech256TestCase(datasets_utils.ImageDatasetTestCase):
...@@ -192,7 +196,7 @@ class Caltech256TestCase(datasets_utils.ImageDatasetTestCase): ...@@ -192,7 +196,7 @@ class Caltech256TestCase(datasets_utils.ImageDatasetTestCase):
def inject_fake_data(self, tmpdir, config): def inject_fake_data(self, tmpdir, config):
tmpdir = pathlib.Path(tmpdir) / "caltech256" / "256_ObjectCategories" tmpdir = pathlib.Path(tmpdir) / "caltech256" / "256_ObjectCategories"
categories = ((1, "ak47"), (127, "laptop-101"), (257, "clutter")) categories = ((1, "ak47"), (2, "american-flag"), (3, "backpack"))
num_images_per_category = 2 num_images_per_category = 2
for idx, category in categories: for idx, category in categories:
...@@ -209,11 +213,11 @@ class Caltech256TestCase(datasets_utils.ImageDatasetTestCase): ...@@ -209,11 +213,11 @@ class Caltech256TestCase(datasets_utils.ImageDatasetTestCase):
class WIDERFaceTestCase(datasets_utils.ImageDatasetTestCase): class WIDERFaceTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.WIDERFace DATASET_CLASS = datasets.WIDERFace
FEATURE_TYPES = (PIL.Image.Image, (dict, type(None))) # test split returns None as target FEATURE_TYPES = (PIL.Image.Image, (dict, type(None))) # test split returns None as target
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=('train', 'val', 'test')) ADDITIONAL_CONFIGS = combinations_grid(split=("train", "val", "test"))
def inject_fake_data(self, tmpdir, config): def inject_fake_data(self, tmpdir, config):
widerface_dir = pathlib.Path(tmpdir) / 'widerface' widerface_dir = pathlib.Path(tmpdir) / "widerface"
annotations_dir = widerface_dir / 'wider_face_split' annotations_dir = widerface_dir / "wider_face_split"
os.makedirs(annotations_dir) os.makedirs(annotations_dir)
split_to_idx = split_to_num_examples = { split_to_idx = split_to_num_examples = {
...@@ -223,21 +227,21 @@ class WIDERFaceTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -223,21 +227,21 @@ class WIDERFaceTestCase(datasets_utils.ImageDatasetTestCase):
} }
# We need to create all folders regardless of the split in config # We need to create all folders regardless of the split in config
for split in ('train', 'val', 'test'): for split in ("train", "val", "test"):
split_idx = split_to_idx[split] split_idx = split_to_idx[split]
num_examples = split_to_num_examples[split] num_examples = split_to_num_examples[split]
datasets_utils.create_image_folder( datasets_utils.create_image_folder(
root=tmpdir, root=tmpdir,
name=widerface_dir / f'WIDER_{split}' / 'images' / '0--Parade', name=widerface_dir / f"WIDER_{split}" / "images" / "0--Parade",
file_name_fn=lambda image_idx: f"0_Parade_marchingband_1_{split_idx + image_idx}.jpg", file_name_fn=lambda image_idx: f"0_Parade_marchingband_1_{split_idx + image_idx}.jpg",
num_examples=num_examples, num_examples=num_examples,
) )
annotation_file_name = { annotation_file_name = {
'train': annotations_dir / 'wider_face_train_bbx_gt.txt', "train": annotations_dir / "wider_face_train_bbx_gt.txt",
'val': annotations_dir / 'wider_face_val_bbx_gt.txt', "val": annotations_dir / "wider_face_val_bbx_gt.txt",
'test': annotations_dir / 'wider_face_test_filelist.txt', "test": annotations_dir / "wider_face_test_filelist.txt",
}[split] }[split]
annotation_content = { annotation_content = {
...@@ -260,6 +264,11 @@ class WIDERFaceTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -260,6 +264,11 @@ class WIDERFaceTestCase(datasets_utils.ImageDatasetTestCase):
return split_to_num_examples[config["split"]] return split_to_num_examples[config["split"]]
def test_transforms_v2_wrapper_spawn(self):
expected_size = (123, 321)
with self.create_dataset(transform=v2.Resize(size=expected_size)) as (dataset, _):
datasets_utils.check_transforms_v2_wrapper_spawn(dataset, expected_size=expected_size)
class CityScapesTestCase(datasets_utils.ImageDatasetTestCase): class CityScapesTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Cityscapes DATASET_CLASS = datasets.Cityscapes
...@@ -270,10 +279,8 @@ class CityScapesTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -270,10 +279,8 @@ class CityScapesTestCase(datasets_utils.ImageDatasetTestCase):
"color", "color",
) )
ADDITIONAL_CONFIGS = ( ADDITIONAL_CONFIGS = (
*datasets_utils.combinations_grid( *combinations_grid(mode=("fine",), split=("train", "test", "val"), target_type=TARGET_TYPES),
mode=("fine",), split=("train", "test", "val"), target_type=TARGET_TYPES *combinations_grid(
),
*datasets_utils.combinations_grid(
mode=("coarse",), mode=("coarse",),
split=("train", "train_extra", "val"), split=("train", "train_extra", "val"),
target_type=TARGET_TYPES, target_type=TARGET_TYPES,
...@@ -327,6 +334,7 @@ class CityScapesTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -327,6 +334,7 @@ class CityScapesTestCase(datasets_utils.ImageDatasetTestCase):
gt_dir = tmpdir / f"gt{mode}" gt_dir = tmpdir / f"gt{mode}"
for split in mode_to_splits[mode]: for split in mode_to_splits[mode]:
for city in cities: for city in cities:
def make_image(name, size=10): def make_image(name, size=10):
datasets_utils.create_image_folder( datasets_utils.create_image_folder(
root=gt_dir / split, root=gt_dir / split,
...@@ -335,6 +343,7 @@ class CityScapesTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -335,6 +343,7 @@ class CityScapesTestCase(datasets_utils.ImageDatasetTestCase):
size=size, size=size,
num_examples=1, num_examples=1,
) )
make_image(f"{city}_000000_000000_gt{mode}_instanceIds.png") make_image(f"{city}_000000_000000_gt{mode}_instanceIds.png")
make_image(f"{city}_000000_000000_gt{mode}_labelIds.png") make_image(f"{city}_000000_000000_gt{mode}_labelIds.png")
make_image(f"{city}_000000_000000_gt{mode}_color.png", size=(4, 10, 10)) make_image(f"{city}_000000_000000_gt{mode}_color.png", size=(4, 10, 10))
...@@ -344,7 +353,7 @@ class CityScapesTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -344,7 +353,7 @@ class CityScapesTestCase(datasets_utils.ImageDatasetTestCase):
json.dump(polygon_target, outfile) json.dump(polygon_target, outfile)
# Create leftImg8bit folder # Create leftImg8bit folder
for split in ['test', 'train_extra', 'train', 'val']: for split in ["test", "train_extra", "train", "val"]:
for city in cities: for city in cities:
datasets_utils.create_image_folder( datasets_utils.create_image_folder(
root=tmpdir / "leftImg8bit" / split, root=tmpdir / "leftImg8bit" / split,
...@@ -353,52 +362,58 @@ class CityScapesTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -353,52 +362,58 @@ class CityScapesTestCase(datasets_utils.ImageDatasetTestCase):
num_examples=1, num_examples=1,
) )
info = {'num_examples': len(cities)} info = {"num_examples": len(cities)}
if config['target_type'] == 'polygon': if config["target_type"] == "polygon":
info['expected_polygon_target'] = polygon_target info["expected_polygon_target"] = polygon_target
return info return info
def test_combined_targets(self): def test_combined_targets(self):
target_types = ['semantic', 'polygon', 'color'] target_types = ["semantic", "polygon", "color"]
with self.create_dataset(target_type=target_types) as (dataset, _): with self.create_dataset(target_type=target_types) as (dataset, _):
output = dataset[0] output = dataset[0]
self.assertTrue(isinstance(output, tuple)) assert isinstance(output, tuple)
self.assertTrue(len(output) == 2) assert len(output) == 2
self.assertTrue(isinstance(output[0], PIL.Image.Image)) assert isinstance(output[0], PIL.Image.Image)
self.assertTrue(isinstance(output[1], tuple)) assert isinstance(output[1], tuple)
self.assertTrue(len(output[1]) == 3) assert len(output[1]) == 3
self.assertTrue(isinstance(output[1][0], PIL.Image.Image)) # semantic assert isinstance(output[1][0], PIL.Image.Image) # semantic
self.assertTrue(isinstance(output[1][1], dict)) # polygon assert isinstance(output[1][1], dict) # polygon
self.assertTrue(isinstance(output[1][2], PIL.Image.Image)) # color assert isinstance(output[1][2], PIL.Image.Image) # color
def test_feature_types_target_color(self): def test_feature_types_target_color(self):
with self.create_dataset(target_type='color') as (dataset, _): with self.create_dataset(target_type="color") as (dataset, _):
color_img, color_target = dataset[0] color_img, color_target = dataset[0]
self.assertTrue(isinstance(color_img, PIL.Image.Image)) assert isinstance(color_img, PIL.Image.Image)
self.assertTrue(np.array(color_target).shape[2] == 4) assert np.array(color_target).shape[2] == 4
def test_feature_types_target_polygon(self): def test_feature_types_target_polygon(self):
with self.create_dataset(target_type='polygon') as (dataset, info): with self.create_dataset(target_type="polygon") as (dataset, info):
polygon_img, polygon_target = dataset[0] polygon_img, polygon_target = dataset[0]
self.assertTrue(isinstance(polygon_img, PIL.Image.Image)) assert isinstance(polygon_img, PIL.Image.Image)
self.assertEqual(polygon_target, info['expected_polygon_target']) (polygon_target, info["expected_polygon_target"])
def test_transforms_v2_wrapper_spawn(self):
expected_size = (123, 321)
for target_type in ["instance", "semantic", ["instance", "semantic"]]:
with self.create_dataset(target_type=target_type, transform=v2.Resize(size=expected_size)) as (dataset, _):
datasets_utils.check_transforms_v2_wrapper_spawn(dataset, expected_size=expected_size)
class ImageNetTestCase(datasets_utils.ImageDatasetTestCase): class ImageNetTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.ImageNet DATASET_CLASS = datasets.ImageNet
REQUIRED_PACKAGES = ('scipy',) REQUIRED_PACKAGES = ("scipy",)
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=('train', 'val')) ADDITIONAL_CONFIGS = combinations_grid(split=("train", "val"))
def inject_fake_data(self, tmpdir, config): def inject_fake_data(self, tmpdir, config):
tmpdir = pathlib.Path(tmpdir) tmpdir = pathlib.Path(tmpdir)
wnid = 'n01234567' wnid = "n01234567"
if config['split'] == 'train': if config["split"] == "train":
num_examples = 3 num_examples = 3
datasets_utils.create_image_folder( datasets_utils.create_image_folder(
root=tmpdir, root=tmpdir,
name=tmpdir / 'train' / wnid / wnid, name=tmpdir / "train" / wnid / wnid,
file_name_fn=lambda image_idx: f"{wnid}_{image_idx}.JPEG", file_name_fn=lambda image_idx: f"{wnid}_{image_idx}.JPEG",
num_examples=num_examples, num_examples=num_examples,
) )
...@@ -406,19 +421,24 @@ class ImageNetTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -406,19 +421,24 @@ class ImageNetTestCase(datasets_utils.ImageDatasetTestCase):
num_examples = 1 num_examples = 1
datasets_utils.create_image_folder( datasets_utils.create_image_folder(
root=tmpdir, root=tmpdir,
name=tmpdir / 'val' / wnid, name=tmpdir / "val" / wnid,
file_name_fn=lambda image_ifx: "ILSVRC2012_val_0000000{image_idx}.JPEG", file_name_fn=lambda image_ifx: "ILSVRC2012_val_0000000{image_idx}.JPEG",
num_examples=num_examples, num_examples=num_examples,
) )
wnid_to_classes = {wnid: [1]} wnid_to_classes = {wnid: [1]}
torch.save((wnid_to_classes, None), tmpdir / 'meta.bin') torch.save((wnid_to_classes, None), tmpdir / "meta.bin")
return num_examples return num_examples
def test_transforms_v2_wrapper_spawn(self):
expected_size = (123, 321)
with self.create_dataset(transform=v2.Resize(size=expected_size)) as (dataset, _):
datasets_utils.check_transforms_v2_wrapper_spawn(dataset, expected_size=expected_size)
class CIFAR10TestCase(datasets_utils.ImageDatasetTestCase): class CIFAR10TestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.CIFAR10 DATASET_CLASS = datasets.CIFAR10
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(train=(True, False)) ADDITIONAL_CONFIGS = combinations_grid(train=(True, False))
_VERSION_CONFIG = dict( _VERSION_CONFIG = dict(
base_folder="cifar-10-batches-py", base_folder="cifar-10-batches-py",
...@@ -447,8 +467,9 @@ class CIFAR10TestCase(datasets_utils.ImageDatasetTestCase): ...@@ -447,8 +467,9 @@ class CIFAR10TestCase(datasets_utils.ImageDatasetTestCase):
) )
def _create_batch_file(self, root, name, num_images): def _create_batch_file(self, root, name, num_images):
np_rng = np.random.RandomState(0)
data = datasets_utils.create_image_or_video_tensor((num_images, 32 * 32 * 3)) 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() labels = np_rng.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}) self._create_binary_file(root, name, {"data": data, self._VERSION_CONFIG["labels_key"]: labels})
def _create_meta_file(self, root): def _create_meta_file(self, root):
...@@ -469,7 +490,7 @@ class CIFAR10TestCase(datasets_utils.ImageDatasetTestCase): ...@@ -469,7 +490,7 @@ class CIFAR10TestCase(datasets_utils.ImageDatasetTestCase):
with self.create_dataset() as (dataset, info): with self.create_dataset() as (dataset, info):
expected = {category: label for label, category in enumerate(info["categories"])} expected = {category: label for label, category in enumerate(info["categories"])}
actual = dataset.class_to_idx actual = dataset.class_to_idx
self.assertEqual(actual, expected) assert actual == expected
class CIFAR100(CIFAR10TestCase): class CIFAR100(CIFAR10TestCase):
...@@ -490,7 +511,7 @@ class CelebATestCase(datasets_utils.ImageDatasetTestCase): ...@@ -490,7 +511,7 @@ class CelebATestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.CelebA DATASET_CLASS = datasets.CelebA
FEATURE_TYPES = (PIL.Image.Image, (torch.Tensor, int, tuple, type(None))) FEATURE_TYPES = (PIL.Image.Image, (torch.Tensor, int, tuple, type(None)))
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( ADDITIONAL_CONFIGS = combinations_grid(
split=("train", "valid", "test", "all"), split=("train", "valid", "test", "all"),
target_type=("attr", "identity", "bbox", "landmarks", ["attr", "identity"]), target_type=("attr", "identity", "bbox", "landmarks", ["attr", "identity"]),
) )
...@@ -514,7 +535,7 @@ class CelebATestCase(datasets_utils.ImageDatasetTestCase): ...@@ -514,7 +535,7 @@ class CelebATestCase(datasets_utils.ImageDatasetTestCase):
return dict(num_examples=num_images_per_split[config["split"]], attr_names=attr_names) return dict(num_examples=num_images_per_split[config["split"]], attr_names=attr_names)
def _create_split_txt(self, root): def _create_split_txt(self, root):
num_images_per_split = dict(train=3, valid=2, test=1) num_images_per_split = dict(train=4, valid=3, test=2)
data = [ data = [
[self._SPLIT_TO_IDX[split]] for split, num_images in num_images_per_split.items() for _ in range(num_images) [self._SPLIT_TO_IDX[split]] for split, num_images in num_images_per_split.items() for _ in range(num_images)
...@@ -573,33 +594,46 @@ class CelebATestCase(datasets_utils.ImageDatasetTestCase): ...@@ -573,33 +594,46 @@ class CelebATestCase(datasets_utils.ImageDatasetTestCase):
actual = len(individual_targets) actual = len(individual_targets)
expected = len(combined_targets) expected = len(combined_targets)
self.assertEqual( assert (
actual, actual == expected
expected, ), "The number of the returned combined targets does not match the the number targets if requested "
f"The number of the returned combined targets does not match the the number targets if requested " f"individually: {actual} != {expected}",
f"individually: {actual} != {expected}",
)
for target_type, combined_target, individual_target in zip(target_types, combined_targets, individual_targets): for target_type, combined_target, individual_target in zip(target_types, combined_targets, individual_targets):
with self.subTest(target_type=target_type): with self.subTest(target_type=target_type):
actual = type(combined_target) actual = type(combined_target)
expected = type(individual_target) expected = type(individual_target)
self.assertIs( assert (
actual, actual is expected
expected, ), "Type of the combined target does not match the type of the corresponding individual target: "
f"Type of the combined target does not match the type of the corresponding individual target: " f"{actual} is not {expected}",
f"{actual} is not {expected}",
)
def test_no_target(self): def test_no_target(self):
with self.create_dataset(target_type=[]) as (dataset, _): with self.create_dataset(target_type=[]) as (dataset, _):
_, target = dataset[0] _, target = dataset[0]
self.assertIsNone(target) assert target is None
def test_attr_names(self): def test_attr_names(self):
with self.create_dataset() as (dataset, info): with self.create_dataset() as (dataset, info):
self.assertEqual(tuple(dataset.attr_names), info["attr_names"]) assert tuple(dataset.attr_names) == info["attr_names"]
def test_images_names_split(self):
with self.create_dataset(split="all") as (dataset, _):
all_imgs_names = set(dataset.filename)
merged_imgs_names = set()
for split in ["train", "valid", "test"]:
with self.create_dataset(split=split) as (dataset, _):
merged_imgs_names.update(dataset.filename)
assert merged_imgs_names == all_imgs_names
def test_transforms_v2_wrapper_spawn(self):
expected_size = (123, 321)
for target_type in ["identity", "bbox", ["identity", "bbox"]]:
with self.create_dataset(target_type=target_type, transform=v2.Resize(size=expected_size)) as (dataset, _):
datasets_utils.check_transforms_v2_wrapper_spawn(dataset, expected_size=expected_size)
class VOCSegmentationTestCase(datasets_utils.ImageDatasetTestCase): class VOCSegmentationTestCase(datasets_utils.ImageDatasetTestCase):
...@@ -607,19 +641,12 @@ class VOCSegmentationTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -607,19 +641,12 @@ class VOCSegmentationTestCase(datasets_utils.ImageDatasetTestCase):
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image) FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image)
ADDITIONAL_CONFIGS = ( ADDITIONAL_CONFIGS = (
*datasets_utils.combinations_grid( *combinations_grid(year=[f"20{year:02d}" for year in range(7, 13)], image_set=("train", "val", "trainval")),
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", image_set="test"),
dict(year="2007-test", image_set="test"),
) )
def inject_fake_data(self, tmpdir, config): def inject_fake_data(self, tmpdir, config):
year, is_test_set = ( year, is_test_set = config["year"], config["image_set"] == "test"
("2007", True)
if config["year"] == "2007-test" or config["image_set"] == "test"
else (config["year"], False)
)
image_set = config["image_set"] image_set = config["image_set"]
base_dir = pathlib.Path(tmpdir) base_dir = pathlib.Path(tmpdir)
...@@ -650,7 +677,7 @@ class VOCSegmentationTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -650,7 +677,7 @@ class VOCSegmentationTestCase(datasets_utils.ImageDatasetTestCase):
shutil.copytree(src, root / "Segmentation") shutil.copytree(src, root / "Segmentation")
num_images = max(itertools.chain(*idcs.values())) + 1 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()]) num_images_per_image_set = {image_set: len(idcs_) for image_set, idcs_ in idcs.items()}
return num_images, num_images_per_image_set return num_images, num_images_per_image_set
def _create_image_set_file(self, root, image_set, idcs): def _create_image_set_file(self, root, image_set, idcs):
...@@ -695,6 +722,11 @@ class VOCSegmentationTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -695,6 +722,11 @@ class VOCSegmentationTestCase(datasets_utils.ImageDatasetTestCase):
return data return data
def test_transforms_v2_wrapper_spawn(self):
expected_size = (123, 321)
with self.create_dataset(transform=v2.Resize(size=expected_size)) as (dataset, _):
datasets_utils.check_transforms_v2_wrapper_spawn(dataset, expected_size=expected_size)
class VOCDetectionTestCase(VOCSegmentationTestCase): class VOCDetectionTestCase(VOCSegmentationTestCase):
DATASET_CLASS = datasets.VOCDetection DATASET_CLASS = datasets.VOCDetection
...@@ -704,16 +736,21 @@ class VOCDetectionTestCase(VOCSegmentationTestCase): ...@@ -704,16 +736,21 @@ class VOCDetectionTestCase(VOCSegmentationTestCase):
with self.create_dataset() as (dataset, info): with self.create_dataset() as (dataset, info):
_, target = dataset[0] _, target = dataset[0]
self.assertIn("annotation", target) assert "annotation" in target
annotation = target["annotation"] annotation = target["annotation"]
self.assertIn("object", annotation) assert "object" in annotation
objects = annotation["object"] objects = annotation["object"]
self.assertEqual(len(objects), 1) assert len(objects) == 1
object = objects[0] object = objects[0]
self.assertEqual(object, info["annotation"]) assert object == info["annotation"]
def test_transforms_v2_wrapper_spawn(self):
expected_size = (123, 321)
with self.create_dataset(transform=v2.Resize(size=expected_size)) as (dataset, _):
datasets_utils.check_transforms_v2_wrapper_spawn(dataset, expected_size=expected_size)
class CocoDetectionTestCase(datasets_utils.ImageDatasetTestCase): class CocoDetectionTestCase(datasets_utils.ImageDatasetTestCase):
...@@ -745,28 +782,52 @@ class CocoDetectionTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -745,28 +782,52 @@ class CocoDetectionTestCase(datasets_utils.ImageDatasetTestCase):
annotation_folder = tmpdir / self._ANNOTATIONS_FOLDER annotation_folder = tmpdir / self._ANNOTATIONS_FOLDER
os.makedirs(annotation_folder) os.makedirs(annotation_folder)
segmentation_kind = config.pop("segmentation_kind", "list")
info = self._create_annotation_file( info = self._create_annotation_file(
annotation_folder, self._ANNOTATIONS_FILE, file_names, num_annotations_per_image annotation_folder,
self._ANNOTATIONS_FILE,
file_names,
num_annotations_per_image,
segmentation_kind=segmentation_kind,
) )
info["num_examples"] = num_images info["num_examples"] = num_images
return info return info
def _create_annotation_file(self, root, name, file_names, num_annotations_per_image): def _create_annotation_file(self, root, name, file_names, num_annotations_per_image, segmentation_kind="list"):
image_ids = [int(file_name.stem) for file_name in file_names] image_ids = [int(file_name.stem) for file_name in file_names]
images = [dict(file_name=str(file_name), id=id) for file_name, id in zip(file_names, image_ids)] images = [dict(file_name=str(file_name), id=id) for file_name, id in zip(file_names, image_ids)]
annotations, info = self._create_annotations(image_ids, num_annotations_per_image) annotations, info = self._create_annotations(image_ids, num_annotations_per_image, segmentation_kind)
self._create_json(root, name, dict(images=images, annotations=annotations)) self._create_json(root, name, dict(images=images, annotations=annotations))
return info return info
def _create_annotations(self, image_ids, num_annotations_per_image): def _create_annotations(self, image_ids, num_annotations_per_image, segmentation_kind="list"):
annotations = datasets_utils.combinations_grid( annotations = []
image_id=image_ids, bbox=([1.0, 2.0, 3.0, 4.0],) * num_annotations_per_image annotion_id = 0
)
for id, annotation in enumerate(annotations): for image_id in itertools.islice(itertools.cycle(image_ids), len(image_ids) * num_annotations_per_image):
annotation["id"] = id segmentation = {
"list": [torch.rand(8).tolist()],
"rle": {"size": [10, 10], "counts": [1]},
"rle_encoded": {"size": [2400, 2400], "counts": "PQRQ2[1\\Y2f0gNVNRhMg2"},
"bad": 123,
}[segmentation_kind]
annotations.append(
dict(
image_id=image_id,
id=annotion_id,
bbox=torch.rand(4).tolist(),
segmentation=segmentation,
category_id=int(torch.randint(91, ())),
area=float(torch.rand(1)),
iscrowd=int(torch.randint(2, size=(1,))),
)
)
annotion_id += 1
return annotations, dict() return annotations, dict()
def _create_json(self, root, name, content): def _create_json(self, root, name, content):
...@@ -775,13 +836,39 @@ class CocoDetectionTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -775,13 +836,39 @@ class CocoDetectionTestCase(datasets_utils.ImageDatasetTestCase):
json.dump(content, fh) json.dump(content, fh)
return file return file
def test_transforms_v2_wrapper_spawn(self):
expected_size = (123, 321)
with self.create_dataset(transform=v2.Resize(size=expected_size)) as (dataset, _):
datasets_utils.check_transforms_v2_wrapper_spawn(dataset, expected_size=expected_size)
def test_slice_error(self):
with self.create_dataset() as (dataset, _):
with pytest.raises(ValueError, match="Index must be of type integer"):
dataset[:2]
def test_segmentation_kind(self):
if isinstance(self, CocoCaptionsTestCase):
return
for segmentation_kind in ("list", "rle", "rle_encoded"):
config = {"segmentation_kind": segmentation_kind}
with self.create_dataset(config) as (dataset, _):
dataset = datasets.wrap_dataset_for_transforms_v2(dataset, target_keys="all")
list(dataset)
config = {"segmentation_kind": "bad"}
with self.create_dataset(config) as (dataset, _):
dataset = datasets.wrap_dataset_for_transforms_v2(dataset, target_keys="all")
with pytest.raises(ValueError, match="COCO segmentation expected to be a dict or a list"):
list(dataset)
class CocoCaptionsTestCase(CocoDetectionTestCase): class CocoCaptionsTestCase(CocoDetectionTestCase):
DATASET_CLASS = datasets.CocoCaptions DATASET_CLASS = datasets.CocoCaptions
def _create_annotations(self, image_ids, num_annotations_per_image): def _create_annotations(self, image_ids, num_annotations_per_image, segmentation_kind="list"):
captions = [str(idx) for idx in range(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) annotations = combinations_grid(image_id=image_ids, caption=captions)
for id, annotation in enumerate(annotations): for id, annotation in enumerate(annotations):
annotation["id"] = id annotation["id"] = id
return annotations, dict(captions=captions) return annotations, dict(captions=captions)
...@@ -789,13 +876,18 @@ class CocoCaptionsTestCase(CocoDetectionTestCase): ...@@ -789,13 +876,18 @@ class CocoCaptionsTestCase(CocoDetectionTestCase):
def test_captions(self): def test_captions(self):
with self.create_dataset() as (dataset, info): with self.create_dataset() as (dataset, info):
_, captions = dataset[0] _, captions = dataset[0]
self.assertEqual(tuple(captions), tuple(info["captions"])) assert tuple(captions) == tuple(info["captions"])
def test_transforms_v2_wrapper_spawn(self):
# We need to define this method, because otherwise the test from the super class will
# be run
pytest.skip("CocoCaptions is currently not supported by the v2 wrapper.")
class UCF101TestCase(datasets_utils.VideoDatasetTestCase): class UCF101TestCase(datasets_utils.VideoDatasetTestCase):
DATASET_CLASS = datasets.UCF101 DATASET_CLASS = datasets.UCF101
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(fold=(1, 2, 3), train=(True, False)) ADDITIONAL_CONFIGS = combinations_grid(fold=(1, 2, 3), train=(True, False))
_VIDEO_FOLDER = "videos" _VIDEO_FOLDER = "videos"
_ANNOTATIONS_FOLDER = "annotations" _ANNOTATIONS_FOLDER = "annotations"
...@@ -849,16 +941,14 @@ class UCF101TestCase(datasets_utils.VideoDatasetTestCase): ...@@ -849,16 +941,14 @@ class UCF101TestCase(datasets_utils.VideoDatasetTestCase):
def _create_annotation_file(self, root, name, video_files): def _create_annotation_file(self, root, name, video_files):
with open(pathlib.Path(root) / name, "w") as fh: with open(pathlib.Path(root) / name, "w") as fh:
fh.writelines(f"{file}\n" for file in sorted(video_files)) fh.writelines(f"{str(file).replace(os.sep, '/')}\n" for file in sorted(video_files))
class LSUNTestCase(datasets_utils.ImageDatasetTestCase): class LSUNTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.LSUN DATASET_CLASS = datasets.LSUN
REQUIRED_PACKAGES = ("lmdb",) REQUIRED_PACKAGES = ("lmdb",)
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( ADDITIONAL_CONFIGS = combinations_grid(classes=("train", "test", "val", ["bedroom_train", "church_outdoor_train"]))
classes=("train", "test", "val", ["bedroom_train", "church_outdoor_train"])
)
_CATEGORIES = ( _CATEGORIES = (
"bedroom", "bedroom",
...@@ -883,10 +973,7 @@ class LSUNTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -883,10 +973,7 @@ class LSUNTestCase(datasets_utils.ImageDatasetTestCase):
return num_images return num_images
@contextlib.contextmanager @contextlib.contextmanager
def create_dataset( def create_dataset(self, *args, **kwargs):
self,
*args, **kwargs
):
with super().create_dataset(*args, **kwargs) as output: with super().create_dataset(*args, **kwargs) as output:
yield output yield output
# Currently datasets.LSUN caches the keys in the current directory rather than in the root directory. Thus, # Currently datasets.LSUN caches the keys in the current directory rather than in the root directory. Thus,
...@@ -940,33 +1027,39 @@ class LSUNTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -940,33 +1027,39 @@ class LSUNTestCase(datasets_utils.ImageDatasetTestCase):
def test_not_found_or_corrupted(self): def test_not_found_or_corrupted(self):
# LSUN does not raise built-in exception, but a custom one. It is expressive enough to not 'cast' it to # LSUN does not raise built-in exception, but a custom one. It is expressive enough to not 'cast' it to
# RuntimeError or FileNotFoundError that are normally checked by this test. # RuntimeError or FileNotFoundError that are normally checked by this test.
with self.assertRaises(datasets_utils.lazy_importer.lmdb.Error): with pytest.raises(datasets_utils.lazy_importer.lmdb.Error):
super().test_not_found_or_corrupted() super().test_not_found_or_corrupted()
class Kinetics400TestCase(datasets_utils.VideoDatasetTestCase): class KineticsTestCase(datasets_utils.VideoDatasetTestCase):
DATASET_CLASS = datasets.Kinetics400 DATASET_CLASS = datasets.Kinetics
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "val"), num_classes=("400", "600", "700"))
def inject_fake_data(self, tmpdir, config): def inject_fake_data(self, tmpdir, config):
classes = ("Abseiling", "Zumba") classes = ("Abseiling", "Zumba")
num_videos_per_class = 2 num_videos_per_class = 2
tmpdir = pathlib.Path(tmpdir) / config["split"]
digits = string.ascii_letters + string.digits + "-_" digits = string.ascii_letters + string.digits + "-_"
for cls in classes: for cls in classes:
datasets_utils.create_video_folder( datasets_utils.create_video_folder(
tmpdir, tmpdir,
cls, cls,
lambda _: f"{datasets_utils.create_random_string(11, digits)}.avi", lambda _: f"{datasets_utils.create_random_string(11, digits)}.mp4",
num_videos_per_class, num_videos_per_class,
) )
return num_videos_per_class * len(classes) return num_videos_per_class * len(classes)
@pytest.mark.xfail(reason="FIXME")
def test_transforms_v2_wrapper_spawn(self):
expected_size = (123, 321)
with self.create_dataset(output_format="TCHW", transform=v2.Resize(size=expected_size)) as (dataset, _):
datasets_utils.check_transforms_v2_wrapper_spawn(dataset, expected_size=expected_size)
class HMDB51TestCase(datasets_utils.VideoDatasetTestCase): class HMDB51TestCase(datasets_utils.VideoDatasetTestCase):
DATASET_CLASS = datasets.HMDB51 DATASET_CLASS = datasets.HMDB51
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(fold=(1, 2, 3), train=(True, False)) ADDITIONAL_CONFIGS = combinations_grid(fold=(1, 2, 3), train=(True, False))
_VIDEO_FOLDER = "videos" _VIDEO_FOLDER = "videos"
_SPLITS_FOLDER = "splits" _SPLITS_FOLDER = "splits"
...@@ -1026,7 +1119,7 @@ class HMDB51TestCase(datasets_utils.VideoDatasetTestCase): ...@@ -1026,7 +1119,7 @@ class HMDB51TestCase(datasets_utils.VideoDatasetTestCase):
class OmniglotTestCase(datasets_utils.ImageDatasetTestCase): class OmniglotTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Omniglot DATASET_CLASS = datasets.Omniglot
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(background=(True, False)) ADDITIONAL_CONFIGS = combinations_grid(background=(True, False))
def inject_fake_data(self, tmpdir, config): def inject_fake_data(self, tmpdir, config):
target_folder = ( target_folder = (
...@@ -1106,7 +1199,7 @@ class SEMEIONTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -1106,7 +1199,7 @@ class SEMEIONTestCase(datasets_utils.ImageDatasetTestCase):
class USPSTestCase(datasets_utils.ImageDatasetTestCase): class USPSTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.USPS DATASET_CLASS = datasets.USPS
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(train=(True, False)) ADDITIONAL_CONFIGS = combinations_grid(train=(True, False))
def inject_fake_data(self, tmpdir, config): def inject_fake_data(self, tmpdir, config):
num_images = 2 if config["train"] else 1 num_images = 2 if config["train"] else 1
...@@ -1128,7 +1221,7 @@ class SBDatasetTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -1128,7 +1221,7 @@ class SBDatasetTestCase(datasets_utils.ImageDatasetTestCase):
REQUIRED_PACKAGES = ("scipy.io", "scipy.sparse") REQUIRED_PACKAGES = ("scipy.io", "scipy.sparse")
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( ADDITIONAL_CONFIGS = combinations_grid(
image_set=("train", "val", "train_noval"), mode=("boundaries", "segmentation") image_set=("train", "val", "train_noval"), mode=("boundaries", "segmentation")
) )
...@@ -1154,7 +1247,7 @@ class SBDatasetTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -1154,7 +1247,7 @@ class SBDatasetTestCase(datasets_utils.ImageDatasetTestCase):
self._create_split_file(root, split, idcs) self._create_split_file(root, split, idcs)
num_images = max(itertools.chain(*splits.values())) + 1 num_images = max(itertools.chain(*splits.values())) + 1
num_images_per_split = dict([(split, len(idcs)) for split, idcs in splits.items()]) num_images_per_split = {split: len(idcs) for split, idcs in splits.items()}
return num_images, num_images_per_split return num_images, num_images_per_split
def _create_split_file(self, root, name, idcs): def _create_split_file(self, root, name, idcs):
...@@ -1189,6 +1282,11 @@ class SBDatasetTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -1189,6 +1282,11 @@ class SBDatasetTestCase(datasets_utils.ImageDatasetTestCase):
def _file_stem(self, idx): def _file_stem(self, idx):
return f"2008_{idx:06d}" return f"2008_{idx:06d}"
def test_transforms_v2_wrapper_spawn(self):
expected_size = (123, 321)
with self.create_dataset(mode="segmentation", transforms=v2.Resize(size=expected_size)) as (dataset, _):
datasets_utils.check_transforms_v2_wrapper_spawn(dataset, expected_size=expected_size)
class FakeDataTestCase(datasets_utils.ImageDatasetTestCase): class FakeDataTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.FakeData DATASET_CLASS = datasets.FakeData
...@@ -1214,7 +1312,7 @@ class PhotoTourTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -1214,7 +1312,7 @@ class PhotoTourTestCase(datasets_utils.ImageDatasetTestCase):
_TRAIN_FEATURE_TYPES = (torch.Tensor,) _TRAIN_FEATURE_TYPES = (torch.Tensor,)
_TEST_FEATURE_TYPES = (torch.Tensor, torch.Tensor, torch.Tensor) _TEST_FEATURE_TYPES = (torch.Tensor, torch.Tensor, torch.Tensor)
datasets_utils.combinations_grid(train=(True, False)) combinations_grid(train=(True, False))
_NAME = "liberty" _NAME = "liberty"
...@@ -1348,7 +1446,8 @@ class Flickr8kTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -1348,7 +1446,8 @@ class Flickr8kTestCase(datasets_utils.ImageDatasetTestCase):
def test_captions(self): def test_captions(self):
with self.create_dataset() as (dataset, info): with self.create_dataset() as (dataset, info):
_, captions = dataset[0] _, captions = dataset[0]
self.assertSequenceEqual(captions, info["captions"]) assert len(captions) == len(info["captions"])
assert all([a == b for a, b in zip(captions, info["captions"])])
class Flickr30kTestCase(Flickr8kTestCase): class Flickr30kTestCase(Flickr8kTestCase):
...@@ -1372,7 +1471,7 @@ class Flickr30kTestCase(Flickr8kTestCase): ...@@ -1372,7 +1471,7 @@ class Flickr30kTestCase(Flickr8kTestCase):
class MNISTTestCase(datasets_utils.ImageDatasetTestCase): class MNISTTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.MNIST DATASET_CLASS = datasets.MNIST
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(train=(True, False)) ADDITIONAL_CONFIGS = combinations_grid(train=(True, False))
_MAGIC_DTYPES = { _MAGIC_DTYPES = {
torch.uint8: 8, torch.uint8: 8,
...@@ -1442,7 +1541,7 @@ class EMNISTTestCase(MNISTTestCase): ...@@ -1442,7 +1541,7 @@ class EMNISTTestCase(MNISTTestCase):
DATASET_CLASS = datasets.EMNIST DATASET_CLASS = datasets.EMNIST
DEFAULT_CONFIG = dict(split="byclass") DEFAULT_CONFIG = dict(split="byclass")
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( ADDITIONAL_CONFIGS = combinations_grid(
split=("byclass", "bymerge", "balanced", "letters", "digits", "mnist"), train=(True, False) split=("byclass", "bymerge", "balanced", "letters", "digits", "mnist"), train=(True, False)
) )
...@@ -1453,7 +1552,7 @@ class EMNISTTestCase(MNISTTestCase): ...@@ -1453,7 +1552,7 @@ class EMNISTTestCase(MNISTTestCase):
class QMNISTTestCase(MNISTTestCase): class QMNISTTestCase(MNISTTestCase):
DATASET_CLASS = datasets.QMNIST DATASET_CLASS = datasets.QMNIST
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(what=("train", "test", "test10k", "nist")) ADDITIONAL_CONFIGS = combinations_grid(what=("train", "test", "test10k", "nist"))
_LABELS_SIZE = (8,) _LABELS_SIZE = (8,)
_LABELS_DTYPE = torch.int32 _LABELS_DTYPE = torch.int32
...@@ -1492,33 +1591,54 @@ class QMNISTTestCase(MNISTTestCase): ...@@ -1492,33 +1591,54 @@ class QMNISTTestCase(MNISTTestCase):
with self.create_dataset(what="test50k") as (dataset, info): with self.create_dataset(what="test50k") as (dataset, info):
# Since the split 'test50k' selects all images beginning from the index 10000, we subtract the number of # Since the split 'test50k' selects all images beginning from the index 10000, we subtract the number of
# created examples by this. # created examples by this.
self.assertEqual(len(dataset), info["num_examples"] - 10000) assert len(dataset) == info["num_examples"] - 10000
class MovingMNISTTestCase(datasets_utils.DatasetTestCase):
DATASET_CLASS = datasets.MovingMNIST
FEATURE_TYPES = (torch.Tensor,)
ADDITIONAL_CONFIGS = combinations_grid(split=(None, "train", "test"), split_ratio=(10, 1, 19))
_NUM_FRAMES = 20
def inject_fake_data(self, tmpdir, config):
base_folder = os.path.join(tmpdir, self.DATASET_CLASS.__name__)
os.makedirs(base_folder, exist_ok=True)
num_samples = 5
data = np.concatenate(
[
np.zeros((config["split_ratio"], num_samples, 64, 64)),
np.ones((self._NUM_FRAMES - config["split_ratio"], num_samples, 64, 64)),
]
)
np.save(os.path.join(base_folder, "mnist_test_seq.npy"), data)
return num_samples
@datasets_utils.test_all_configs
def test_split(self, config):
with self.create_dataset(config) as (dataset, _):
if config["split"] == "train":
assert (dataset.data == 0).all()
elif config["split"] == "test":
assert (dataset.data == 1).all()
else:
assert dataset.data.size()[1] == self._NUM_FRAMES
class DatasetFolderTestCase(datasets_utils.ImageDatasetTestCase): class DatasetFolderTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.DatasetFolder DATASET_CLASS = datasets.DatasetFolder
# The dataset has no fixed return type since it is defined by the loader parameter. For testing, we use a loader _EXTENSIONS = ("jpg", "png")
# that simply returns the path as type 'str' instead of loading anything. See the 'dataset_args()' method.
FEATURE_TYPES = (str, int)
_IMAGE_EXTENSIONS = ("jpg", "png")
_VIDEO_EXTENSIONS = ("avi", "mp4")
_EXTENSIONS = (*_IMAGE_EXTENSIONS, *_VIDEO_EXTENSIONS)
# DatasetFolder has two mutually exclusive parameters: 'extensions' and 'is_valid_file'. One of both is required. # DatasetFolder has two mutually exclusive parameters: 'extensions' and 'is_valid_file'. One of both is required.
# We only iterate over different 'extensions' here and handle the tests for 'is_valid_file' in the # We only iterate over different 'extensions' here and handle the tests for 'is_valid_file' in the
# 'test_is_valid_file()' method. # 'test_is_valid_file()' method.
DEFAULT_CONFIG = dict(extensions=_EXTENSIONS) DEFAULT_CONFIG = dict(extensions=_EXTENSIONS)
ADDITIONAL_CONFIGS = ( ADDITIONAL_CONFIGS = combinations_grid(extensions=[(ext,) for ext in _EXTENSIONS])
*datasets_utils.combinations_grid(extensions=[(ext,) for ext in _IMAGE_EXTENSIONS]),
dict(extensions=_IMAGE_EXTENSIONS),
*datasets_utils.combinations_grid(extensions=[(ext,) for ext in _VIDEO_EXTENSIONS]),
dict(extensions=_VIDEO_EXTENSIONS),
)
def dataset_args(self, tmpdir, config): def dataset_args(self, tmpdir, config):
return tmpdir, lambda x: x return tmpdir, datasets.folder.pil_loader
def inject_fake_data(self, tmpdir, config): def inject_fake_data(self, tmpdir, config):
extensions = config["extensions"] or self._is_valid_file_to_extensions(config["is_valid_file"]) extensions = config["extensions"] or self._is_valid_file_to_extensions(config["is_valid_file"])
...@@ -1529,18 +1649,16 @@ class DatasetFolderTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -1529,18 +1649,16 @@ class DatasetFolderTestCase(datasets_utils.ImageDatasetTestCase):
if ext not in extensions: if ext not in extensions:
continue continue
create_example_folder = (
datasets_utils.create_image_folder
if ext in self._IMAGE_EXTENSIONS
else datasets_utils.create_video_folder
)
num_examples = torch.randint(1, 3, size=()).item() num_examples = torch.randint(1, 3, size=()).item()
create_example_folder(tmpdir, cls, lambda idx: self._file_name_fn(cls, ext, idx), num_examples) datasets_utils.create_image_folder(tmpdir, cls, lambda idx: self._file_name_fn(cls, ext, idx), num_examples)
num_examples_total += num_examples num_examples_total += num_examples
classes.append(cls) classes.append(cls)
if config.pop("make_empty_class", False):
os.makedirs(pathlib.Path(tmpdir) / "empty_class")
classes.append("empty_class")
return dict(num_examples=num_examples_total, classes=classes) return dict(num_examples=num_examples_total, classes=classes)
def _file_name_fn(self, cls, ext, idx): def _file_name_fn(self, cls, ext, idx):
...@@ -1555,14 +1673,32 @@ class DatasetFolderTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -1555,14 +1673,32 @@ class DatasetFolderTestCase(datasets_utils.ImageDatasetTestCase):
# We need to explicitly pass extensions=None here or otherwise it would be filled by the value from the # We need to explicitly pass extensions=None here or otherwise it would be filled by the value from the
# DEFAULT_CONFIG. # DEFAULT_CONFIG.
with self.create_dataset( with self.create_dataset(
config, extensions=None, is_valid_file=lambda file: pathlib.Path(file).suffix[1:] in extensions config, extensions=None, is_valid_file=lambda file: pathlib.Path(file).suffix[1:] in extensions
) as (dataset, info): ) as (dataset, info):
self.assertEqual(len(dataset), info["num_examples"]) assert len(dataset) == info["num_examples"]
@datasets_utils.test_all_configs @datasets_utils.test_all_configs
def test_classes(self, config): def test_classes(self, config):
with self.create_dataset(config) as (dataset, info): with self.create_dataset(config) as (dataset, info):
self.assertSequenceEqual(dataset.classes, info["classes"]) assert len(dataset.classes) == len(info["classes"])
assert all([a == b for a, b in zip(dataset.classes, info["classes"])])
def test_allow_empty(self):
config = {
"extensions": self._EXTENSIONS,
"make_empty_class": True,
}
config["allow_empty"] = True
with self.create_dataset(config) as (dataset, info):
assert "empty_class" in dataset.classes
assert len(dataset.classes) == len(info["classes"])
assert all([a == b for a, b in zip(dataset.classes, info["classes"])])
config["allow_empty"] = False
with pytest.raises(FileNotFoundError, match="Found no valid file"):
with self.create_dataset(config) as (dataset, info):
pass
class ImageFolderTestCase(datasets_utils.ImageDatasetTestCase): class ImageFolderTestCase(datasets_utils.ImageDatasetTestCase):
...@@ -1582,13 +1718,14 @@ class ImageFolderTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -1582,13 +1718,14 @@ class ImageFolderTestCase(datasets_utils.ImageDatasetTestCase):
@datasets_utils.test_all_configs @datasets_utils.test_all_configs
def test_classes(self, config): def test_classes(self, config):
with self.create_dataset(config) as (dataset, info): with self.create_dataset(config) as (dataset, info):
self.assertSequenceEqual(dataset.classes, info["classes"]) assert len(dataset.classes) == len(info["classes"])
assert all([a == b for a, b in zip(dataset.classes, info["classes"])])
class KittiTestCase(datasets_utils.ImageDatasetTestCase): class KittiTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Kitti DATASET_CLASS = datasets.Kitti
FEATURE_TYPES = (PIL.Image.Image, (list, type(None))) # test split returns None as target FEATURE_TYPES = (PIL.Image.Image, (list, type(None))) # test split returns None as target
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(train=(True, False)) ADDITIONAL_CONFIGS = combinations_grid(train=(True, False))
def inject_fake_data(self, tmpdir, config): def inject_fake_data(self, tmpdir, config):
kitti_dir = os.path.join(tmpdir, "Kitti", "raw") kitti_dir = os.path.join(tmpdir, "Kitti", "raw")
...@@ -1620,11 +1757,16 @@ class KittiTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -1620,11 +1757,16 @@ class KittiTestCase(datasets_utils.ImageDatasetTestCase):
return split_to_num_examples[config["train"]] return split_to_num_examples[config["train"]]
def test_transforms_v2_wrapper_spawn(self):
expected_size = (123, 321)
with self.create_dataset(transform=v2.Resize(size=expected_size)) as (dataset, _):
datasets_utils.check_transforms_v2_wrapper_spawn(dataset, expected_size=expected_size)
class SvhnTestCase(datasets_utils.ImageDatasetTestCase): class SvhnTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.SVHN DATASET_CLASS = datasets.SVHN
REQUIRED_PACKAGES = ("scipy",) REQUIRED_PACKAGES = ("scipy",)
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid(split=("train", "test", "extra")) ADDITIONAL_CONFIGS = combinations_grid(split=("train", "test", "extra"))
def inject_fake_data(self, tmpdir, config): def inject_fake_data(self, tmpdir, config):
import scipy.io as sio import scipy.io as sio
...@@ -1639,13 +1781,13 @@ class SvhnTestCase(datasets_utils.ImageDatasetTestCase): ...@@ -1639,13 +1781,13 @@ class SvhnTestCase(datasets_utils.ImageDatasetTestCase):
file = f"{split}_32x32.mat" file = f"{split}_32x32.mat"
images = np.zeros((32, 32, 3, num_examples), dtype=np.uint8) images = np.zeros((32, 32, 3, num_examples), dtype=np.uint8)
targets = np.zeros((num_examples,), dtype=np.uint8) targets = np.zeros((num_examples,), dtype=np.uint8)
sio.savemat(os.path.join(tmpdir, file), {'X': images, 'y': targets}) sio.savemat(os.path.join(tmpdir, file), {"X": images, "y": targets})
return num_examples return num_examples
class Places365TestCase(datasets_utils.ImageDatasetTestCase): class Places365TestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Places365 DATASET_CLASS = datasets.Places365
ADDITIONAL_CONFIGS = datasets_utils.combinations_grid( ADDITIONAL_CONFIGS = combinations_grid(
split=("train-standard", "train-challenge", "val"), split=("train-standard", "train-challenge", "val"),
small=(False, True), small=(False, True),
) )
...@@ -1674,8 +1816,7 @@ class Places365TestCase(datasets_utils.ImageDatasetTestCase): ...@@ -1674,8 +1816,7 @@ class Places365TestCase(datasets_utils.ImageDatasetTestCase):
# (file, idx) # (file, idx)
_FILE_LIST_CONTENT = ( _FILE_LIST_CONTENT = (
("Places365_val_00000001.png", 0), ("Places365_val_00000001.png", 0),
*((f"{category}/Places365_train_00000001.png", idx) *((f"{category}/Places365_train_00000001.png", idx) for category, idx in _CATEGORIES_CONTENT),
for category, idx in _CATEGORIES_CONTENT),
) )
@staticmethod @staticmethod
...@@ -1715,24 +1856,1699 @@ class Places365TestCase(datasets_utils.ImageDatasetTestCase): ...@@ -1715,24 +1856,1699 @@ class Places365TestCase(datasets_utils.ImageDatasetTestCase):
return [(os.path.join(root, folder_name, image), idx) for image, idx in zip(images, idcs)] return [(os.path.join(root, folder_name, image), idx) for image, idx in zip(images, idcs)]
def inject_fake_data(self, tmpdir, config): def inject_fake_data(self, tmpdir, config):
self._make_devkit_archive(tmpdir, config['split']) self._make_devkit_archive(tmpdir, config["split"])
return len(self._make_images_archive(tmpdir, config['split'], config['small'])) return len(self._make_images_archive(tmpdir, config["split"], config["small"]))
def test_classes(self): def test_classes(self):
classes = list(map(lambda x: x[0], self._CATEGORIES_CONTENT)) classes = list(map(lambda x: x[0], self._CATEGORIES_CONTENT))
with self.create_dataset() as (dataset, _): with self.create_dataset() as (dataset, _):
self.assertEqual(dataset.classes, classes) assert dataset.classes == classes
def test_class_to_idx(self): def test_class_to_idx(self):
class_to_idx = dict(self._CATEGORIES_CONTENT) class_to_idx = dict(self._CATEGORIES_CONTENT)
with self.create_dataset() as (dataset, _): with self.create_dataset() as (dataset, _):
self.assertEqual(dataset.class_to_idx, class_to_idx) assert dataset.class_to_idx == class_to_idx
def test_images_download_preexisting(self): def test_images_download_preexisting(self):
with self.assertRaises(RuntimeError): with pytest.raises(RuntimeError):
with self.create_dataset({'download': True}): with self.create_dataset({"download": True}):
pass
class INaturalistTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.INaturalist
FEATURE_TYPES = (PIL.Image.Image, (int, tuple))
ADDITIONAL_CONFIGS = combinations_grid(
target_type=("kingdom", "full", "genus", ["kingdom", "phylum", "class", "order", "family", "genus", "full"]),
version=("2021_train",),
)
def inject_fake_data(self, tmpdir, config):
categories = [
"00000_Akingdom_0phylum_Aclass_Aorder_Afamily_Agenus_Aspecies",
"00001_Akingdom_1phylum_Aclass_Border_Afamily_Bgenus_Aspecies",
"00002_Akingdom_2phylum_Cclass_Corder_Cfamily_Cgenus_Cspecies",
]
num_images_per_category = 3
for category in categories:
datasets_utils.create_image_folder(
root=os.path.join(tmpdir, config["version"]),
name=category,
file_name_fn=lambda idx: f"image_{idx + 1:04d}.jpg",
num_examples=num_images_per_category,
)
return num_images_per_category * len(categories)
def test_targets(self):
target_types = ["kingdom", "phylum", "class", "order", "family", "genus", "full"]
with self.create_dataset(target_type=target_types, version="2021_valid") as (dataset, _):
items = [d[1] for d in dataset]
for i, item in enumerate(items):
assert dataset.category_name("kingdom", item[0]) == "Akingdom"
assert dataset.category_name("phylum", item[1]) == f"{i // 3}phylum"
assert item[6] == i // 3
class LFWPeopleTestCase(datasets_utils.DatasetTestCase):
DATASET_CLASS = datasets.LFWPeople
FEATURE_TYPES = (PIL.Image.Image, int)
ADDITIONAL_CONFIGS = combinations_grid(
split=("10fold", "train", "test"), image_set=("original", "funneled", "deepfunneled")
)
_IMAGES_DIR = {"original": "lfw", "funneled": "lfw_funneled", "deepfunneled": "lfw-deepfunneled"}
_file_id = {"10fold": "", "train": "DevTrain", "test": "DevTest"}
def inject_fake_data(self, tmpdir, config):
tmpdir = pathlib.Path(tmpdir) / "lfw-py"
os.makedirs(tmpdir, exist_ok=True)
return dict(
num_examples=self._create_images_dir(tmpdir, self._IMAGES_DIR[config["image_set"]], config["split"]),
split=config["split"],
)
def _create_images_dir(self, root, idir, split):
idir = os.path.join(root, idir)
os.makedirs(idir, exist_ok=True)
n, flines = (10, ["10\n"]) if split == "10fold" else (1, [])
num_examples = 0
names = []
for _ in range(n):
num_people = random.randint(2, 5)
flines.append(f"{num_people}\n")
for i in range(num_people):
name = self._create_random_id()
no = random.randint(1, 10)
flines.append(f"{name}\t{no}\n")
names.append(f"{name}\t{no}\n")
datasets_utils.create_image_folder(idir, name, lambda n: f"{name}_{n+1:04d}.jpg", no, 250)
num_examples += no
with open(pathlib.Path(root) / f"people{self._file_id[split]}.txt", "w") as f:
f.writelines(flines)
with open(pathlib.Path(root) / "lfw-names.txt", "w") as f:
f.writelines(sorted(names))
return num_examples
def _create_random_id(self):
part1 = datasets_utils.create_random_string(random.randint(5, 7))
part2 = datasets_utils.create_random_string(random.randint(4, 7))
return f"{part1}_{part2}"
class LFWPairsTestCase(LFWPeopleTestCase):
DATASET_CLASS = datasets.LFWPairs
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, int)
def _create_images_dir(self, root, idir, split):
idir = os.path.join(root, idir)
os.makedirs(idir, exist_ok=True)
num_pairs = 7 # effectively 7*2*n = 14*n
n, self.flines = (10, [f"10\t{num_pairs}"]) if split == "10fold" else (1, [str(num_pairs)])
for _ in range(n):
self._inject_pairs(idir, num_pairs, True)
self._inject_pairs(idir, num_pairs, False)
with open(pathlib.Path(root) / f"pairs{self._file_id[split]}.txt", "w") as f:
f.writelines(self.flines)
return num_pairs * 2 * n
def _inject_pairs(self, root, num_pairs, same):
for i in range(num_pairs):
name1 = self._create_random_id()
name2 = name1 if same else self._create_random_id()
no1, no2 = random.randint(1, 100), random.randint(1, 100)
if same:
self.flines.append(f"\n{name1}\t{no1}\t{no2}")
else:
self.flines.append(f"\n{name1}\t{no1}\t{name2}\t{no2}")
datasets_utils.create_image_folder(root, name1, lambda _: f"{name1}_{no1:04d}.jpg", 1, 250)
datasets_utils.create_image_folder(root, name2, lambda _: f"{name2}_{no2:04d}.jpg", 1, 250)
class SintelTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Sintel
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "test"), pass_name=("clean", "final", "both"))
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)))
FLOW_H, FLOW_W = 3, 4
def inject_fake_data(self, tmpdir, config):
root = pathlib.Path(tmpdir) / "Sintel"
num_images_per_scene = 3 if config["split"] == "train" else 4
num_scenes = 2
for split_dir in ("training", "test"):
for pass_name in ("clean", "final"):
image_root = root / split_dir / pass_name
for scene_id in range(num_scenes):
scene_dir = image_root / f"scene_{scene_id}"
datasets_utils.create_image_folder(
image_root,
name=str(scene_dir),
file_name_fn=lambda image_idx: f"frame_000{image_idx}.png",
num_examples=num_images_per_scene,
)
flow_root = root / "training" / "flow"
for scene_id in range(num_scenes):
scene_dir = flow_root / f"scene_{scene_id}"
os.makedirs(scene_dir)
for i in range(num_images_per_scene - 1):
file_name = str(scene_dir / f"frame_000{i}.flo")
datasets_utils.make_fake_flo_file(h=self.FLOW_H, w=self.FLOW_W, file_name=file_name)
# with e.g. num_images_per_scene = 3, for a single scene with have 3 images
# which are frame_0000, frame_0001 and frame_0002
# They will be consecutively paired as (frame_0000, frame_0001), (frame_0001, frame_0002),
# that is 3 - 1 = 2 examples. Hence the formula below
num_passes = 2 if config["pass_name"] == "both" else 1
num_examples = (num_images_per_scene - 1) * num_scenes * num_passes
return num_examples
def test_flow(self):
# Make sure flow exists for train split, and make sure there are as many flow values as (pairs of) images
h, w = self.FLOW_H, self.FLOW_W
expected_flow = np.arange(2 * h * w).reshape(h, w, 2).transpose(2, 0, 1)
with self.create_dataset(split="train") as (dataset, _):
assert dataset._flow_list and len(dataset._flow_list) == len(dataset._image_list)
for _, _, flow in dataset:
assert flow.shape == (2, h, w)
np.testing.assert_allclose(flow, expected_flow)
# Make sure flow is always None for test split
with self.create_dataset(split="test") as (dataset, _):
assert dataset._image_list and not dataset._flow_list
for _, _, flow in dataset:
assert flow is None
def test_bad_input(self):
with pytest.raises(ValueError, match="Unknown value 'bad' for argument split"):
with self.create_dataset(split="bad"):
pass
with pytest.raises(ValueError, match="Unknown value 'bad' for argument pass_name"):
with self.create_dataset(pass_name="bad"):
pass
class KittiFlowTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.KittiFlow
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "test"))
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)), (np.ndarray, type(None)))
def inject_fake_data(self, tmpdir, config):
root = pathlib.Path(tmpdir) / "KittiFlow"
num_examples = 2 if config["split"] == "train" else 3
for split_dir in ("training", "testing"):
datasets_utils.create_image_folder(
root / split_dir,
name="image_2",
file_name_fn=lambda image_idx: f"{image_idx}_10.png",
num_examples=num_examples,
)
datasets_utils.create_image_folder(
root / split_dir,
name="image_2",
file_name_fn=lambda image_idx: f"{image_idx}_11.png",
num_examples=num_examples,
)
# For kitti the ground truth flows are encoded as 16-bits pngs.
# create_image_folder() will actually create 8-bits pngs, but it doesn't
# matter much: the flow reader will still be able to read the files, it
# will just be garbage flow value - but we don't care about that here.
datasets_utils.create_image_folder(
root / "training",
name="flow_occ",
file_name_fn=lambda image_idx: f"{image_idx}_10.png",
num_examples=num_examples,
)
return num_examples
def test_flow_and_valid(self):
# Make sure flow exists for train split, and make sure there are as many flow values as (pairs of) images
# Also assert flow and valid are of the expected shape
with self.create_dataset(split="train") as (dataset, _):
assert dataset._flow_list and len(dataset._flow_list) == len(dataset._image_list)
for _, _, flow, valid in dataset:
two, h, w = flow.shape
assert two == 2
assert valid.shape == (h, w)
# Make sure flow and valid are always None for test split
with self.create_dataset(split="test") as (dataset, _):
assert dataset._image_list and not dataset._flow_list
for _, _, flow, valid in dataset:
assert flow is None
assert valid is None
def test_bad_input(self):
with pytest.raises(ValueError, match="Unknown value 'bad' for argument split"):
with self.create_dataset(split="bad"):
pass
class FlyingChairsTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.FlyingChairs
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "val"))
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)))
FLOW_H, FLOW_W = 3, 4
def _make_split_file(self, root, num_examples):
# We create a fake split file here, but users are asked to download the real one from the authors website
split_ids = [1] * num_examples["train"] + [2] * num_examples["val"]
random.shuffle(split_ids)
with open(str(root / "FlyingChairs_train_val.txt"), "w+") as split_file:
for split_id in split_ids:
split_file.write(f"{split_id}\n")
def inject_fake_data(self, tmpdir, config):
root = pathlib.Path(tmpdir) / "FlyingChairs"
num_examples = {"train": 5, "val": 3}
num_examples_total = sum(num_examples.values())
datasets_utils.create_image_folder( # img1
root,
name="data",
file_name_fn=lambda image_idx: f"00{image_idx}_img1.ppm",
num_examples=num_examples_total,
)
datasets_utils.create_image_folder( # img2
root,
name="data",
file_name_fn=lambda image_idx: f"00{image_idx}_img2.ppm",
num_examples=num_examples_total,
)
for i in range(num_examples_total):
file_name = str(root / "data" / f"00{i}_flow.flo")
datasets_utils.make_fake_flo_file(h=self.FLOW_H, w=self.FLOW_W, file_name=file_name)
self._make_split_file(root, num_examples)
return num_examples[config["split"]]
@datasets_utils.test_all_configs
def test_flow(self, config):
# Make sure flow always exists, and make sure there are as many flow values as (pairs of) images
# Also make sure the flow is properly decoded
h, w = self.FLOW_H, self.FLOW_W
expected_flow = np.arange(2 * h * w).reshape(h, w, 2).transpose(2, 0, 1)
with self.create_dataset(config=config) as (dataset, _):
assert dataset._flow_list and len(dataset._flow_list) == len(dataset._image_list)
for _, _, flow in dataset:
assert flow.shape == (2, h, w)
np.testing.assert_allclose(flow, expected_flow)
class FlyingThings3DTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.FlyingThings3D
ADDITIONAL_CONFIGS = combinations_grid(
split=("train", "test"), pass_name=("clean", "final", "both"), camera=("left", "right", "both")
)
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)))
FLOW_H, FLOW_W = 3, 4
def inject_fake_data(self, tmpdir, config):
root = pathlib.Path(tmpdir) / "FlyingThings3D"
num_images_per_camera = 3 if config["split"] == "train" else 4
passes = ("frames_cleanpass", "frames_finalpass")
splits = ("TRAIN", "TEST")
letters = ("A", "B", "C")
subfolders = ("0000", "0001")
cameras = ("left", "right")
for pass_name, split, letter, subfolder, camera in itertools.product(
passes, splits, letters, subfolders, cameras
):
current_folder = root / pass_name / split / letter / subfolder
datasets_utils.create_image_folder(
current_folder,
name=camera,
file_name_fn=lambda image_idx: f"00{image_idx}.png",
num_examples=num_images_per_camera,
)
directions = ("into_future", "into_past")
for split, letter, subfolder, direction, camera in itertools.product(
splits, letters, subfolders, directions, cameras
):
current_folder = root / "optical_flow" / split / letter / subfolder / direction / camera
os.makedirs(str(current_folder), exist_ok=True)
for i in range(num_images_per_camera):
datasets_utils.make_fake_pfm_file(self.FLOW_H, self.FLOW_W, file_name=str(current_folder / f"{i}.pfm"))
num_cameras = 2 if config["camera"] == "both" else 1
num_passes = 2 if config["pass_name"] == "both" else 1
num_examples = (
(num_images_per_camera - 1) * num_cameras * len(subfolders) * len(letters) * len(splits) * num_passes
)
return num_examples
@datasets_utils.test_all_configs
def test_flow(self, config):
h, w = self.FLOW_H, self.FLOW_W
expected_flow = np.arange(3 * h * w).reshape(h, w, 3).transpose(2, 0, 1)
expected_flow = np.flip(expected_flow, axis=1)
expected_flow = expected_flow[:2, :, :]
with self.create_dataset(config=config) as (dataset, _):
assert dataset._flow_list and len(dataset._flow_list) == len(dataset._image_list)
for _, _, flow in dataset:
assert flow.shape == (2, self.FLOW_H, self.FLOW_W)
np.testing.assert_allclose(flow, expected_flow)
def test_bad_input(self):
with pytest.raises(ValueError, match="Unknown value 'bad' for argument split"):
with self.create_dataset(split="bad"):
pass
with pytest.raises(ValueError, match="Unknown value 'bad' for argument pass_name"):
with self.create_dataset(pass_name="bad"):
pass
with pytest.raises(ValueError, match="Unknown value 'bad' for argument camera"):
with self.create_dataset(camera="bad"):
pass pass
class HD1KTestCase(KittiFlowTestCase):
DATASET_CLASS = datasets.HD1K
def inject_fake_data(self, tmpdir, config):
root = pathlib.Path(tmpdir) / "hd1k"
num_sequences = 4 if config["split"] == "train" else 3
num_examples_per_train_sequence = 3
for seq_idx in range(num_sequences):
# Training data
datasets_utils.create_image_folder(
root / "hd1k_input",
name="image_2",
file_name_fn=lambda image_idx: f"{seq_idx:06d}_{image_idx}.png",
num_examples=num_examples_per_train_sequence,
)
datasets_utils.create_image_folder(
root / "hd1k_flow_gt",
name="flow_occ",
file_name_fn=lambda image_idx: f"{seq_idx:06d}_{image_idx}.png",
num_examples=num_examples_per_train_sequence,
)
# Test data
datasets_utils.create_image_folder(
root / "hd1k_challenge",
name="image_2",
file_name_fn=lambda _: f"{seq_idx:06d}_10.png",
num_examples=1,
)
datasets_utils.create_image_folder(
root / "hd1k_challenge",
name="image_2",
file_name_fn=lambda _: f"{seq_idx:06d}_11.png",
num_examples=1,
)
num_examples_per_sequence = num_examples_per_train_sequence if config["split"] == "train" else 2
return num_sequences * (num_examples_per_sequence - 1)
class EuroSATTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.EuroSAT
FEATURE_TYPES = (PIL.Image.Image, int)
def inject_fake_data(self, tmpdir, config):
data_folder = os.path.join(tmpdir, "eurosat", "2750")
os.makedirs(data_folder)
num_examples_per_class = 3
classes = ("AnnualCrop", "Forest")
for cls in classes:
datasets_utils.create_image_folder(
root=data_folder,
name=cls,
file_name_fn=lambda idx: f"{cls}_{idx}.jpg",
num_examples=num_examples_per_class,
)
return len(classes) * num_examples_per_class
class Food101TestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Food101
FEATURE_TYPES = (PIL.Image.Image, int)
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "test"))
def inject_fake_data(self, tmpdir: str, config):
root_folder = pathlib.Path(tmpdir) / "food-101"
image_folder = root_folder / "images"
meta_folder = root_folder / "meta"
image_folder.mkdir(parents=True)
meta_folder.mkdir()
num_images_per_class = 5
metadata = {}
n_samples_per_class = 3 if config["split"] == "train" else 2
sampled_classes = ("apple_pie", "crab_cakes", "gyoza")
for cls in sampled_classes:
im_fnames = datasets_utils.create_image_folder(
image_folder,
cls,
file_name_fn=lambda idx: f"{idx}.jpg",
num_examples=num_images_per_class,
)
metadata[cls] = [
"/".join(fname.relative_to(image_folder).with_suffix("").parts)
for fname in random.choices(im_fnames, k=n_samples_per_class)
]
with open(meta_folder / f"{config['split']}.json", "w") as file:
file.write(json.dumps(metadata))
return len(sampled_classes * n_samples_per_class)
class FGVCAircraftTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.FGVCAircraft
ADDITIONAL_CONFIGS = combinations_grid(
split=("train", "val", "trainval", "test"), annotation_level=("variant", "family", "manufacturer")
)
def inject_fake_data(self, tmpdir: str, config):
split = config["split"]
annotation_level = config["annotation_level"]
annotation_level_to_file = {
"variant": "variants.txt",
"family": "families.txt",
"manufacturer": "manufacturers.txt",
}
root_folder = pathlib.Path(tmpdir) / "fgvc-aircraft-2013b"
data_folder = root_folder / "data"
classes = ["707-320", "Hawk T1", "Tornado"]
num_images_per_class = 5
datasets_utils.create_image_folder(
data_folder,
"images",
file_name_fn=lambda idx: f"{idx}.jpg",
num_examples=num_images_per_class * len(classes),
)
annotation_file = data_folder / annotation_level_to_file[annotation_level]
with open(annotation_file, "w") as file:
file.write("\n".join(classes))
num_samples_per_class = 4 if split == "trainval" else 2
images_classes = []
for i in range(len(classes)):
images_classes.extend(
[
f"{idx} {classes[i]}"
for idx in random.sample(
range(i * num_images_per_class, (i + 1) * num_images_per_class), num_samples_per_class
)
]
)
images_annotation_file = data_folder / f"images_{annotation_level}_{split}.txt"
with open(images_annotation_file, "w") as file:
file.write("\n".join(images_classes))
return len(classes * num_samples_per_class)
class SUN397TestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.SUN397
def inject_fake_data(self, tmpdir: str, config):
data_dir = pathlib.Path(tmpdir) / "SUN397"
data_dir.mkdir()
num_images_per_class = 5
sampled_classes = ("abbey", "airplane_cabin", "airport_terminal")
im_paths = []
for cls in sampled_classes:
image_folder = data_dir / cls[0]
im_paths.extend(
datasets_utils.create_image_folder(
image_folder,
image_folder / cls,
file_name_fn=lambda idx: f"sun_{idx}.jpg",
num_examples=num_images_per_class,
)
)
with open(data_dir / "ClassName.txt", "w") as file:
file.writelines("\n".join(f"/{cls[0]}/{cls}" for cls in sampled_classes))
num_samples = len(im_paths)
return num_samples
class DTDTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.DTD
FEATURE_TYPES = (PIL.Image.Image, int)
ADDITIONAL_CONFIGS = combinations_grid(
split=("train", "test", "val"),
# There is no need to test the whole matrix here, since each fold is treated exactly the same
partition=(1, 5, 10),
)
def inject_fake_data(self, tmpdir: str, config):
data_folder = pathlib.Path(tmpdir) / "dtd" / "dtd"
num_images_per_class = 3
image_folder = data_folder / "images"
image_files = []
for cls in ("banded", "marbled", "zigzagged"):
image_files.extend(
datasets_utils.create_image_folder(
image_folder,
cls,
file_name_fn=lambda idx: f"{cls}_{idx:04d}.jpg",
num_examples=num_images_per_class,
)
)
meta_folder = data_folder / "labels"
meta_folder.mkdir()
image_ids = [str(path.relative_to(path.parents[1])).replace(os.sep, "/") for path in image_files]
image_ids_in_config = random.choices(image_ids, k=len(image_files) // 2)
with open(meta_folder / f"{config['split']}{config['partition']}.txt", "w") as file:
file.write("\n".join(image_ids_in_config) + "\n")
return len(image_ids_in_config)
class FER2013TestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.FER2013
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "test"))
FEATURE_TYPES = (PIL.Image.Image, (int, type(None)))
def inject_fake_data(self, tmpdir, config):
base_folder = os.path.join(tmpdir, "fer2013")
os.makedirs(base_folder)
use_icml = config.pop("use_icml", False)
use_fer = config.pop("use_fer", False)
num_samples = 5
if use_icml or use_fer:
pixels_key, usage_key = (" pixels", " Usage") if use_icml else ("pixels", "Usage")
fieldnames = ("emotion", usage_key, pixels_key) if use_icml else ("emotion", pixels_key, usage_key)
filename = "icml_face_data.csv" if use_icml else "fer2013.csv"
with open(os.path.join(base_folder, filename), "w", newline="") as file:
writer = csv.DictWriter(
file,
fieldnames=fieldnames,
quoting=csv.QUOTE_NONNUMERIC,
quotechar='"',
)
writer.writeheader()
for i in range(num_samples):
row = {
"emotion": str(int(torch.randint(0, 7, ()))),
usage_key: "Training" if i % 2 else "PublicTest",
pixels_key: " ".join(
str(pixel)
for pixel in datasets_utils.create_image_or_video_tensor((48, 48)).view(-1).tolist()
),
}
writer.writerow(row)
else:
with open(os.path.join(base_folder, f"{config['split']}.csv"), "w", newline="") as file:
writer = csv.DictWriter(
file,
fieldnames=("emotion", "pixels") if config["split"] == "train" else ("pixels",),
quoting=csv.QUOTE_NONNUMERIC,
quotechar='"',
)
writer.writeheader()
for _ in range(num_samples):
row = dict(
pixels=" ".join(
str(pixel)
for pixel in datasets_utils.create_image_or_video_tensor((48, 48)).view(-1).tolist()
)
)
if config["split"] == "train":
row["emotion"] = str(int(torch.randint(0, 7, ())))
writer.writerow(row)
return num_samples
def test_icml_file(self):
config = {"split": "test"}
with self.create_dataset(config=config) as (dataset, _):
assert all(s[1] is None for s in dataset)
for split in ("train", "test"):
for d in ({"use_icml": True}, {"use_fer": True}):
config = {"split": split, **d}
with self.create_dataset(config=config) as (dataset, _):
assert all(s[1] is not None for s in dataset)
class GTSRBTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.GTSRB
FEATURE_TYPES = (PIL.Image.Image, int)
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "test"))
def inject_fake_data(self, tmpdir: str, config):
root_folder = os.path.join(tmpdir, "gtsrb")
os.makedirs(root_folder, exist_ok=True)
# Train data
train_folder = os.path.join(root_folder, "GTSRB", "Training")
os.makedirs(train_folder, exist_ok=True)
num_examples = 3 if config["split"] == "train" else 4
classes = ("00000", "00042", "00012")
for class_idx in classes:
datasets_utils.create_image_folder(
train_folder,
name=class_idx,
file_name_fn=lambda image_idx: f"{class_idx}_{image_idx:05d}.ppm",
num_examples=num_examples,
)
total_number_of_examples = num_examples * len(classes)
# Test data
test_folder = os.path.join(root_folder, "GTSRB", "Final_Test", "Images")
os.makedirs(test_folder, exist_ok=True)
with open(os.path.join(root_folder, "GT-final_test.csv"), "w") as csv_file:
csv_file.write("Filename;Width;Height;Roi.X1;Roi.Y1;Roi.X2;Roi.Y2;ClassId\n")
for _ in range(total_number_of_examples):
image_file = datasets_utils.create_random_string(5, string.digits) + ".ppm"
datasets_utils.create_image_file(test_folder, image_file)
row = [
image_file,
torch.randint(1, 100, size=()).item(),
torch.randint(1, 100, size=()).item(),
torch.randint(1, 100, size=()).item(),
torch.randint(1, 100, size=()).item(),
torch.randint(1, 100, size=()).item(),
torch.randint(1, 100, size=()).item(),
torch.randint(0, 43, size=()).item(),
]
csv_file.write(";".join(map(str, row)) + "\n")
return total_number_of_examples
class CLEVRClassificationTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.CLEVRClassification
FEATURE_TYPES = (PIL.Image.Image, (int, type(None)))
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "val", "test"))
def inject_fake_data(self, tmpdir, config):
data_folder = pathlib.Path(tmpdir) / "clevr" / "CLEVR_v1.0"
images_folder = data_folder / "images"
image_files = datasets_utils.create_image_folder(
images_folder, config["split"], lambda idx: f"CLEVR_{config['split']}_{idx:06d}.png", num_examples=5
)
scenes_folder = data_folder / "scenes"
scenes_folder.mkdir()
if config["split"] != "test":
with open(scenes_folder / f"CLEVR_{config['split']}_scenes.json", "w") as file:
json.dump(
dict(
info=dict(),
scenes=[
dict(image_filename=image_file.name, objects=[dict()] * int(torch.randint(10, ())))
for image_file in image_files
],
),
file,
)
return len(image_files)
class OxfordIIITPetTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.OxfordIIITPet
FEATURE_TYPES = (PIL.Image.Image, (int, PIL.Image.Image, tuple, type(None)))
ADDITIONAL_CONFIGS = combinations_grid(
split=("trainval", "test"),
target_types=("category", "binary-category", "segmentation", ["category", "segmentation"], []),
)
def inject_fake_data(self, tmpdir, config):
base_folder = os.path.join(tmpdir, "oxford-iiit-pet")
classification_anns_meta = (
dict(cls="Abyssinian", label=0, species="cat"),
dict(cls="Keeshond", label=18, species="dog"),
dict(cls="Yorkshire Terrier", label=37, species="dog"),
)
split_and_classification_anns = [
self._meta_to_split_and_classification_ann(meta, idx)
for meta, idx in itertools.product(classification_anns_meta, (1, 2, 10))
]
image_ids, *_ = zip(*split_and_classification_anns)
image_files = datasets_utils.create_image_folder(
base_folder, "images", file_name_fn=lambda idx: f"{image_ids[idx]}.jpg", num_examples=len(image_ids)
)
anns_folder = os.path.join(base_folder, "annotations")
os.makedirs(anns_folder)
split_and_classification_anns_in_split = random.choices(split_and_classification_anns, k=len(image_ids) // 2)
with open(os.path.join(anns_folder, f"{config['split']}.txt"), "w", newline="") as file:
writer = csv.writer(file, delimiter=" ")
for split_and_classification_ann in split_and_classification_anns_in_split:
writer.writerow(split_and_classification_ann)
segmentation_files = datasets_utils.create_image_folder(
anns_folder, "trimaps", file_name_fn=lambda idx: f"{image_ids[idx]}.png", num_examples=len(image_ids)
)
# The dataset has some rogue files
for path in image_files[:2]:
path.with_suffix(".mat").touch()
for path in segmentation_files:
path.with_name(f".{path.name}").touch()
return len(split_and_classification_anns_in_split)
def _meta_to_split_and_classification_ann(self, meta, idx):
image_id = "_".join(
[
*[(str.title if meta["species"] == "cat" else str.lower)(part) for part in meta["cls"].split()],
str(idx),
]
)
class_id = str(meta["label"] + 1)
species = "1" if meta["species"] == "cat" else "2"
breed_id = "-1"
return (image_id, class_id, species, breed_id)
def test_transforms_v2_wrapper_spawn(self):
expected_size = (123, 321)
with self.create_dataset(transform=v2.Resize(size=expected_size)) as (dataset, _):
datasets_utils.check_transforms_v2_wrapper_spawn(dataset, expected_size=expected_size)
class StanfordCarsTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.StanfordCars
REQUIRED_PACKAGES = ("scipy",)
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "test"))
def inject_fake_data(self, tmpdir, config):
import scipy.io as io
from numpy.core.records import fromarrays
num_examples = {"train": 5, "test": 7}[config["split"]]
num_classes = 3
base_folder = pathlib.Path(tmpdir) / "stanford_cars"
devkit = base_folder / "devkit"
devkit.mkdir(parents=True)
if config["split"] == "train":
images_folder_name = "cars_train"
annotations_mat_path = devkit / "cars_train_annos.mat"
else:
images_folder_name = "cars_test"
annotations_mat_path = base_folder / "cars_test_annos_withlabels.mat"
datasets_utils.create_image_folder(
root=base_folder,
name=images_folder_name,
file_name_fn=lambda image_index: f"{image_index:5d}.jpg",
num_examples=num_examples,
)
classes = np.random.randint(1, num_classes + 1, num_examples, dtype=np.uint8)
fnames = [f"{i:5d}.jpg" for i in range(num_examples)]
rec_array = fromarrays(
[classes, fnames],
names=["class", "fname"],
)
io.savemat(annotations_mat_path, {"annotations": rec_array})
random_class_names = ["random_name"] * num_classes
io.savemat(devkit / "cars_meta.mat", {"class_names": random_class_names})
return num_examples
class Country211TestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Country211
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "valid", "test"))
def inject_fake_data(self, tmpdir: str, config):
split_folder = pathlib.Path(tmpdir) / "country211" / config["split"]
split_folder.mkdir(parents=True, exist_ok=True)
num_examples = {
"train": 3,
"valid": 4,
"test": 5,
}[config["split"]]
classes = ("AD", "BS", "GR")
for cls in classes:
datasets_utils.create_image_folder(
split_folder,
name=cls,
file_name_fn=lambda idx: f"{idx}.jpg",
num_examples=num_examples,
)
return num_examples * len(classes)
class Flowers102TestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Flowers102
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "val", "test"))
REQUIRED_PACKAGES = ("scipy",)
def inject_fake_data(self, tmpdir: str, config):
base_folder = pathlib.Path(tmpdir) / "flowers-102"
num_classes = 3
num_images_per_split = dict(train=5, val=4, test=3)
num_images_total = sum(num_images_per_split.values())
datasets_utils.create_image_folder(
base_folder,
"jpg",
file_name_fn=lambda idx: f"image_{idx + 1:05d}.jpg",
num_examples=num_images_total,
)
label_dict = dict(
labels=np.random.randint(1, num_classes + 1, size=(1, num_images_total), dtype=np.uint8),
)
datasets_utils.lazy_importer.scipy.io.savemat(str(base_folder / "imagelabels.mat"), label_dict)
setid_mat = np.arange(1, num_images_total + 1, dtype=np.uint16)
np.random.shuffle(setid_mat)
setid_dict = dict(
trnid=setid_mat[: num_images_per_split["train"]].reshape(1, -1),
valid=setid_mat[num_images_per_split["train"] : -num_images_per_split["test"]].reshape(1, -1),
tstid=setid_mat[-num_images_per_split["test"] :].reshape(1, -1),
)
datasets_utils.lazy_importer.scipy.io.savemat(str(base_folder / "setid.mat"), setid_dict)
return num_images_per_split[config["split"]]
class PCAMTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.PCAM
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "val", "test"))
REQUIRED_PACKAGES = ("h5py",)
def inject_fake_data(self, tmpdir: str, config):
base_folder = pathlib.Path(tmpdir) / "pcam"
base_folder.mkdir()
num_images = {"train": 2, "test": 3, "val": 4}[config["split"]]
images_file = datasets.PCAM._FILES[config["split"]]["images"][0]
with datasets_utils.lazy_importer.h5py.File(str(base_folder / images_file), "w") as f:
f["x"] = np.random.randint(0, 256, size=(num_images, 10, 10, 3), dtype=np.uint8)
targets_file = datasets.PCAM._FILES[config["split"]]["targets"][0]
with datasets_utils.lazy_importer.h5py.File(str(base_folder / targets_file), "w") as f:
f["y"] = np.random.randint(0, 2, size=(num_images, 1, 1, 1), dtype=np.uint8)
return num_images
class RenderedSST2TestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.RenderedSST2
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "val", "test"))
SPLIT_TO_FOLDER = {"train": "train", "val": "valid", "test": "test"}
def inject_fake_data(self, tmpdir: str, config):
root_folder = pathlib.Path(tmpdir) / "rendered-sst2"
image_folder = root_folder / self.SPLIT_TO_FOLDER[config["split"]]
num_images_per_class = {"train": 5, "test": 6, "val": 7}
sampled_classes = ["positive", "negative"]
for cls in sampled_classes:
datasets_utils.create_image_folder(
image_folder,
cls,
file_name_fn=lambda idx: f"{idx}.png",
num_examples=num_images_per_class[config["split"]],
)
return len(sampled_classes) * num_images_per_class[config["split"]]
class Kitti2012StereoTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Kitti2012Stereo
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "test"))
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)), (np.ndarray, type(None)))
def inject_fake_data(self, tmpdir, config):
kitti_dir = pathlib.Path(tmpdir) / "Kitti2012"
os.makedirs(kitti_dir, exist_ok=True)
split_dir = kitti_dir / (config["split"] + "ing")
os.makedirs(split_dir, exist_ok=True)
num_examples = {"train": 4, "test": 3}.get(config["split"], 0)
datasets_utils.create_image_folder(
root=split_dir,
name="colored_0",
file_name_fn=lambda i: f"{i:06d}_10.png",
num_examples=num_examples,
size=(3, 100, 200),
)
datasets_utils.create_image_folder(
root=split_dir,
name="colored_1",
file_name_fn=lambda i: f"{i:06d}_10.png",
num_examples=num_examples,
size=(3, 100, 200),
)
if config["split"] == "train":
datasets_utils.create_image_folder(
root=split_dir,
name="disp_noc",
file_name_fn=lambda i: f"{i:06d}.png",
num_examples=num_examples,
# Kitti2012 uses a single channel image for disparities
size=(1, 100, 200),
)
return num_examples
def test_train_splits(self):
for split in ["train"]:
with self.create_dataset(split=split) as (dataset, _):
for left, right, disparity, mask in dataset:
assert mask is None
datasets_utils.shape_test_for_stereo(left, right, disparity)
def test_test_split(self):
for split in ["test"]:
with self.create_dataset(split=split) as (dataset, _):
for left, right, disparity, mask in dataset:
assert mask is None
assert disparity is None
datasets_utils.shape_test_for_stereo(left, right)
def test_bad_input(self):
with pytest.raises(ValueError, match="Unknown value 'bad' for argument split"):
with self.create_dataset(split="bad"):
pass
class Kitti2015StereoTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Kitti2015Stereo
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "test"))
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)), (np.ndarray, type(None)))
def inject_fake_data(self, tmpdir, config):
kitti_dir = pathlib.Path(tmpdir) / "Kitti2015"
os.makedirs(kitti_dir, exist_ok=True)
split_dir = kitti_dir / (config["split"] + "ing")
os.makedirs(split_dir, exist_ok=True)
num_examples = {"train": 4, "test": 6}.get(config["split"], 0)
datasets_utils.create_image_folder(
root=split_dir,
name="image_2",
file_name_fn=lambda i: f"{i:06d}_10.png",
num_examples=num_examples,
size=(3, 100, 200),
)
datasets_utils.create_image_folder(
root=split_dir,
name="image_3",
file_name_fn=lambda i: f"{i:06d}_10.png",
num_examples=num_examples,
size=(3, 100, 200),
)
if config["split"] == "train":
datasets_utils.create_image_folder(
root=split_dir,
name="disp_occ_0",
file_name_fn=lambda i: f"{i:06d}.png",
num_examples=num_examples,
# Kitti2015 uses a single channel image for disparities
size=(1, 100, 200),
)
datasets_utils.create_image_folder(
root=split_dir,
name="disp_occ_1",
file_name_fn=lambda i: f"{i:06d}.png",
num_examples=num_examples,
# Kitti2015 uses a single channel image for disparities
size=(1, 100, 200),
)
return num_examples
def test_train_splits(self):
for split in ["train"]:
with self.create_dataset(split=split) as (dataset, _):
for left, right, disparity, mask in dataset:
assert mask is None
datasets_utils.shape_test_for_stereo(left, right, disparity)
def test_test_split(self):
for split in ["test"]:
with self.create_dataset(split=split) as (dataset, _):
for left, right, disparity, mask in dataset:
assert mask is None
assert disparity is None
datasets_utils.shape_test_for_stereo(left, right)
def test_bad_input(self):
with pytest.raises(ValueError, match="Unknown value 'bad' for argument split"):
with self.create_dataset(split="bad"):
pass
class CarlaStereoTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.CarlaStereo
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, None))
@staticmethod
def _create_scene_folders(num_examples: int, root_dir: Union[str, pathlib.Path]):
# make the root_dir if it does not exits
os.makedirs(root_dir, exist_ok=True)
for i in range(num_examples):
scene_dir = pathlib.Path(root_dir) / f"scene_{i}"
os.makedirs(scene_dir, exist_ok=True)
# populate with left right images
datasets_utils.create_image_file(root=scene_dir, name="im0.png", size=(100, 100))
datasets_utils.create_image_file(root=scene_dir, name="im1.png", size=(100, 100))
datasets_utils.make_fake_pfm_file(100, 100, file_name=str(scene_dir / "disp0GT.pfm"))
datasets_utils.make_fake_pfm_file(100, 100, file_name=str(scene_dir / "disp1GT.pfm"))
def inject_fake_data(self, tmpdir, config):
carla_dir = pathlib.Path(tmpdir) / "carla-highres"
os.makedirs(carla_dir, exist_ok=True)
split_dir = pathlib.Path(carla_dir) / "trainingF"
os.makedirs(split_dir, exist_ok=True)
num_examples = 6
self._create_scene_folders(num_examples=num_examples, root_dir=split_dir)
return num_examples
def test_train_splits(self):
with self.create_dataset() as (dataset, _):
for left, right, disparity in dataset:
datasets_utils.shape_test_for_stereo(left, right, disparity)
class CREStereoTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.CREStereo
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, np.ndarray, type(None))
def inject_fake_data(self, tmpdir, config):
crestereo_dir = pathlib.Path(tmpdir) / "CREStereo"
os.makedirs(crestereo_dir, exist_ok=True)
examples = {"tree": 2, "shapenet": 3, "reflective": 6, "hole": 5}
for category_name in ["shapenet", "reflective", "tree", "hole"]:
split_dir = crestereo_dir / category_name
os.makedirs(split_dir, exist_ok=True)
num_examples = examples[category_name]
for idx in range(num_examples):
datasets_utils.create_image_file(root=split_dir, name=f"{idx}_left.jpg", size=(100, 100))
datasets_utils.create_image_file(root=split_dir, name=f"{idx}_right.jpg", size=(100, 100))
# these are going to end up being gray scale images
datasets_utils.create_image_file(root=split_dir, name=f"{idx}_left.disp.png", size=(1, 100, 100))
datasets_utils.create_image_file(root=split_dir, name=f"{idx}_right.disp.png", size=(1, 100, 100))
return sum(examples.values())
def test_splits(self):
with self.create_dataset() as (dataset, _):
for left, right, disparity, mask in dataset:
assert mask is None
datasets_utils.shape_test_for_stereo(left, right, disparity)
class FallingThingsStereoTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.FallingThingsStereo
ADDITIONAL_CONFIGS = combinations_grid(variant=("single", "mixed", "both"))
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)))
@staticmethod
def _make_dummy_depth_map(root: str, name: str, size: Tuple[int, int]):
file = pathlib.Path(root) / name
image = np.ones((size[0], size[1]), dtype=np.uint8)
PIL.Image.fromarray(image).save(file)
@staticmethod
def _make_scene_folder(root: str, scene_name: str, size: Tuple[int, int]) -> None:
root = pathlib.Path(root) / scene_name
os.makedirs(root, exist_ok=True)
# jpg images
datasets_utils.create_image_file(root, "image1.left.jpg", size=(3, size[1], size[0]))
datasets_utils.create_image_file(root, "image1.right.jpg", size=(3, size[1], size[0]))
# single channel depth maps
FallingThingsStereoTestCase._make_dummy_depth_map(root, "image1.left.depth.png", size=(size[0], size[1]))
FallingThingsStereoTestCase._make_dummy_depth_map(root, "image1.right.depth.png", size=(size[0], size[1]))
# camera settings json. Minimal example for _read_disparity function testing
settings_json = {"camera_settings": [{"intrinsic_settings": {"fx": 1}}]}
with open(root / "_camera_settings.json", "w") as f:
json.dump(settings_json, f)
def inject_fake_data(self, tmpdir, config):
fallingthings_dir = pathlib.Path(tmpdir) / "FallingThings"
os.makedirs(fallingthings_dir, exist_ok=True)
num_examples = {"single": 2, "mixed": 3, "both": 4}.get(config["variant"], 0)
variants = {
"single": ["single"],
"mixed": ["mixed"],
"both": ["single", "mixed"],
}.get(config["variant"], [])
variant_dir_prefixes = {
"single": 1,
"mixed": 0,
}
for variant_name in variants:
variant_dir = pathlib.Path(fallingthings_dir) / variant_name
os.makedirs(variant_dir, exist_ok=True)
for i in range(variant_dir_prefixes[variant_name]):
variant_dir = variant_dir / f"{i:02d}"
os.makedirs(variant_dir, exist_ok=True)
for i in range(num_examples):
self._make_scene_folder(
root=variant_dir,
scene_name=f"scene_{i:06d}",
size=(100, 200),
)
if config["variant"] == "both":
num_examples *= 2
return num_examples
def test_splits(self):
for variant_name in ["single", "mixed"]:
with self.create_dataset(variant=variant_name) as (dataset, _):
for left, right, disparity in dataset:
datasets_utils.shape_test_for_stereo(left, right, disparity)
def test_bad_input(self):
with pytest.raises(ValueError, match="Unknown value 'bad' for argument variant"):
with self.create_dataset(variant="bad"):
pass
class SceneFlowStereoTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.SceneFlowStereo
ADDITIONAL_CONFIGS = combinations_grid(
variant=("FlyingThings3D", "Driving", "Monkaa"), pass_name=("clean", "final", "both")
)
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)))
@staticmethod
def _create_pfm_folder(
root: str, name: str, file_name_fn: Callable[..., str], num_examples: int, size: Tuple[int, int]
) -> None:
root = pathlib.Path(root) / name
os.makedirs(root, exist_ok=True)
for i in range(num_examples):
datasets_utils.make_fake_pfm_file(size[0], size[1], root / file_name_fn(i))
def inject_fake_data(self, tmpdir, config):
scene_flow_dir = pathlib.Path(tmpdir) / "SceneFlow"
os.makedirs(scene_flow_dir, exist_ok=True)
variant_dir = scene_flow_dir / config["variant"]
variant_dir_prefixes = {
"Monkaa": 0,
"Driving": 2,
"FlyingThings3D": 2,
}
os.makedirs(variant_dir, exist_ok=True)
num_examples = {"FlyingThings3D": 4, "Driving": 6, "Monkaa": 5}.get(config["variant"], 0)
passes = {
"clean": ["frames_cleanpass"],
"final": ["frames_finalpass"],
"both": ["frames_cleanpass", "frames_finalpass"],
}.get(config["pass_name"], [])
for pass_dir_name in passes:
# create pass directories
pass_dir = variant_dir / pass_dir_name
disp_dir = variant_dir / "disparity"
os.makedirs(pass_dir, exist_ok=True)
os.makedirs(disp_dir, exist_ok=True)
for i in range(variant_dir_prefixes.get(config["variant"], 0)):
pass_dir = pass_dir / str(i)
disp_dir = disp_dir / str(i)
os.makedirs(pass_dir, exist_ok=True)
os.makedirs(disp_dir, exist_ok=True)
for direction in ["left", "right"]:
for scene_idx in range(num_examples):
os.makedirs(pass_dir / f"scene_{scene_idx:06d}", exist_ok=True)
datasets_utils.create_image_folder(
root=pass_dir / f"scene_{scene_idx:06d}",
name=direction,
file_name_fn=lambda i: f"{i:06d}.png",
num_examples=1,
size=(3, 200, 100),
)
os.makedirs(disp_dir / f"scene_{scene_idx:06d}", exist_ok=True)
self._create_pfm_folder(
root=disp_dir / f"scene_{scene_idx:06d}",
name=direction,
file_name_fn=lambda i: f"{i:06d}.pfm",
num_examples=1,
size=(100, 200),
)
if config["pass_name"] == "both":
num_examples *= 2
return num_examples
def test_splits(self):
for variant_name, pass_name in itertools.product(["FlyingThings3D", "Driving", "Monkaa"], ["clean", "final"]):
with self.create_dataset(variant=variant_name, pass_name=pass_name) as (dataset, _):
for left, right, disparity in dataset:
datasets_utils.shape_test_for_stereo(left, right, disparity)
def test_bad_input(self):
with pytest.raises(ValueError, match="Unknown value 'bad' for argument variant"):
with self.create_dataset(variant="bad"):
pass
class InStereo2k(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.InStereo2k
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)))
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "test"))
@staticmethod
def _make_scene_folder(root: str, name: str, size: Tuple[int, int]):
root = pathlib.Path(root) / name
os.makedirs(root, exist_ok=True)
datasets_utils.create_image_file(root=root, name="left.png", size=(3, size[0], size[1]))
datasets_utils.create_image_file(root=root, name="right.png", size=(3, size[0], size[1]))
datasets_utils.create_image_file(root=root, name="left_disp.png", size=(1, size[0], size[1]))
datasets_utils.create_image_file(root=root, name="right_disp.png", size=(1, size[0], size[1]))
def inject_fake_data(self, tmpdir, config):
in_stereo_dir = pathlib.Path(tmpdir) / "InStereo2k"
os.makedirs(in_stereo_dir, exist_ok=True)
split_dir = pathlib.Path(in_stereo_dir) / config["split"]
os.makedirs(split_dir, exist_ok=True)
num_examples = {"train": 4, "test": 5}.get(config["split"], 0)
for i in range(num_examples):
self._make_scene_folder(split_dir, f"scene_{i:06d}", (100, 200))
return num_examples
def test_splits(self):
for split_name in ["train", "test"]:
with self.create_dataset(split=split_name) as (dataset, _):
for left, right, disparity in dataset:
datasets_utils.shape_test_for_stereo(left, right, disparity)
def test_bad_input(self):
with pytest.raises(
ValueError, match="Unknown value 'bad' for argument split. Valid values are {'train', 'test'}."
):
with self.create_dataset(split="bad"):
pass
class SintelStereoTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.SintelStereo
ADDITIONAL_CONFIGS = combinations_grid(pass_name=("final", "clean", "both"))
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)), (np.ndarray, type(None)))
def inject_fake_data(self, tmpdir, config):
sintel_dir = pathlib.Path(tmpdir) / "Sintel"
os.makedirs(sintel_dir, exist_ok=True)
split_dir = pathlib.Path(sintel_dir) / "training"
os.makedirs(split_dir, exist_ok=True)
# a single setting, since there are no splits
num_examples = {"final": 2, "clean": 3}
pass_names = {
"final": ["final"],
"clean": ["clean"],
"both": ["final", "clean"],
}.get(config["pass_name"], [])
for p in pass_names:
for view in [f"{p}_left", f"{p}_right"]:
root = split_dir / view
os.makedirs(root, exist_ok=True)
datasets_utils.create_image_folder(
root=root,
name="scene1",
file_name_fn=lambda i: f"{i:06d}.png",
num_examples=num_examples[p],
size=(3, 100, 200),
)
datasets_utils.create_image_folder(
root=split_dir / "occlusions",
name="scene1",
file_name_fn=lambda i: f"{i:06d}.png",
num_examples=max(num_examples.values()),
size=(1, 100, 200),
)
datasets_utils.create_image_folder(
root=split_dir / "outofframe",
name="scene1",
file_name_fn=lambda i: f"{i:06d}.png",
num_examples=max(num_examples.values()),
size=(1, 100, 200),
)
datasets_utils.create_image_folder(
root=split_dir / "disparities",
name="scene1",
file_name_fn=lambda i: f"{i:06d}.png",
num_examples=max(num_examples.values()),
size=(3, 100, 200),
)
if config["pass_name"] == "both":
num_examples = sum(num_examples.values())
else:
num_examples = num_examples.get(config["pass_name"], 0)
return num_examples
def test_splits(self):
for pass_name in ["final", "clean", "both"]:
with self.create_dataset(pass_name=pass_name) as (dataset, _):
for left, right, disparity, valid_mask in dataset:
datasets_utils.shape_test_for_stereo(left, right, disparity, valid_mask)
def test_bad_input(self):
with pytest.raises(ValueError, match="Unknown value 'bad' for argument pass_name"):
with self.create_dataset(pass_name="bad"):
pass
class ETH3DStereoestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.ETH3DStereo
ADDITIONAL_CONFIGS = combinations_grid(split=("train", "test"))
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)), (np.ndarray, type(None)))
@staticmethod
def _create_scene_folder(num_examples: int, root_dir: str):
# make the root_dir if it does not exits
root_dir = pathlib.Path(root_dir)
os.makedirs(root_dir, exist_ok=True)
for i in range(num_examples):
scene_dir = root_dir / f"scene_{i}"
os.makedirs(scene_dir, exist_ok=True)
# populate with left right images
datasets_utils.create_image_file(root=scene_dir, name="im0.png", size=(100, 100))
datasets_utils.create_image_file(root=scene_dir, name="im1.png", size=(100, 100))
@staticmethod
def _create_annotation_folder(num_examples: int, root_dir: str):
# make the root_dir if it does not exits
root_dir = pathlib.Path(root_dir)
os.makedirs(root_dir, exist_ok=True)
# create scene directories
for i in range(num_examples):
scene_dir = root_dir / f"scene_{i}"
os.makedirs(scene_dir, exist_ok=True)
# populate with a random png file for occlusion mask, and a pfm file for disparity
datasets_utils.create_image_file(root=scene_dir, name="mask0nocc.png", size=(1, 100, 100))
pfm_path = scene_dir / "disp0GT.pfm"
datasets_utils.make_fake_pfm_file(h=100, w=100, file_name=pfm_path)
def inject_fake_data(self, tmpdir, config):
eth3d_dir = pathlib.Path(tmpdir) / "ETH3D"
num_examples = 2 if config["split"] == "train" else 3
split_name = "two_view_training" if config["split"] == "train" else "two_view_test"
split_dir = eth3d_dir / split_name
self._create_scene_folder(num_examples, split_dir)
if config["split"] == "train":
annot_dir = eth3d_dir / "two_view_training_gt"
self._create_annotation_folder(num_examples, annot_dir)
return num_examples
def test_training_splits(self):
with self.create_dataset(split="train") as (dataset, _):
for left, right, disparity, valid_mask in dataset:
datasets_utils.shape_test_for_stereo(left, right, disparity, valid_mask)
def test_testing_splits(self):
with self.create_dataset(split="test") as (dataset, _):
assert all(d == (None, None) for d in dataset._disparities)
for left, right, disparity, valid_mask in dataset:
assert valid_mask is None
datasets_utils.shape_test_for_stereo(left, right, disparity)
def test_bad_input(self):
with pytest.raises(ValueError, match="Unknown value 'bad' for argument split"):
with self.create_dataset(split="bad"):
pass
class Middlebury2014StereoTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Middlebury2014Stereo
ADDITIONAL_CONFIGS = combinations_grid(
split=("train", "additional"),
calibration=("perfect", "imperfect", "both"),
use_ambient_views=(True, False),
)
FEATURE_TYPES = (PIL.Image.Image, PIL.Image.Image, (np.ndarray, type(None)), (np.ndarray, type(None)))
@staticmethod
def _make_scene_folder(root_dir: str, scene_name: str, split: str) -> None:
calibrations = [None] if split == "test" else ["-perfect", "-imperfect"]
root_dir = pathlib.Path(root_dir)
for c in calibrations:
scene_dir = root_dir / f"{scene_name}{c}"
os.makedirs(scene_dir, exist_ok=True)
# make normal images first
datasets_utils.create_image_file(root=scene_dir, name="im0.png", size=(3, 100, 100))
datasets_utils.create_image_file(root=scene_dir, name="im1.png", size=(3, 100, 100))
datasets_utils.create_image_file(root=scene_dir, name="im1E.png", size=(3, 100, 100))
datasets_utils.create_image_file(root=scene_dir, name="im1L.png", size=(3, 100, 100))
# these are going to end up being gray scale images
datasets_utils.make_fake_pfm_file(h=100, w=100, file_name=scene_dir / "disp0.pfm")
datasets_utils.make_fake_pfm_file(h=100, w=100, file_name=scene_dir / "disp1.pfm")
def inject_fake_data(self, tmpdir, config):
split_scene_map = {
"train": ["Adirondack", "Jadeplant", "Motorcycle", "Piano"],
"additional": ["Backpack", "Bicycle1", "Cable", "Classroom1"],
"test": ["Plants", "Classroom2E", "Classroom2", "Australia"],
}
middlebury_dir = pathlib.Path(tmpdir, "Middlebury2014")
os.makedirs(middlebury_dir, exist_ok=True)
split_dir = middlebury_dir / config["split"]
os.makedirs(split_dir, exist_ok=True)
num_examples = {"train": 2, "additional": 3, "test": 4}.get(config["split"], 0)
for idx in range(num_examples):
scene_name = split_scene_map[config["split"]][idx]
self._make_scene_folder(root_dir=split_dir, scene_name=scene_name, split=config["split"])
if config["calibration"] == "both":
num_examples *= 2
return num_examples
def test_train_splits(self):
for split, calibration in itertools.product(["train", "additional"], ["perfect", "imperfect", "both"]):
with self.create_dataset(split=split, calibration=calibration) as (dataset, _):
for left, right, disparity, mask in dataset:
datasets_utils.shape_test_for_stereo(left, right, disparity, mask)
def test_test_split(self):
for split in ["test"]:
with self.create_dataset(split=split, calibration=None) as (dataset, _):
for left, right, disparity, mask in dataset:
datasets_utils.shape_test_for_stereo(left, right)
def test_augmented_view_usage(self):
with self.create_dataset(split="train", use_ambient_views=True) as (dataset, _):
for left, right, disparity, mask in dataset:
datasets_utils.shape_test_for_stereo(left, right, disparity, mask)
def test_value_err_train(self):
# train set invalid
split = "train"
calibration = None
with pytest.raises(
ValueError,
match=f"Split '{split}' has calibration settings, however None was provided as an argument."
f"\nSetting calibration to 'perfect' for split '{split}'. Available calibration settings are: 'perfect', 'imperfect', 'both'.",
):
with self.create_dataset(split=split, calibration=calibration):
pass
def test_value_err_test(self):
# test set invalid
split = "test"
calibration = "perfect"
with pytest.raises(
ValueError, match="Split 'test' has only no calibration settings, please set `calibration=None`."
):
with self.create_dataset(split=split, calibration=calibration):
pass
def test_bad_input(self):
with pytest.raises(ValueError, match="Unknown value 'bad' for argument split"):
with self.create_dataset(split="bad"):
pass
class ImagenetteTestCase(datasets_utils.ImageDatasetTestCase):
DATASET_CLASS = datasets.Imagenette
ADDITIONAL_CONFIGS = combinations_grid(split=["train", "val"], size=["full", "320px", "160px"])
_WNIDS = [
"n01440764",
"n02102040",
"n02979186",
"n03000684",
"n03028079",
"n03394916",
"n03417042",
"n03425413",
"n03445777",
"n03888257",
]
def inject_fake_data(self, tmpdir, config):
archive_root = "imagenette2"
if config["size"] != "full":
archive_root += f"-{config['size'].replace('px', '')}"
image_root = pathlib.Path(tmpdir) / archive_root / config["split"]
num_images_per_class = 3
for wnid in self._WNIDS:
datasets_utils.create_image_folder(
root=image_root,
name=wnid,
file_name_fn=lambda idx: f"{wnid}_{idx}.JPEG",
num_examples=num_images_per_class,
)
return num_images_per_class * len(self._WNIDS)
class TestDatasetWrapper:
def test_unknown_type(self):
unknown_object = object()
with pytest.raises(
TypeError, match=re.escape("is meant for subclasses of `torchvision.datasets.VisionDataset`")
):
datasets.wrap_dataset_for_transforms_v2(unknown_object)
def test_unknown_dataset(self):
class MyVisionDataset(datasets.VisionDataset):
pass
dataset = MyVisionDataset("root")
with pytest.raises(TypeError, match="No wrapper exist"):
datasets.wrap_dataset_for_transforms_v2(dataset)
def test_missing_wrapper(self):
dataset = datasets.FakeData()
with pytest.raises(TypeError, match="please open an issue"):
datasets.wrap_dataset_for_transforms_v2(dataset)
def test_subclass(self, mocker):
from torchvision import tv_tensors
sentinel = object()
mocker.patch.dict(
tv_tensors._dataset_wrapper.WRAPPER_FACTORIES,
clear=False,
values={datasets.FakeData: lambda dataset, target_keys: lambda idx, sample: sentinel},
)
class MyFakeData(datasets.FakeData):
pass
dataset = MyFakeData()
wrapped_dataset = datasets.wrap_dataset_for_transforms_v2(dataset)
assert wrapped_dataset[0] is sentinel
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()
test/test_datasets_download.py
View file @ bf491463
import contextlib import contextlib
import itertools import itertools
import shutil
import tempfile
import time import time
import traceback
import unittest.mock import unittest.mock
import warnings
from datetime import datetime from datetime import datetime
from distutils import dir_util
from os import path from os import path
from urllib.error import HTTPError, URLError from urllib.error import HTTPError, URLError
from urllib.parse import urlparse from urllib.parse import urlparse
from urllib.request import urlopen, Request from urllib.request import Request, urlopen
import tempfile
import warnings
import pytest import pytest
from torchvision import datasets from torchvision import datasets
from torchvision.datasets.utils import ( from torchvision.datasets.utils import _get_redirect_url, USER_AGENT
download_url,
check_integrity,
download_file_from_google_drive,
_get_redirect_url,
USER_AGENT,
)
from common_utils import get_tmp_dir
def limit_requests_per_time(min_secs_between_requests=2.0): def limit_requests_per_time(min_secs_between_requests=2.0):
...@@ -86,63 +78,65 @@ urlopen = resolve_redirects()(urlopen) ...@@ -86,63 +78,65 @@ urlopen = resolve_redirects()(urlopen)
@contextlib.contextmanager @contextlib.contextmanager
def log_download_attempts( def log_download_attempts(
urls_and_md5s=None, urls,
file="utils", *,
patch=True, dataset_module,
mock_auxiliaries=None,
): ):
def add_mock(stack, name, file, **kwargs): def maybe_add_mock(*, module, name, stack, lst=None):
patcher = unittest.mock.patch(f"torchvision.datasets.{module}.{name}")
try: try:
return stack.enter_context(unittest.mock.patch(f"torchvision.datasets.{file}.{name}", **kwargs)) mock = stack.enter_context(patcher)
except AttributeError as error: except AttributeError:
if file != "utils": return
return add_mock(stack, name, "utils", **kwargs)
else:
raise pytest.UsageError from error
if urls_and_md5s is None:
urls_and_md5s = set()
if mock_auxiliaries is None:
mock_auxiliaries = patch
with contextlib.ExitStack() as stack: if lst is not None:
url_mock = add_mock(stack, "download_url", file, wraps=None if patch else download_url) lst.append(mock)
google_drive_mock = add_mock(
stack, "download_file_from_google_drive", file, wraps=None if patch else download_file_from_google_drive
)
if mock_auxiliaries: with contextlib.ExitStack() as stack:
add_mock(stack, "extract_archive", file) download_url_mocks = []
download_file_from_google_drive_mocks = []
for module in [dataset_module, "utils"]:
maybe_add_mock(module=module, name="download_url", stack=stack, lst=download_url_mocks)
maybe_add_mock(
module=module,
name="download_file_from_google_drive",
stack=stack,
lst=download_file_from_google_drive_mocks,
)
maybe_add_mock(module=module, name="extract_archive", stack=stack)
try: try:
yield urls_and_md5s yield
finally: finally:
for args, kwargs in url_mock.call_args_list: for download_url_mock in download_url_mocks:
url = args[0] for args, kwargs in download_url_mock.call_args_list:
md5 = args[-1] if len(args) == 4 else kwargs.get("md5") urls.append(args[0] if args else kwargs["url"])
urls_and_md5s.add((url, md5))
for args, kwargs in google_drive_mock.call_args_list: for download_file_from_google_drive_mock in download_file_from_google_drive_mocks:
id = args[0] for args, kwargs in download_file_from_google_drive_mock.call_args_list:
url = f"https://drive.google.com/file/d/{id}" file_id = args[0] if args else kwargs["file_id"]
md5 = args[3] if len(args) == 4 else kwargs.get("md5") urls.append(f"https://drive.google.com/file/d/{file_id}")
urls_and_md5s.add((url, md5))
def retry(fn, times=1, wait=5.0): def retry(fn, times=1, wait=5.0):
msgs = [] tbs = []
for _ in range(times + 1): for _ in range(times + 1):
try: try:
return fn() return fn()
except AssertionError as error: except AssertionError as error:
msgs.append(str(error)) tbs.append("".join(traceback.format_exception(type(error), error, error.__traceback__)))
time.sleep(wait) time.sleep(wait)
else: else:
raise AssertionError( raise AssertionError(
"\n".join( "\n".join(
( (
f"Assertion failed {times + 1} times with {wait:.1f} seconds intermediate wait time.\n", "\n",
*(f"{idx}: {error}" for idx, error in enumerate(msgs, 1)), *[f"{'_' * 40} {idx:2d} {'_' * 40}\n\n{tb}" for idx, tb in enumerate(tbs, 1)],
(
f"Assertion failed {times + 1} times with {wait:.1f} seconds intermediate wait time. "
f"You can find the the full tracebacks above."
),
) )
) )
) )
...@@ -152,10 +146,12 @@ def retry(fn, times=1, wait=5.0): ...@@ -152,10 +146,12 @@ def retry(fn, times=1, wait=5.0):
def assert_server_response_ok(): def assert_server_response_ok():
try: try:
yield yield
except URLError as error:
raise AssertionError("The request timed out.") from error
except HTTPError as error: except HTTPError as error:
raise AssertionError(f"The server returned {error.code}: {error.reason}.") from error raise AssertionError(f"The server returned {error.code}: {error.reason}.") from error
except URLError as error:
raise AssertionError(
"Connection not possible due to SSL." if "SSL" in str(error) else "The request timed out."
) from error
except RecursionError as error: except RecursionError as error:
raise AssertionError(str(error)) from error raise AssertionError(str(error)) from error
...@@ -166,46 +162,14 @@ def assert_url_is_accessible(url, timeout=5.0): ...@@ -166,46 +162,14 @@ def assert_url_is_accessible(url, timeout=5.0):
urlopen(request, timeout=timeout) urlopen(request, timeout=timeout)
def assert_file_downloads_correctly(url, md5, timeout=5.0): def collect_urls(dataset_cls, *args, **kwargs):
with get_tmp_dir() as root: urls = []
file = path.join(root, path.basename(url)) with contextlib.suppress(Exception), log_download_attempts(
with assert_server_response_ok(): urls, dataset_module=dataset_cls.__module__.split(".")[-1]
with open(file, "wb") as fh: ):
request = Request(url, headers={"User-Agent": USER_AGENT}) dataset_cls(*args, **kwargs)
response = urlopen(request, timeout=timeout)
fh.write(response.read())
assert check_integrity(file, md5=md5), "The MD5 checksums mismatch"
class DownloadConfig:
def __init__(self, url, md5=None, id=None):
self.url = url
self.md5 = md5
self.id = id or url
def __repr__(self):
return self.id
def make_download_configs(urls_and_md5s, name=None):
return [
DownloadConfig(url, md5=md5, id=f"{name}, {url}" if name is not None else None) for url, md5 in urls_and_md5s
]
def collect_download_configs(dataset_loader, name=None, **kwargs):
urls_and_md5s = set()
try:
with log_download_attempts(urls_and_md5s=urls_and_md5s, **kwargs):
dataset = dataset_loader()
except Exception:
dataset = None
if name is None and dataset is not None:
name = type(dataset).__name__
return make_download_configs(urls_and_md5s, name) return [(url, f"{dataset_cls.__name__}, {url}") for url in urls]
# This is a workaround since fixtures, such as the built-in tmp_dir, can only be used within a test but not within a # This is a workaround since fixtures, such as the built-in tmp_dir, can only be used within a test but not within a
...@@ -216,16 +180,18 @@ ROOT = tempfile.mkdtemp() ...@@ -216,16 +180,18 @@ ROOT = tempfile.mkdtemp()
@pytest.fixture(scope="module", autouse=True) @pytest.fixture(scope="module", autouse=True)
def root(): def root():
yield ROOT yield ROOT
dir_util.remove_tree(ROOT) shutil.rmtree(ROOT)
def places365(): def places365():
return itertools.chain( return itertools.chain.from_iterable(
*[ [
collect_download_configs( collect_urls(
lambda: datasets.Places365(ROOT, split=split, small=small, download=True), datasets.Places365,
name=f"Places365, {split}, {'small' if small else 'large'}", ROOT,
file="places365", split=split,
small=small,
download=True,
) )
for split, small in itertools.product(("train-standard", "train-challenge", "val"), (False, True)) for split, small in itertools.product(("train-standard", "train-challenge", "val"), (False, True))
] ]
...@@ -233,85 +199,69 @@ def places365(): ...@@ -233,85 +199,69 @@ def places365():
def caltech101(): def caltech101():
return collect_download_configs(lambda: datasets.Caltech101(ROOT, download=True), name="Caltech101") return collect_urls(datasets.Caltech101, ROOT, download=True)
def caltech256(): def caltech256():
return collect_download_configs(lambda: datasets.Caltech256(ROOT, download=True), name="Caltech256") return collect_urls(datasets.Caltech256, ROOT, download=True)
def cifar10(): def cifar10():
return collect_download_configs(lambda: datasets.CIFAR10(ROOT, download=True), name="CIFAR10") return collect_urls(datasets.CIFAR10, ROOT, download=True)
def cifar100(): def cifar100():
return collect_download_configs(lambda: datasets.CIFAR100(ROOT, download=True), name="CIFAR100") return collect_urls(datasets.CIFAR100, ROOT, download=True)
def voc(): def voc():
return itertools.chain( # TODO: Also test the "2007-test" key
*[ return itertools.chain.from_iterable(
collect_download_configs( [
lambda: datasets.VOCSegmentation(ROOT, year=year, download=True), collect_urls(datasets.VOCSegmentation, ROOT, year=year, download=True)
name=f"VOC, {year}", for year in ("2007", "2008", "2009", "2010", "2011", "2012")
file="voc",
)
for year in ("2007", "2007-test", "2008", "2009", "2010", "2011", "2012")
] ]
) )
def mnist(): def mnist():
with unittest.mock.patch.object(datasets.MNIST, "mirrors", datasets.MNIST.mirrors[-1:]): with unittest.mock.patch.object(datasets.MNIST, "mirrors", datasets.MNIST.mirrors[-1:]):
return collect_download_configs(lambda: datasets.MNIST(ROOT, download=True), name="MNIST") return collect_urls(datasets.MNIST, ROOT, download=True)
def fashion_mnist(): def fashion_mnist():
return collect_download_configs(lambda: datasets.FashionMNIST(ROOT, download=True), name="FashionMNIST") return collect_urls(datasets.FashionMNIST, ROOT, download=True)
def kmnist(): def kmnist():
return collect_download_configs(lambda: datasets.KMNIST(ROOT, download=True), name="KMNIST") return collect_urls(datasets.KMNIST, ROOT, download=True)
def emnist(): def emnist():
# the 'split' argument can be any valid one, since everything is downloaded anyway # the 'split' argument can be any valid one, since everything is downloaded anyway
return collect_download_configs(lambda: datasets.EMNIST(ROOT, split="byclass", download=True), name="EMNIST") return collect_urls(datasets.EMNIST, ROOT, split="byclass", download=True)
def qmnist(): def qmnist():
return itertools.chain( return itertools.chain.from_iterable(
*[ [collect_urls(datasets.QMNIST, ROOT, what=what, download=True) for what in ("train", "test", "nist")]
collect_download_configs(
lambda: datasets.QMNIST(ROOT, what=what, download=True),
name=f"QMNIST, {what}",
file="mnist",
)
for what in ("train", "test", "nist")
]
) )
def moving_mnist():
return collect_urls(datasets.MovingMNIST, ROOT, download=True)
def omniglot(): def omniglot():
return itertools.chain( return itertools.chain.from_iterable(
*[ [collect_urls(datasets.Omniglot, ROOT, background=background, download=True) for background in (True, False)]
collect_download_configs(
lambda: datasets.Omniglot(ROOT, background=background, download=True),
name=f"Omniglot, {'background' if background else 'evaluation'}",
)
for background in (True, False)
]
) )
def phototour(): def phototour():
return itertools.chain( return itertools.chain.from_iterable(
*[ [
collect_download_configs( collect_urls(datasets.PhotoTour, ROOT, name=name, download=True)
lambda: datasets.PhotoTour(ROOT, name=name, download=True),
name=f"PhotoTour, {name}",
file="phototour",
)
# The names postfixed with '_harris' point to the domain 'matthewalunbrown.com'. For some reason all # The names postfixed with '_harris' point to the domain 'matthewalunbrown.com'. For some reason all
# requests timeout from within CI. They are disabled until this is resolved. # requests timeout from within CI. They are disabled until this is resolved.
for name in ("notredame", "yosemite", "liberty") # "notredame_harris", "yosemite_harris", "liberty_harris" for name in ("notredame", "yosemite", "liberty") # "notredame_harris", "yosemite_harris", "liberty_harris"
...@@ -320,134 +270,119 @@ def phototour(): ...@@ -320,134 +270,119 @@ def phototour():
def sbdataset(): def sbdataset():
return collect_download_configs( return collect_urls(datasets.SBDataset, ROOT, download=True)
lambda: datasets.SBDataset(ROOT, download=True),
name="SBDataset",
file="voc",
)
def sbu(): def sbu():
return collect_download_configs( return collect_urls(datasets.SBU, ROOT, download=True)
lambda: datasets.SBU(ROOT, download=True),
name="SBU",
file="sbu",
)
def semeion(): def semeion():
return collect_download_configs( return collect_urls(datasets.SEMEION, ROOT, download=True)
lambda: datasets.SEMEION(ROOT, download=True),
name="SEMEION",
file="semeion",
)
def stl10(): def stl10():
return collect_download_configs( return collect_urls(datasets.STL10, ROOT, download=True)
lambda: datasets.STL10(ROOT, download=True),
name="STL10",
)
def svhn(): def svhn():
return itertools.chain( return itertools.chain.from_iterable(
*[ [collect_urls(datasets.SVHN, ROOT, split=split, download=True) for split in ("train", "test", "extra")]
collect_download_configs(
lambda: datasets.SVHN(ROOT, split=split, download=True),
name=f"SVHN, {split}",
file="svhn",
)
for split in ("train", "test", "extra")
]
) )
def usps(): def usps():
return itertools.chain( return itertools.chain.from_iterable(
*[ [collect_urls(datasets.USPS, ROOT, train=train, download=True) for train in (True, False)]
collect_download_configs(
lambda: datasets.USPS(ROOT, train=train, download=True),
name=f"USPS, {'train' if train else 'test'}",
file="usps",
)
for train in (True, False)
]
) )
def celeba(): def celeba():
return collect_download_configs( return collect_urls(datasets.CelebA, ROOT, download=True)
lambda: datasets.CelebA(ROOT, download=True),
name="CelebA",
file="celeba",
)
def widerface(): def widerface():
return collect_download_configs( return collect_urls(datasets.WIDERFace, ROOT, download=True)
lambda: datasets.WIDERFace(ROOT, download=True),
name="WIDERFace",
file="widerface", def kinetics():
return itertools.chain.from_iterable(
[
collect_urls(
datasets.Kinetics,
path.join(ROOT, f"Kinetics{num_classes}"),
frames_per_clip=1,
num_classes=num_classes,
split=split,
download=True,
)
for num_classes, split in itertools.product(("400", "600", "700"), ("train", "val"))
]
) )
def kitti(): def kitti():
return itertools.chain( return itertools.chain.from_iterable(
*[ [collect_urls(datasets.Kitti, ROOT, train=train, download=True) for train in (True, False)]
collect_download_configs(
lambda train=train: datasets.Kitti(ROOT, train=train, download=True),
name=f"Kitti, {'train' if train else 'test'}",
file="kitti",
)
for train in (True, False)
]
) )
def make_parametrize_kwargs(download_configs): def url_parametrization(*dataset_urls_and_ids_fns):
argvalues = [] return pytest.mark.parametrize(
ids = [] "url",
for config in download_configs: [
argvalues.append((config.url, config.md5)) pytest.param(url, id=id)
ids.append(config.id) for dataset_urls_and_ids_fn in dataset_urls_and_ids_fns
for url, id in sorted(set(dataset_urls_and_ids_fn()))
return dict(argnames=("url", "md5"), argvalues=argvalues, ids=ids) ],
@pytest.mark.parametrize(
**make_parametrize_kwargs(
itertools.chain(
places365(),
caltech101(),
caltech256(),
cifar10(),
cifar100(),
# The VOC download server is unstable. See https://github.com/pytorch/vision/issues/2953 for details.
# voc(),
mnist(),
fashion_mnist(),
kmnist(),
emnist(),
qmnist(),
omniglot(),
phototour(),
sbdataset(),
sbu(),
semeion(),
stl10(),
svhn(),
usps(),
celeba(),
widerface(),
kitti(),
)
) )
@url_parametrization(
caltech101,
caltech256,
cifar10,
cifar100,
# The VOC download server is unstable. See https://github.com/pytorch/vision/issues/2953 for details.
# voc,
mnist,
fashion_mnist,
kmnist,
emnist,
qmnist,
omniglot,
phototour,
sbdataset,
semeion,
stl10,
svhn,
usps,
celeba,
widerface,
kinetics,
kitti,
places365,
sbu,
) )
def test_url_is_accessible(url, md5): def test_url_is_accessible(url):
"""
If you see this test failing, find the offending dataset in the parametrization and move it to
``test_url_is_not_accessible`` and link an issue detailing the problem.
"""
retry(lambda: assert_url_is_accessible(url)) retry(lambda: assert_url_is_accessible(url))
@pytest.mark.parametrize(**make_parametrize_kwargs(itertools.chain())) # TODO: if e.g. caltech101 starts failing, remove the pytest.mark.parametrize below and use
def test_file_downloads_correctly(url, md5): # @url_parametrization(caltech101)
retry(lambda: assert_file_downloads_correctly(url, md5)) @pytest.mark.parametrize("url", ("http://url_that_doesnt_exist.com",)) # here until we actually have a failing dataset
@pytest.mark.xfail
def test_url_is_not_accessible(url):
"""
As the name implies, this test is the 'inverse' of ``test_url_is_accessible``. Since the download servers are
beyond our control, some files might not be accessible for longer stretches of time. Still, we want to know if they
come back up, or if we need to remove the download functionality of the dataset for good.
If you see this test failing, find the offending dataset in the parametrization and move it to
``test_url_is_accessible``.
"""
assert_url_is_accessible(url)
test/test_datasets_samplers.py
View file @ bf491463
import contextlib import pytest
import sys
import os
import torch import torch
import unittest from common_utils import assert_equal, get_list_of_videos
from torchvision import io from torchvision import io
from torchvision.datasets.samplers import ( from torchvision.datasets.samplers import DistributedSampler, RandomClipSampler, UniformClipSampler
DistributedSampler, from torchvision.datasets.video_utils import VideoClips
RandomClipSampler,
UniformClipSampler,
)
from torchvision.datasets.video_utils import VideoClips, unfold
from torchvision import get_video_backend
from common_utils import get_tmp_dir
from _assert_utils import assert_equal
@contextlib.contextmanager
def get_list_of_videos(num_videos=5, sizes=None, fps=None):
with get_tmp_dir() as tmp_dir:
names = []
for i in range(num_videos):
if sizes is None:
size = 5 * (i + 1)
else:
size = sizes[i]
if fps is None:
f = 5
else:
f = fps[i]
data = torch.randint(0, 256, (size, 300, 400, 3), dtype=torch.uint8)
name = os.path.join(tmp_dir, "{}.mp4".format(i))
names.append(name)
io.write_video(name, data, fps=f)
yield names
@unittest.skipIf(not io.video._av_available(), "this test requires av") @pytest.mark.skipif(not io.video._av_available(), reason="this test requires av")
class Tester(unittest.TestCase): class TestDatasetsSamplers:
def test_random_clip_sampler(self): def test_random_clip_sampler(self, tmpdir):
with get_list_of_videos(num_videos=3, sizes=[25, 25, 25]) as video_list: video_list = get_list_of_videos(tmpdir, num_videos=3, sizes=[25, 25, 25])
video_clips = VideoClips(video_list, 5, 5) video_clips = VideoClips(video_list, 5, 5)
sampler = RandomClipSampler(video_clips, 3) sampler = RandomClipSampler(video_clips, 3)
self.assertEqual(len(sampler), 3 * 3) assert len(sampler) == 3 * 3
indices = torch.tensor(list(iter(sampler))) indices = torch.tensor(list(iter(sampler)))
videos = torch.div(indices, 5, rounding_mode='floor') videos = torch.div(indices, 5, rounding_mode="floor")
v_idxs, count = torch.unique(videos, return_counts=True) v_idxs, count = torch.unique(videos, return_counts=True)
assert_equal(v_idxs, torch.tensor([0, 1, 2])) assert_equal(v_idxs, torch.tensor([0, 1, 2]))
assert_equal(count, torch.tensor([3, 3, 3])) assert_equal(count, torch.tensor([3, 3, 3]))
def test_random_clip_sampler_unequal(self): def test_random_clip_sampler_unequal(self, tmpdir):
with get_list_of_videos(num_videos=3, sizes=[10, 25, 25]) as video_list: video_list = get_list_of_videos(tmpdir, num_videos=3, sizes=[10, 25, 25])
video_clips = VideoClips(video_list, 5, 5) video_clips = VideoClips(video_list, 5, 5)
sampler = RandomClipSampler(video_clips, 3) sampler = RandomClipSampler(video_clips, 3)
self.assertEqual(len(sampler), 2 + 3 + 3) assert len(sampler) == 2 + 3 + 3
indices = list(iter(sampler)) indices = list(iter(sampler))
self.assertIn(0, indices) assert 0 in indices
self.assertIn(1, indices) assert 1 in indices
# remove elements of the first video, to simplify testing # remove elements of the first video, to simplify testing
indices.remove(0) indices.remove(0)
indices.remove(1) indices.remove(1)
indices = torch.tensor(indices) - 2 indices = torch.tensor(indices) - 2
videos = torch.div(indices, 5, rounding_mode='floor') videos = torch.div(indices, 5, rounding_mode="floor")
v_idxs, count = torch.unique(videos, return_counts=True) v_idxs, count = torch.unique(videos, return_counts=True)
assert_equal(v_idxs, torch.tensor([0, 1])) assert_equal(v_idxs, torch.tensor([0, 1]))
assert_equal(count, torch.tensor([3, 3])) assert_equal(count, torch.tensor([3, 3]))
def test_uniform_clip_sampler(self): def test_uniform_clip_sampler(self, tmpdir):
with get_list_of_videos(num_videos=3, sizes=[25, 25, 25]) as video_list: video_list = get_list_of_videos(tmpdir, num_videos=3, sizes=[25, 25, 25])
video_clips = VideoClips(video_list, 5, 5) video_clips = VideoClips(video_list, 5, 5)
sampler = UniformClipSampler(video_clips, 3) sampler = UniformClipSampler(video_clips, 3)
self.assertEqual(len(sampler), 3 * 3) assert len(sampler) == 3 * 3
indices = torch.tensor(list(iter(sampler))) indices = torch.tensor(list(iter(sampler)))
videos = torch.div(indices, 5, rounding_mode='floor') videos = torch.div(indices, 5, rounding_mode="floor")
v_idxs, count = torch.unique(videos, return_counts=True) v_idxs, count = torch.unique(videos, return_counts=True)
assert_equal(v_idxs, torch.tensor([0, 1, 2])) assert_equal(v_idxs, torch.tensor([0, 1, 2]))
assert_equal(count, torch.tensor([3, 3, 3])) assert_equal(count, torch.tensor([3, 3, 3]))
assert_equal(indices, torch.tensor([0, 2, 4, 5, 7, 9, 10, 12, 14])) assert_equal(indices, torch.tensor([0, 2, 4, 5, 7, 9, 10, 12, 14]))
def test_uniform_clip_sampler_insufficient_clips(self): def test_uniform_clip_sampler_insufficient_clips(self, tmpdir):
with get_list_of_videos(num_videos=3, sizes=[10, 25, 25]) as video_list: video_list = get_list_of_videos(tmpdir, num_videos=3, sizes=[10, 25, 25])
video_clips = VideoClips(video_list, 5, 5) video_clips = VideoClips(video_list, 5, 5)
sampler = UniformClipSampler(video_clips, 3) sampler = UniformClipSampler(video_clips, 3)
self.assertEqual(len(sampler), 3 * 3) assert len(sampler) == 3 * 3
indices = torch.tensor(list(iter(sampler))) indices = torch.tensor(list(iter(sampler)))
assert_equal(indices, torch.tensor([0, 0, 1, 2, 4, 6, 7, 9, 11])) assert_equal(indices, torch.tensor([0, 0, 1, 2, 4, 6, 7, 9, 11]))
def test_distributed_sampler_and_uniform_clip_sampler(self): def test_distributed_sampler_and_uniform_clip_sampler(self, tmpdir):
with get_list_of_videos(num_videos=3, sizes=[25, 25, 25]) as video_list: video_list = get_list_of_videos(tmpdir, num_videos=3, sizes=[25, 25, 25])
video_clips = VideoClips(video_list, 5, 5) video_clips = VideoClips(video_list, 5, 5)
clip_sampler = UniformClipSampler(video_clips, 3) clip_sampler = UniformClipSampler(video_clips, 3)
distributed_sampler_rank0 = DistributedSampler( distributed_sampler_rank0 = DistributedSampler(
clip_sampler, clip_sampler,
num_replicas=2, num_replicas=2,
rank=0, rank=0,
group_size=3, group_size=3,
) )
indices = torch.tensor(list(iter(distributed_sampler_rank0))) indices = torch.tensor(list(iter(distributed_sampler_rank0)))
self.assertEqual(len(distributed_sampler_rank0), 6) assert len(distributed_sampler_rank0) == 6
assert_equal(indices, torch.tensor([0, 2, 4, 10, 12, 14])) assert_equal(indices, torch.tensor([0, 2, 4, 10, 12, 14]))
distributed_sampler_rank1 = DistributedSampler( distributed_sampler_rank1 = DistributedSampler(
clip_sampler, clip_sampler,
num_replicas=2, num_replicas=2,
rank=1, rank=1,
group_size=3, group_size=3,
) )
indices = torch.tensor(list(iter(distributed_sampler_rank1))) indices = torch.tensor(list(iter(distributed_sampler_rank1)))
self.assertEqual(len(distributed_sampler_rank1), 6) assert len(distributed_sampler_rank1) == 6
assert_equal(indices, torch.tensor([5, 7, 9, 0, 2, 4])) assert_equal(indices, torch.tensor([5, 7, 9, 0, 2, 4]))
if __name__ == '__main__': if __name__ == "__main__":
unittest.main() pytest.main([__file__])
test/test_datasets_utils.py
View file @ bf491463
import contextlib
import gzip
import os import os
import torchvision.datasets.utils as utils import pathlib
import unittest import re
import unittest.mock
import zipfile
import tarfile import tarfile
import gzip import zipfile
import warnings
import pytest
import torch
import torchvision.datasets.utils as utils
from common_utils import assert_equal
from torch._utils_internal import get_file_path_2 from torch._utils_internal import get_file_path_2
from urllib.error import URLError from torchvision.datasets.folder import make_dataset
import itertools from torchvision.datasets.utils import _COMPRESSED_FILE_OPENERS
import lzma
from common_utils import get_tmp_dir, call_args_to_kwargs_only TEST_FILE = get_file_path_2(
os.path.dirname(os.path.abspath(__file__)), "assets", "encode_jpeg", "grace_hopper_517x606.jpg"
)
TEST_FILE = get_file_path_2( def patch_url_redirection(mocker, redirect_url):
os.path.dirname(os.path.abspath(__file__)), 'assets', 'encode_jpeg', 'grace_hopper_517x606.jpg') class Response:
def __init__(self, url):
self.url = url
@contextlib.contextmanager
def patched_opener(*args, **kwargs):
yield Response(redirect_url)
return mocker.patch("torchvision.datasets.utils.urllib.request.urlopen", side_effect=patched_opener)
class TestDatasetsUtils:
def test_get_redirect_url(self, mocker):
url = "https://url.org"
expected_redirect_url = "https://redirect.url.org"
mock = patch_url_redirection(mocker, expected_redirect_url)
actual = utils._get_redirect_url(url)
assert actual == expected_redirect_url
class Tester(unittest.TestCase): assert mock.call_count == 2
call_args_1, call_args_2 = mock.call_args_list
assert call_args_1[0][0].full_url == url
assert call_args_2[0][0].full_url == expected_redirect_url
def test_check_md5(self): def test_get_redirect_url_max_hops_exceeded(self, mocker):
url = "https://url.org"
redirect_url = "https://redirect.url.org"
mock = patch_url_redirection(mocker, redirect_url)
with pytest.raises(RecursionError):
utils._get_redirect_url(url, max_hops=0)
assert mock.call_count == 1
assert mock.call_args[0][0].full_url == url
@pytest.mark.parametrize("use_pathlib", (True, False))
def test_check_md5(self, use_pathlib):
fpath = TEST_FILE fpath = TEST_FILE
correct_md5 = '9c0bb82894bb3af7f7675ef2b3b6dcdc' if use_pathlib:
false_md5 = '' fpath = pathlib.Path(fpath)
self.assertTrue(utils.check_md5(fpath, correct_md5)) correct_md5 = "9c0bb82894bb3af7f7675ef2b3b6dcdc"
self.assertFalse(utils.check_md5(fpath, false_md5)) false_md5 = ""
assert utils.check_md5(fpath, correct_md5)
assert not utils.check_md5(fpath, false_md5)
def test_check_integrity(self): def test_check_integrity(self):
existing_fpath = TEST_FILE existing_fpath = TEST_FILE
nonexisting_fpath = '' nonexisting_fpath = ""
correct_md5 = '9c0bb82894bb3af7f7675ef2b3b6dcdc' correct_md5 = "9c0bb82894bb3af7f7675ef2b3b6dcdc"
false_md5 = '' false_md5 = ""
self.assertTrue(utils.check_integrity(existing_fpath, correct_md5)) assert utils.check_integrity(existing_fpath, correct_md5)
self.assertFalse(utils.check_integrity(existing_fpath, false_md5)) assert not utils.check_integrity(existing_fpath, false_md5)
self.assertTrue(utils.check_integrity(existing_fpath)) assert utils.check_integrity(existing_fpath)
self.assertFalse(utils.check_integrity(nonexisting_fpath)) assert not utils.check_integrity(nonexisting_fpath)
def test_get_google_drive_file_id(self): def test_get_google_drive_file_id(self):
url = "https://drive.google.com/file/d/1hbzc_P1FuxMkcabkgn9ZKinBwW683j45/view" url = "https://drive.google.com/file/d/1GO-BHUYRuvzr1Gtp2_fqXRsr9TIeYbhV/view"
expected = "1hbzc_P1FuxMkcabkgn9ZKinBwW683j45" expected = "1GO-BHUYRuvzr1Gtp2_fqXRsr9TIeYbhV"
actual = utils._get_google_drive_file_id(url) actual = utils._get_google_drive_file_id(url)
assert actual == expected assert actual == expected
...@@ -49,84 +90,64 @@ class Tester(unittest.TestCase): ...@@ -49,84 +90,64 @@ class Tester(unittest.TestCase):
assert utils._get_google_drive_file_id(url) is None assert utils._get_google_drive_file_id(url) is None
def test_detect_file_type(self): @pytest.mark.parametrize(
for file, expected in [ "file, expected",
[
("foo.tar.bz2", (".tar.bz2", ".tar", ".bz2")),
("foo.tar.xz", (".tar.xz", ".tar", ".xz")), ("foo.tar.xz", (".tar.xz", ".tar", ".xz")),
("foo.tar", (".tar", ".tar", None)), ("foo.tar", (".tar", ".tar", None)),
("foo.tar.gz", (".tar.gz", ".tar", ".gz")), ("foo.tar.gz", (".tar.gz", ".tar", ".gz")),
("foo.tbz", (".tbz", ".tar", ".bz2")),
("foo.tbz2", (".tbz2", ".tar", ".bz2")),
("foo.tgz", (".tgz", ".tar", ".gz")), ("foo.tgz", (".tgz", ".tar", ".gz")),
("foo.bz2", (".bz2", None, ".bz2")),
("foo.gz", (".gz", None, ".gz")), ("foo.gz", (".gz", None, ".gz")),
("foo.zip", (".zip", ".zip", None)), ("foo.zip", (".zip", ".zip", None)),
("foo.xz", (".xz", None, ".xz")), ("foo.xz", (".xz", None, ".xz")),
]: ("foo.bar.tar.gz", (".tar.gz", ".tar", ".gz")),
with self.subTest(file=file): ("foo.bar.gz", (".gz", None, ".gz")),
self.assertSequenceEqual(utils._detect_file_type(file), expected) ("foo.bar.zip", (".zip", ".zip", None)),
],
def test_detect_file_type_no_ext(self): )
with self.assertRaises(RuntimeError): def test_detect_file_type(self, file, expected):
utils._detect_file_type("foo") assert utils._detect_file_type(file) == expected
def test_detect_file_type_to_many_exts(self): @pytest.mark.parametrize("file", ["foo", "foo.tar.baz", "foo.bar"])
with self.assertRaises(RuntimeError): def test_detect_file_type_incompatible(self, file):
utils._detect_file_type("foo.bar.tar.gz") # tests detect file type for no extension, unknown compression and unknown partial extension
with pytest.raises(RuntimeError):
def test_detect_file_type_unknown_archive_type(self): utils._detect_file_type(file)
with self.assertRaises(RuntimeError):
utils._detect_file_type("foo.bar.gz") @pytest.mark.parametrize("extension", [".bz2", ".gz", ".xz"])
@pytest.mark.parametrize("use_pathlib", (True, False))
def test_detect_file_type_unknown_compression(self): def test_decompress(self, extension, tmpdir, use_pathlib):
with self.assertRaises(RuntimeError):
utils._detect_file_type("foo.tar.baz")
def test_detect_file_type_unknown_partial_ext(self):
with self.assertRaises(RuntimeError):
utils._detect_file_type("foo.bar")
def test_decompress_gzip(self):
def create_compressed(root, content="this is the content"):
file = os.path.join(root, "file")
compressed = f"{file}.gz"
with gzip.open(compressed, "wb") as fh:
fh.write(content.encode())
return compressed, file, content
with get_tmp_dir() as temp_dir:
compressed, file, content = create_compressed(temp_dir)
utils._decompress(compressed)
self.assertTrue(os.path.exists(file))
with open(file, "r") as fh:
self.assertEqual(fh.read(), content)
def test_decompress_lzma(self):
def create_compressed(root, content="this is the content"): def create_compressed(root, content="this is the content"):
file = os.path.join(root, "file") file = os.path.join(root, "file")
compressed = f"{file}.xz" compressed = f"{file}{extension}"
compressed_file_opener = _COMPRESSED_FILE_OPENERS[extension]
with lzma.open(compressed, "wb") as fh: with compressed_file_opener(compressed, "wb") as fh:
fh.write(content.encode()) fh.write(content.encode())
return compressed, file, content return compressed, file, content
with get_tmp_dir() as temp_dir: compressed, file, content = create_compressed(tmpdir)
compressed, file, content = create_compressed(temp_dir) if use_pathlib:
compressed = pathlib.Path(compressed)
utils.extract_archive(compressed, temp_dir) utils._decompress(compressed)
self.assertTrue(os.path.exists(file)) assert os.path.exists(file)
with open(file, "r") as fh: with open(file) as fh:
self.assertEqual(fh.read(), content) assert fh.read() == content
def test_decompress_no_compression(self): def test_decompress_no_compression(self):
with self.assertRaises(RuntimeError): with pytest.raises(RuntimeError):
utils._decompress("foo.tar") utils._decompress("foo.tar")
def test_decompress_remove_finished(self): @pytest.mark.parametrize("use_pathlib", (True, False))
def test_decompress_remove_finished(self, tmpdir, use_pathlib):
def create_compressed(root, content="this is the content"): def create_compressed(root, content="this is the content"):
file = os.path.join(root, "file") file = os.path.join(root, "file")
compressed = f"{file}.gz" compressed = f"{file}.gz"
...@@ -136,28 +157,35 @@ class Tester(unittest.TestCase): ...@@ -136,28 +157,35 @@ class Tester(unittest.TestCase):
return compressed, file, content return compressed, file, content
with get_tmp_dir() as temp_dir: compressed, file, content = create_compressed(tmpdir)
compressed, file, content = create_compressed(temp_dir) print(f"{type(compressed)=}")
if use_pathlib:
compressed = pathlib.Path(compressed)
tmpdir = pathlib.Path(tmpdir)
extracted_dir = utils.extract_archive(compressed, tmpdir, remove_finished=True)
assert not os.path.exists(compressed)
if use_pathlib:
assert isinstance(extracted_dir, pathlib.Path)
assert isinstance(compressed, pathlib.Path)
else:
assert isinstance(extracted_dir, str)
assert isinstance(compressed, str)
@pytest.mark.parametrize("extension", [".gz", ".xz"])
@pytest.mark.parametrize("remove_finished", [True, False])
def test_extract_archive_defer_to_decompress(self, extension, remove_finished, mocker):
filename = "foo"
file = f"{filename}{extension}"
utils.extract_archive(compressed, temp_dir, remove_finished=True) mocked = mocker.patch("torchvision.datasets.utils._decompress")
utils.extract_archive(file, remove_finished=remove_finished)
self.assertFalse(os.path.exists(compressed)) mocked.assert_called_once_with(file, filename, remove_finished=remove_finished)
def test_extract_archive_defer_to_decompress(self): @pytest.mark.parametrize("use_pathlib", (True, False))
filename = "foo" def test_extract_zip(self, tmpdir, use_pathlib):
for ext, remove_finished in itertools.product((".gz", ".xz"), (True, False)):
with self.subTest(ext=ext, remove_finished=remove_finished):
with unittest.mock.patch("torchvision.datasets.utils._decompress") as mock:
file = f"{filename}{ext}"
utils.extract_archive(file, remove_finished=remove_finished)
mock.assert_called_once()
self.assertEqual(
call_args_to_kwargs_only(mock.call_args, utils._decompress),
dict(from_path=file, to_path=filename, remove_finished=remove_finished),
)
def test_extract_zip(self):
def create_archive(root, content="this is the content"): def create_archive(root, content="this is the content"):
file = os.path.join(root, "dst.txt") file = os.path.join(root, "dst.txt")
archive = os.path.join(root, "archive.zip") archive = os.path.join(root, "archive.zip")
...@@ -167,46 +195,26 @@ class Tester(unittest.TestCase): ...@@ -167,46 +195,26 @@ class Tester(unittest.TestCase):
return archive, file, content return archive, file, content
with get_tmp_dir() as temp_dir: if use_pathlib:
archive, file, content = create_archive(temp_dir) tmpdir = pathlib.Path(tmpdir)
archive, file, content = create_archive(tmpdir)
utils.extract_archive(archive, temp_dir)
self.assertTrue(os.path.exists(file))
with open(file, "r") as fh:
self.assertEqual(fh.read(), content)
def test_extract_tar(self): utils.extract_archive(archive, tmpdir)
def create_archive(root, ext, mode, content="this is the content"):
src = os.path.join(root, "src.txt")
dst = os.path.join(root, "dst.txt")
archive = os.path.join(root, f"archive{ext}")
with open(src, "w") as fh:
fh.write(content)
with tarfile.open(archive, mode=mode) as fh:
fh.add(src, arcname=os.path.basename(dst))
return archive, dst, content
for ext, mode in zip(['.tar', '.tar.gz', '.tgz'], ['w', 'w:gz', 'w:gz']):
with get_tmp_dir() as temp_dir:
archive, file, content = create_archive(temp_dir, ext, mode)
utils.extract_archive(archive, temp_dir) assert os.path.exists(file)
self.assertTrue(os.path.exists(file)) with open(file) as fh:
assert fh.read() == content
with open(file, "r") as fh: @pytest.mark.parametrize(
self.assertEqual(fh.read(), content) "extension, mode", [(".tar", "w"), (".tar.gz", "w:gz"), (".tgz", "w:gz"), (".tar.xz", "w:xz")]
)
def test_extract_tar_xz(self): @pytest.mark.parametrize("use_pathlib", (True, False))
def create_archive(root, ext, mode, content="this is the content"): def test_extract_tar(self, extension, mode, tmpdir, use_pathlib):
def create_archive(root, extension, mode, content="this is the content"):
src = os.path.join(root, "src.txt") src = os.path.join(root, "src.txt")
dst = os.path.join(root, "dst.txt") dst = os.path.join(root, "dst.txt")
archive = os.path.join(root, f"archive{ext}") archive = os.path.join(root, f"archive{extension}")
with open(src, "w") as fh: with open(src, "w") as fh:
fh.write(content) fh.write(content)
...@@ -216,22 +224,64 @@ class Tester(unittest.TestCase): ...@@ -216,22 +224,64 @@ class Tester(unittest.TestCase):
return archive, dst, content return archive, dst, content
for ext, mode in zip(['.tar.xz'], ['w:xz']): if use_pathlib:
with get_tmp_dir() as temp_dir: tmpdir = pathlib.Path(tmpdir)
archive, file, content = create_archive(temp_dir, ext, mode) archive, file, content = create_archive(tmpdir, extension, mode)
utils.extract_archive(archive, temp_dir) utils.extract_archive(archive, tmpdir)
self.assertTrue(os.path.exists(file)) assert os.path.exists(file)
with open(file, "r") as fh: with open(file) as fh:
self.assertEqual(fh.read(), content) assert fh.read() == content
def test_verify_str_arg(self): def test_verify_str_arg(self):
self.assertEqual("a", utils.verify_str_arg("a", "arg", ("a",))) assert "a" == utils.verify_str_arg("a", "arg", ("a",))
self.assertRaises(ValueError, utils.verify_str_arg, 0, ("a",), "arg") pytest.raises(ValueError, utils.verify_str_arg, 0, ("a",), "arg")
self.assertRaises(ValueError, utils.verify_str_arg, "b", ("a",), "arg") pytest.raises(ValueError, utils.verify_str_arg, "b", ("a",), "arg")
@pytest.mark.parametrize(
if __name__ == '__main__': ("dtype", "actual_hex", "expected_hex"),
unittest.main() [
(torch.uint8, "01 23 45 67 89 AB CD EF", "01 23 45 67 89 AB CD EF"),
(torch.float16, "01 23 45 67 89 AB CD EF", "23 01 67 45 AB 89 EF CD"),
(torch.int32, "01 23 45 67 89 AB CD EF", "67 45 23 01 EF CD AB 89"),
(torch.float64, "01 23 45 67 89 AB CD EF", "EF CD AB 89 67 45 23 01"),
],
)
def test_flip_byte_order(self, dtype, actual_hex, expected_hex):
def to_tensor(hex):
return torch.frombuffer(bytes.fromhex(hex), dtype=dtype)
assert_equal(
utils._flip_byte_order(to_tensor(actual_hex)),
to_tensor(expected_hex),
)
@pytest.mark.parametrize(
("kwargs", "expected_error_msg"),
[
(dict(is_valid_file=lambda path: pathlib.Path(path).suffix in {".png", ".jpeg"}), "classes c"),
(dict(extensions=".png"), re.escape("classes b, c. Supported extensions are: .png")),
(dict(extensions=(".png", ".jpeg")), re.escape("classes c. Supported extensions are: .png, .jpeg")),
],
)
def test_make_dataset_no_valid_files(tmpdir, kwargs, expected_error_msg):
tmpdir = pathlib.Path(tmpdir)
(tmpdir / "a").mkdir()
(tmpdir / "a" / "a.png").touch()
(tmpdir / "b").mkdir()
(tmpdir / "b" / "b.jpeg").touch()
(tmpdir / "c").mkdir()
(tmpdir / "c" / "c.unknown").touch()
with pytest.raises(FileNotFoundError, match=expected_error_msg):
make_dataset(str(tmpdir), **kwargs)
if __name__ == "__main__":
pytest.main([__file__])
test/test_datasets_video_utils.py
View file @ bf491463
import contextlib import pytest
import os
import torch import torch
import unittest from common_utils import assert_equal, get_list_of_videos
from torchvision import io from torchvision import io
from torchvision.datasets.video_utils import VideoClips, unfold from torchvision.datasets.video_utils import unfold, VideoClips
from common_utils import get_tmp_dir
from _assert_utils import assert_equal
@contextlib.contextmanager
def get_list_of_videos(num_videos=5, sizes=None, fps=None):
with get_tmp_dir() as tmp_dir:
names = []
for i in range(num_videos):
if sizes is None:
size = 5 * (i + 1)
else:
size = sizes[i]
if fps is None:
f = 5
else:
f = fps[i]
data = torch.randint(0, 256, (size, 300, 400, 3), dtype=torch.uint8)
name = os.path.join(tmp_dir, "{}.mp4".format(i))
names.append(name)
io.write_video(name, data, fps=f)
yield names
class Tester(unittest.TestCase):
class TestVideo:
def test_unfold(self): def test_unfold(self):
a = torch.arange(7) a = torch.arange(7)
r = unfold(a, 3, 3, 1) r = unfold(a, 3, 3, 1)
expected = torch.tensor([ expected = torch.tensor(
[0, 1, 2], [
[3, 4, 5], [0, 1, 2],
]) [3, 4, 5],
assert_equal(r, expected, check_stride=False) ]
)
assert_equal(r, expected)
r = unfold(a, 3, 2, 1) r = unfold(a, 3, 2, 1)
expected = torch.tensor([ expected = torch.tensor([[0, 1, 2], [2, 3, 4], [4, 5, 6]])
[0, 1, 2], assert_equal(r, expected)
[2, 3, 4],
[4, 5, 6]
])
assert_equal(r, expected, check_stride=False)
r = unfold(a, 3, 2, 2) r = unfold(a, 3, 2, 2)
expected = torch.tensor([ expected = torch.tensor(
[0, 2, 4], [
[2, 4, 6], [0, 2, 4],
]) [2, 4, 6],
assert_equal(r, expected, check_stride=False) ]
)
@unittest.skipIf(not io.video._av_available(), "this test requires av") assert_equal(r, expected)
def test_video_clips(self):
with get_list_of_videos(num_videos=3) as video_list: @pytest.mark.skipif(not io.video._av_available(), reason="this test requires av")
video_clips = VideoClips(video_list, 5, 5, num_workers=2) def test_video_clips(self, tmpdir):
assert video_clips.num_clips() == 1 + 2 + 3 video_list = get_list_of_videos(tmpdir, num_videos=3)
for i, (v_idx, c_idx) in enumerate([(0, 0), (1, 0), (1, 1), (2, 0), (2, 1), (2, 2)]): video_clips = VideoClips(video_list, 5, 5, num_workers=2)
video_idx, clip_idx = video_clips.get_clip_location(i) assert video_clips.num_clips() == 1 + 2 + 3
assert video_idx == v_idx for i, (v_idx, c_idx) in enumerate([(0, 0), (1, 0), (1, 1), (2, 0), (2, 1), (2, 2)]):
assert clip_idx == c_idx video_idx, clip_idx = video_clips.get_clip_location(i)
assert video_idx == v_idx
video_clips = VideoClips(video_list, 6, 6) assert clip_idx == c_idx
assert video_clips.num_clips() == 0 + 1 + 2
for i, (v_idx, c_idx) in enumerate([(1, 0), (2, 0), (2, 1)]): video_clips = VideoClips(video_list, 6, 6)
video_idx, clip_idx = video_clips.get_clip_location(i) assert video_clips.num_clips() == 0 + 1 + 2
assert video_idx == v_idx for i, (v_idx, c_idx) in enumerate([(1, 0), (2, 0), (2, 1)]):
assert clip_idx == c_idx video_idx, clip_idx = video_clips.get_clip_location(i)
assert video_idx == v_idx
video_clips = VideoClips(video_list, 6, 1) assert clip_idx == c_idx
assert video_clips.num_clips() == 0 + (10 - 6 + 1) + (15 - 6 + 1)
for i, v_idx, c_idx in [(0, 1, 0), (4, 1, 4), (5, 2, 0), (6, 2, 1)]: video_clips = VideoClips(video_list, 6, 1)
video_idx, clip_idx = video_clips.get_clip_location(i) assert video_clips.num_clips() == 0 + (10 - 6 + 1) + (15 - 6 + 1)
assert video_idx == v_idx for i, v_idx, c_idx in [(0, 1, 0), (4, 1, 4), (5, 2, 0), (6, 2, 1)]:
assert clip_idx == c_idx video_idx, clip_idx = video_clips.get_clip_location(i)
assert video_idx == v_idx
@unittest.skipIf(not io.video._av_available(), "this test requires av") assert clip_idx == c_idx
def test_video_clips_custom_fps(self):
with get_list_of_videos(num_videos=3, sizes=[12, 12, 12], fps=[3, 4, 6]) as video_list: @pytest.mark.skipif(not io.video._av_available(), reason="this test requires av")
num_frames = 4 def test_video_clips_custom_fps(self, tmpdir):
for fps in [1, 3, 4, 10]: video_list = get_list_of_videos(tmpdir, num_videos=3, sizes=[12, 12, 12], fps=[3, 4, 6])
video_clips = VideoClips(video_list, num_frames, num_frames, fps, num_workers=2) num_frames = 4
for i in range(video_clips.num_clips()): for fps in [1, 3, 4, 10]:
video, audio, info, video_idx = video_clips.get_clip(i) video_clips = VideoClips(video_list, num_frames, num_frames, fps)
assert video.shape[0] == num_frames for i in range(video_clips.num_clips()):
assert info["video_fps"] == fps video, audio, info, video_idx = video_clips.get_clip(i)
# TODO add tests checking that the content is right assert video.shape[0] == num_frames
assert info["video_fps"] == fps
# TODO add tests checking that the content is right
def test_compute_clips_for_video(self): def test_compute_clips_for_video(self):
video_pts = torch.arange(30) video_pts = torch.arange(30)
...@@ -101,8 +74,7 @@ class Tester(unittest.TestCase): ...@@ -101,8 +74,7 @@ class Tester(unittest.TestCase):
orig_fps = 30 orig_fps = 30
duration = float(len(video_pts)) / orig_fps duration = float(len(video_pts)) / orig_fps
new_fps = 13 new_fps = 13
clips, idxs = VideoClips.compute_clips_for_video(video_pts, num_frames, num_frames, clips, idxs = VideoClips.compute_clips_for_video(video_pts, num_frames, num_frames, orig_fps, new_fps)
orig_fps, new_fps)
resampled_idxs = VideoClips._resample_video_idx(int(duration * new_fps), orig_fps, new_fps) resampled_idxs = VideoClips._resample_video_idx(int(duration * new_fps), orig_fps, new_fps)
assert len(clips) == 1 assert len(clips) == 1
assert_equal(clips, idxs) assert_equal(clips, idxs)
...@@ -113,8 +85,7 @@ class Tester(unittest.TestCase): ...@@ -113,8 +85,7 @@ class Tester(unittest.TestCase):
orig_fps = 30 orig_fps = 30
duration = float(len(video_pts)) / orig_fps duration = float(len(video_pts)) / orig_fps
new_fps = 12 new_fps = 12
clips, idxs = VideoClips.compute_clips_for_video(video_pts, num_frames, num_frames, clips, idxs = VideoClips.compute_clips_for_video(video_pts, num_frames, num_frames, orig_fps, new_fps)
orig_fps, new_fps)
resampled_idxs = VideoClips._resample_video_idx(int(duration * new_fps), orig_fps, new_fps) resampled_idxs = VideoClips._resample_video_idx(int(duration * new_fps), orig_fps, new_fps)
assert len(clips) == 3 assert len(clips) == 3
assert_equal(clips, idxs) assert_equal(clips, idxs)
...@@ -124,12 +95,11 @@ class Tester(unittest.TestCase): ...@@ -124,12 +95,11 @@ class Tester(unittest.TestCase):
num_frames = 32 num_frames = 32
orig_fps = 30 orig_fps = 30
new_fps = 13 new_fps = 13
with self.assertWarns(UserWarning): with pytest.warns(UserWarning):
clips, idxs = VideoClips.compute_clips_for_video(video_pts, num_frames, num_frames, clips, idxs = VideoClips.compute_clips_for_video(video_pts, num_frames, num_frames, orig_fps, new_fps)
orig_fps, new_fps)
assert len(clips) == 0 assert len(clips) == 0
assert len(idxs) == 0 assert len(idxs) == 0
if __name__ == '__main__': if __name__ == "__main__":
unittest.main() pytest.main([__file__])
test/test_datasets_video_utils_opt.py
View file @ bf491463
import unittest import unittest
from torchvision import set_video_backend
import test_datasets_video_utils import test_datasets_video_utils
from torchvision import set_video_backend # noqa: 401
# Disabling the video backend switching temporarily # Disabling the video backend switching temporarily
# set_video_backend('video_reader') # set_video_backend('video_reader')
if __name__ == '__main__': if __name__ == "__main__":
suite = unittest.TestLoader().loadTestsFromModule(test_datasets_video_utils) suite = unittest.TestLoader().loadTestsFromModule(test_datasets_video_utils)
unittest.TextTestRunner(verbosity=1).run(suite) unittest.TextTestRunner(verbosity=1).run(suite)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment