def test_extract_gzip(self): with get_tmp_dir() as temp_dir: with tempfile.NamedTemporaryFile(suffix='.gz') as f: with gzip.GzipFile(f.name, 'wb') as zf: zf.write('this is the content'.encode()) utils.extract_archive(f.name, temp_dir) f_name = os.path.join( temp_dir, os.path.splitext(os.path.basename(f.name))[0]) self.assertTrue(os.path.exists(f_name)) with open(os.path.join(f_name), 'r') as nf: data = nf.read() self.assertEqual(data, 'this is the content')
def test_write_png(img_path): with get_tmp_dir() as d: pil_image = Image.open(img_path) img_pil = torch.from_numpy(np.array(pil_image)) img_pil = img_pil.permute(2, 0, 1) filename, _ = os.path.splitext(os.path.basename(img_path)) torch_png = os.path.join(d, '{0}_torch.png'.format(filename)) write_png(img_pil, torch_png, compression_level=6) saved_image = torch.from_numpy(np.array(Image.open(torch_png))) saved_image = saved_image.permute(2, 0, 1) assert_equal(img_pil, saved_image)
def test_read_file(): with get_tmp_dir() as d: fname, content = 'test1.bin', b'TorchVision\211\n' fpath = os.path.join(d, fname) with open(fpath, 'wb') as f: f.write(content) data = read_file(fpath) expected = torch.tensor(list(content), dtype=torch.uint8) os.unlink(fpath) assert_equal(data, expected) with pytest.raises(RuntimeError, match="No such file or directory: 'tst'"): read_file('tst')
def create_dataset( self, config: Optional[Dict[str, Any]] = None, inject_fake_data: bool = True, patch_checks: Optional[bool] = None, **kwargs: Any, ) -> Iterator[Tuple[torchvision.datasets.VisionDataset, Dict[str, Any]]]: r"""Create the dataset in a temporary directory. Args: config (Optional[Dict[str, Any]]): Configuration that will be used to create the dataset. If omitted, the default configuration is used. inject_fake_data (bool): If ``True`` (default) inject the fake data with :meth:`.inject_fake_data` before creating the dataset. patch_checks (Optional[bool]): If ``True`` disable integrity check logic while creating the dataset. If omitted defaults to the same value as ``inject_fake_data``. **kwargs (Any): Additional parameters passed to the dataset. These parameters take precedence in case they overlap with ``config``. Yields: dataset (torchvision.dataset.VisionDataset): Dataset. info (Dict[str, Any]): Additional information about the injected fake data. See :meth:`.inject_fake_data` for details. """ default_config = self._DEFAULT_CONFIG.copy() if config is not None: default_config.update(config) config = default_config if patch_checks is None: patch_checks = inject_fake_data special_kwargs, other_kwargs = self._split_kwargs(kwargs) if "download" in self._HAS_SPECIAL_KWARG and special_kwargs.get("download", False): # override download param to False param if its default is truthy special_kwargs["download"] = False config.update(other_kwargs) patchers = self._patch_download_extract() if patch_checks: patchers.update(self._patch_checks()) with get_tmp_dir() as tmpdir: args = self.dataset_args(tmpdir, config) info = self._inject_fake_data(tmpdir, config) if inject_fake_data else None with self._maybe_apply_patches(patchers), disable_console_output(): dataset = self.DATASET_CLASS(*args, **config, **special_kwargs) yield dataset, info
def test_resize(self): # TODO: Minimal check for bug-fix, improve this later x = torch.rand(3, 32, 46) t = T.Resize(size=38) y = t(x) # If size is an int, smaller edge of the image will be matched to this number. # i.e, if height > width, then image will be rescaled to (size * height / width, size). self.assertTrue(isinstance(y, torch.Tensor)) self.assertEqual(y.shape[1], 38) self.assertEqual(y.shape[2], int(38 * 46 / 32)) tensor, _ = _create_data(height=34, width=36, device=self.device) batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device) for dt in [None, torch.float32, torch.float64]: if dt is not None: # This is a trivial cast to float of uint8 data to test all cases tensor = tensor.to(dt) for size in [32, 34, [ 32, ], [32, 32], (32, 32), [34, 35]]: for max_size in (None, 35, 1000): if max_size is not None and isinstance( size, Sequence) and len(size) != 1: continue # Not supported for interpolation in [BILINEAR, BICUBIC, NEAREST]: if isinstance(size, int): script_size = [ size, ] else: script_size = size transform = T.Resize(size=script_size, interpolation=interpolation, max_size=max_size) s_transform = torch.jit.script(transform) _test_transform_vs_scripted(transform, s_transform, tensor) _test_transform_vs_scripted_on_batch( transform, s_transform, batch_tensors) with get_tmp_dir() as tmp_dir: s_transform.save(os.path.join(tmp_dir, "t_resize.pt"))
def test_read_packed_b_frames_divx_file(self): with get_tmp_dir() as temp_dir: name = "hmdb51_Turnk_r_Pippi_Michel_cartwheel_f_cm_np2_le_med_6.avi" f_name = os.path.join(temp_dir, name) url = "https://download.pytorch.org/vision_tests/io/" + name try: utils.download_url(url, temp_dir) pts, fps = io.read_video_timestamps(f_name) self.assertEqual(pts, sorted(pts)) self.assertEqual(fps, 30) except URLError: msg = "could not download test file '{}'".format(url) warnings.warn(msg, RuntimeWarning) raise unittest.SkipTest(msg)
def svhn_root(): import scipy.io as sio def _make_mat(file): images = np.zeros((32, 32, 3, 2), dtype=np.uint8) targets = np.zeros((2, ), dtype=np.uint8) sio.savemat(file, {'X': images, 'y': targets}) with get_tmp_dir() as root: _make_mat(os.path.join(root, "train_32x32.mat")) _make_mat(os.path.join(root, "test_32x32.mat")) _make_mat(os.path.join(root, "extra_32x32.mat")) yield root
def test_normalize(self): tensor, _ = self._create_data(26, 34, device=self.device) batch_tensors = torch.rand(4, 3, 44, 56, device=self.device) tensor = tensor.to(dtype=torch.float32) / 255.0 # test for class interface fn = T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) scripted_fn = torch.jit.script(fn) self._test_transform_vs_scripted(fn, scripted_fn, tensor) self._test_transform_vs_scripted_on_batch(fn, scripted_fn, batch_tensors) with get_tmp_dir() as tmp_dir: scripted_fn.save(os.path.join(tmp_dir, "t_norm.pt"))
def test_extract_tar(self): for ext, mode in zip(['.tar', '.tar.gz', '.tgz'], ['w', 'w:gz', 'w:gz']): with get_tmp_dir() as temp_dir: with tempfile.NamedTemporaryFile() as bf: bf.write("this is the content".encode()) bf.seek(0) with tempfile.NamedTemporaryFile(suffix=ext) as f: with tarfile.open(f.name, mode=mode) as zf: zf.add(bf.name, arcname='file.tst') utils.extract_archive(f.name, temp_dir) self.assertTrue(os.path.exists(os.path.join(temp_dir, 'file.tst'))) with open(os.path.join(temp_dir, 'file.tst'), 'r') as nf: data = nf.read() self.assertEqual(data, 'this is the content')
def test_imagefolder(self): # TODO: create the fake data on-the-fly FAKEDATA_DIR = get_file_path_2( os.path.dirname(os.path.abspath(__file__)), 'assets', 'fakedata') with get_tmp_dir(src=os.path.join(FAKEDATA_DIR, 'imagefolder')) as root: classes = sorted(['a', 'b']) class_a_image_files = [ os.path.join(root, 'a', file) for file in ('a1.png', 'a2.png', 'a3.png') ] class_b_image_files = [ os.path.join(root, 'b', file) for file in ('b1.png', 'b2.png', 'b3.png', 'b4.png') ] dataset = torchvision.datasets.ImageFolder(root, loader=lambda x: x) # test if all classes are present self.assertEqual(classes, sorted(dataset.classes)) # test if combination of classes and class_to_index functions correctly for cls in classes: self.assertEqual(cls, dataset.classes[dataset.class_to_idx[cls]]) # test if all images were detected correctly class_a_idx = dataset.class_to_idx['a'] class_b_idx = dataset.class_to_idx['b'] imgs_a = [(img_file, class_a_idx) for img_file in class_a_image_files] imgs_b = [(img_file, class_b_idx) for img_file in class_b_image_files] imgs = sorted(imgs_a + imgs_b) self.assertEqual(imgs, dataset.imgs) # test if the datasets outputs all images correctly outputs = sorted([dataset[i] for i in range(len(dataset))]) self.assertEqual(imgs, outputs) # redo all tests with specified valid image files dataset = torchvision.datasets.ImageFolder( root, loader=lambda x: x, is_valid_file=lambda x: '3' in x) self.assertEqual(classes, sorted(dataset.classes)) class_a_idx = dataset.class_to_idx['a'] class_b_idx = dataset.class_to_idx['b'] imgs_a = [(img_file, class_a_idx) for img_file in class_a_image_files if '3' in img_file] imgs_b = [(img_file, class_b_idx) for img_file in class_b_image_files if '3' in img_file] imgs = sorted(imgs_a + imgs_b) self.assertEqual(imgs, dataset.imgs) outputs = sorted([dataset[i] for i in range(len(dataset))]) self.assertEqual(imgs, outputs)
def test_autoaugment(self): tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=self.device) batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device) for policy in T.AutoAugmentPolicy: for fill in [None, 85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [1, ], 1]: for _ in range(100): transform = T.AutoAugment(policy=policy, fill=fill) s_transform = torch.jit.script(transform) self._test_transform_vs_scripted(transform, s_transform, tensor) self._test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors) with get_tmp_dir() as tmp_dir: s_transform.save(os.path.join(tmp_dir, "t_autoaugment.pt"))
def test_read_file(self): with get_tmp_dir() as d: fname, content = 'test1.bin', b'TorchVision\211\n' fpath = os.path.join(d, fname) with open(fpath, 'wb') as f: f.write(content) data = read_file(fpath) expected = torch.tensor(list(content), dtype=torch.uint8) self.assertTrue(data.equal(expected)) os.unlink(fpath) with self.assertRaisesRegex( RuntimeError, "No such file or directory: 'tst'"): read_file('tst')
def _make_data_archive(root): with get_tmp_dir() as tmp_dir: base_dir = 'ADEChallengeData2016' for folder_name in ['images', 'annotations']: folder_dir = os.path.join(tmp_dir, base_dir, folder_name) os.makedirs(folder_dir) for split_name in ['training', 'validation']: split_dir = os.path.join(folder_dir, split_name) os.makedirs(os.path.join(split_dir)) _make_image(os.path.join(split_dir, 'ADE_train_00000000.png')) archive = os.path.join(root, 'ADEChallengeData2016.zip') _make_zip(archive=archive, tmp_dir=tmp_dir, content=os.path.join(tmp_dir, base_dir))
def test_random_affine(self): tensor = torch.randint(0, 256, size=(3, 44, 56), dtype=torch.uint8, device=self.device) batch_tensors = torch.randint(0, 256, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device) def _test(**kwargs): transform = T.RandomAffine(**kwargs) s_transform = torch.jit.script(transform) _test_transform_vs_scripted(transform, s_transform, tensor) _test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors) return s_transform for interpolation in [NEAREST, BILINEAR]: for shear in [ 15, 10.0, (5.0, 10.0), [-15, 15], [-10.0, 10.0, -11.0, 11.0] ]: _test(degrees=0.0, interpolation=interpolation, shear=shear) for scale in [(0.7, 1.2), [0.7, 1.2]]: _test(degrees=0.0, interpolation=interpolation, scale=scale) for translate in [(0.1, 0.2), [0.2, 0.1]]: _test(degrees=0.0, interpolation=interpolation, translate=translate) for degrees in [45, 35.0, (-45, 45), [-90.0, 90.0]]: _test(degrees=degrees, interpolation=interpolation) for fill in [85, (10, -10, 10), 0.7, [0.0, 0.0, 0.0], [ 1, ], 1]: _test(degrees=0.0, interpolation=interpolation, fill=fill) s_transform = _test(degrees=0.0) with get_tmp_dir() as tmp_dir: s_transform.save(os.path.join(tmp_dir, "t_random_affine.pt"))
def test_decompress_remove_finished(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.extract_archive(compressed, temp_dir, remove_finished=True) self.assertFalse(os.path.exists(compressed))
def test_normalize(device): fn = T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) tensor, _ = _create_data(26, 34, device=device) with pytest.raises(TypeError, match="Input tensor should be a float tensor"): fn(tensor) batch_tensors = torch.rand(4, 3, 44, 56, device=device) tensor = tensor.to(dtype=torch.float32) / 255.0 # test for class interface scripted_fn = torch.jit.script(fn) _test_transform_vs_scripted(fn, scripted_fn, tensor) _test_transform_vs_scripted_on_batch(fn, scripted_fn, batch_tensors) with get_tmp_dir() as tmp_dir: scripted_fn.save(os.path.join(tmp_dir, "t_norm.pt"))
def test_random_perspective(self): tensor = torch.randint(0, 255, size=(3, 44, 56), dtype=torch.uint8, device=self.device) batch_tensors = torch.randint(0, 255, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device) for distortion_scale in np.linspace(0.1, 1.0, num=20): for interpolation in [NEAREST, BILINEAR]: transform = T.RandomPerspective( distortion_scale=distortion_scale, interpolation=interpolation ) s_transform = torch.jit.script(transform) self._test_transform_vs_scripted(transform, s_transform, tensor) self._test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors) with get_tmp_dir() as tmp_dir: s_transform.save(os.path.join(tmp_dir, "t_perspective.pt"))
def test_resized_crop(self): tensor = torch.randint(0, 255, size=(3, 44, 56), dtype=torch.uint8, device=self.device) batch_tensors = torch.randint(0, 255, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device) for scale in [(0.7, 1.2), [0.7, 1.2]]: for ratio in [(0.75, 1.333), [0.75, 1.333]]: for size in [(32, ), [44, ], [32, ], [32, 32], (32, 32), [44, 55]]: for interpolation in [NEAREST, BILINEAR, BICUBIC]: transform = T.RandomResizedCrop( size=size, scale=scale, ratio=ratio, interpolation=interpolation ) s_transform = torch.jit.script(transform) self._test_transform_vs_scripted(transform, s_transform, tensor) self._test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors) with get_tmp_dir() as tmp_dir: s_transform.save(os.path.join(tmp_dir, "t_resized_crop.pt"))
def _test_op_list_output(self, func, method, out_length, fn_kwargs=None, meth_kwargs=None): if fn_kwargs is None: fn_kwargs = {} if meth_kwargs is None: meth_kwargs = {} fn = getattr(F, func) scripted_fn = torch.jit.script(fn) tensor, pil_img = self._create_data(height=20, width=20, device=self.device) transformed_t_list = fn(tensor, **fn_kwargs) transformed_p_list = fn(pil_img, **fn_kwargs) self.assertEqual(len(transformed_t_list), len(transformed_p_list)) self.assertEqual(len(transformed_t_list), out_length) for transformed_tensor, transformed_pil_img in zip(transformed_t_list, transformed_p_list): self.compareTensorToPIL(transformed_tensor, transformed_pil_img) transformed_t_list_script = scripted_fn(tensor.detach().clone(), **fn_kwargs) self.assertEqual(len(transformed_t_list), len(transformed_t_list_script)) self.assertEqual(len(transformed_t_list_script), out_length) for transformed_tensor, transformed_tensor_script in zip(transformed_t_list, transformed_t_list_script): self.assertTrue(transformed_tensor.equal(transformed_tensor_script), msg="{} vs {}".format(transformed_tensor, transformed_tensor_script)) # test for class interface fn = getattr(T, method)(**meth_kwargs) scripted_fn = torch.jit.script(fn) output = scripted_fn(tensor) self.assertEqual(len(output), len(transformed_t_list_script)) # test on batch of tensors batch_tensors = self._create_data_batch(height=23, width=34, channels=3, num_samples=4, device=self.device) torch.manual_seed(12) transformed_batch_list = fn(batch_tensors) for i in range(len(batch_tensors)): img_tensor = batch_tensors[i, ...] torch.manual_seed(12) transformed_img_list = fn(img_tensor) for transformed_img, transformed_batch in zip(transformed_img_list, transformed_batch_list): self.assertTrue(transformed_img.equal(transformed_batch[i, ...]), msg="{} vs {}".format(transformed_img, transformed_batch[i, ...])) with get_tmp_dir() as tmp_dir: scripted_fn.save(os.path.join(tmp_dir, "t_op_list_{}.pt".format(method)))
def test_random_rotate(self): tensor = torch.randint(0, 255, size=(3, 44, 56), dtype=torch.uint8, device=self.device) batch_tensors = torch.randint(0, 255, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device) for center in [(0, 0), [10, 10], None, (56, 44)]: for expand in [True, False]: for degrees in [45, 35.0, (-45, 45), [-90.0, 90.0]]: for interpolation in [NEAREST, BILINEAR]: transform = T.RandomRotation( degrees=degrees, resample=interpolation, expand=expand, center=center ) s_transform = torch.jit.script(transform) self._test_transform_vs_scripted(transform, s_transform, tensor) self._test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors) with get_tmp_dir() as tmp_dir: s_transform.save(os.path.join(tmp_dir, "t_random_rotate.pt"))
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, 255, (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
def test_resize(self): tensor, _ = self._create_data(height=34, width=36, device=self.device) batch_tensors = torch.randint(0, 255, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device) script_fn = torch.jit.script(F.resize) for dt in [None, torch.float32, torch.float64]: if dt is not None: # This is a trivial cast to float of uint8 data to test all cases tensor = tensor.to(dt) for size in [32, 34, [ 32, ], [32, 32], (32, 32), [34, 35]]: for interpolation in [BILINEAR, BICUBIC, NEAREST]: resized_tensor = F.resize(tensor, size=size, interpolation=interpolation) if isinstance(size, int): script_size = [ size, ] else: script_size = size s_resized_tensor = script_fn(tensor, size=script_size, interpolation=interpolation) self.assertTrue(s_resized_tensor.equal(resized_tensor)) transform = T.Resize(size=script_size, interpolation=interpolation) s_transform = torch.jit.script(transform) self._test_transform_vs_scripted(transform, s_transform, tensor) self._test_transform_vs_scripted_on_batch( transform, s_transform, batch_tensors) with get_tmp_dir() as tmp_dir: script_fn.save(os.path.join(tmp_dir, "t_resize.pt"))
def test_read_file(self): with get_tmp_dir() as d: fname, content = 'test1.bin', b'TorchVision\211\n' fpath = os.path.join(d, fname) with open(fpath, 'wb') as f: f.write(content) data = read_file(fpath) expected = torch.tensor(list(content), dtype=torch.uint8) self.assertTrue(data.equal(expected)) # Windows holds into the file until the tensor is alive # so need to del the tensor before deleting the file see # https://github.com/pytorch/vision/issues/2743#issuecomment-703817293 del data os.unlink(fpath) with self.assertRaisesRegex(RuntimeError, "No such file or directory: 'tst'"): read_file('tst')
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_write_jpeg(img_path): with get_tmp_dir() as d: d = Path(d) img = read_image(img_path) pil_img = F.to_pil_image(img) torch_jpeg = str(d / 'torch.jpg') pil_jpeg = str(d / 'pil.jpg') write_jpeg(img, torch_jpeg, quality=75) pil_img.save(pil_jpeg, quality=75) with open(torch_jpeg, 'rb') as f: torch_bytes = f.read() with open(pil_jpeg, 'rb') as f: pil_bytes = f.read() assert_equal(torch_bytes, pil_bytes)
def test_extract_zip(self): def create_archive(root, content="this is the content"): file = os.path.join(root, "dst.txt") archive = os.path.join(root, "archive.zip") with zipfile.ZipFile(archive, "w") as zf: zf.writestr(os.path.basename(file), content) return archive, file, content with get_tmp_dir() as temp_dir: archive, file, content = create_archive(temp_dir) utils.extract_archive(archive, temp_dir) self.assertTrue(os.path.exists(file)) with open(file, "r") as fh: self.assertEqual(fh.read(), content)
def ucf101_root(): with get_tmp_dir() as tmp_dir: ucf_dir = os.path.join(tmp_dir, 'UCF-101') video_dir = os.path.join(ucf_dir, 'video') annotations = os.path.join(ucf_dir, 'annotations') os.makedirs(ucf_dir) os.makedirs(video_dir) os.makedirs(annotations) fold_files = [] for split in {'train', 'test'}: for fold in range(1, 4): fold_file = '{:s}list{:02d}.txt'.format(split, fold) fold_files.append(os.path.join(annotations, fold_file)) file_handles = [open(x, 'w') for x in fold_files] file_iter = cycle(file_handles) for i in range(0, 2): current_class = 'class_{0}'.format(i + 1) class_dir = os.path.join(video_dir, current_class) os.makedirs(class_dir) for group in range(0, 3): for clip in range(0, 4): # Save sample file clip_name = 'v_{0}_g{1}_c{2}.avi'.format( current_class, group, clip) clip_path = os.path.join(class_dir, clip_name) length = random.randrange(10, 21) this_clip = torch.randint(0, 256, (length * 25, 320, 240, 3), dtype=torch.uint8) write_video(clip_path, this_clip, 25) # Add to annotations ann_file = next(file_iter) ann_file.write('{0}\n'.format( os.path.join(current_class, clip_name))) # Close all file descriptors for f in file_handles: f.close() yield (video_dir, annotations)
def test_decompress(self, extension): def create_compressed(root, content="this is the content"): file = os.path.join(root, "file") compressed = f"{file}{extension}" compressed_file_opener = _COMPRESSED_FILE_OPENERS[extension] with compressed_file_opener(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) assert os.path.exists(file) with open(file, "r") as fh: assert fh.read() == content
def test_random_affine(self): tensor = torch.randint(0, 255, size=(3, 44, 56), dtype=torch.uint8, device=self.device) batch_tensors = torch.randint(0, 255, size=(4, 3, 44, 56), dtype=torch.uint8, device=self.device) for shear in [15, 10.0, (5.0, 10.0), [-15, 15], [-10.0, 10.0, -11.0, 11.0]]: for scale in [(0.7, 1.2), [0.7, 1.2]]: for translate in [(0.1, 0.2), [0.2, 0.1]]: for degrees in [45, 35.0, (-45, 45), [-90.0, 90.0]]: for interpolation in [NEAREST, BILINEAR]: transform = T.RandomAffine( degrees=degrees, translate=translate, scale=scale, shear=shear, resample=interpolation ) s_transform = torch.jit.script(transform) self._test_transform_vs_scripted(transform, s_transform, tensor) self._test_transform_vs_scripted_on_batch(transform, s_transform, batch_tensors) with get_tmp_dir() as tmp_dir: s_transform.save(os.path.join(tmp_dir, "t_random_affine.pt"))
def test_write_jpeg(self): with get_tmp_dir() as d: for img_path in get_images(ENCODE_JPEG, ".jpg"): data = read_file(img_path) img = decode_jpeg(data) basedir = os.path.dirname(img_path) filename, _ = os.path.splitext(os.path.basename(img_path)) torch_jpeg = os.path.join(d, '{0}_torch.jpg'.format(filename)) pil_jpeg = os.path.join(basedir, 'jpeg_write', '{0}_pil.jpg'.format(filename)) write_jpeg(img, torch_jpeg, quality=75) with open(torch_jpeg, 'rb') as f: torch_bytes = f.read() with open(pil_jpeg, 'rb') as f: pil_bytes = f.read() self.assertEqual(torch_bytes, pil_bytes)