def test_features_image(self, p): input, expected = self.input_expected_image_tensor(p) transform = transforms.RandomVerticalFlip(p=p) actual = transform(features.Image(input)) assert_equal(features.Image(expected), actual)
def make_image(size=None, *, color_space, extra_dims=(), dtype=torch.float32, constant_alpha=True): size = size or torch.randint(16, 33, (2, )).tolist() try: num_channels = { features.ColorSpace.GRAY: 1, features.ColorSpace.GRAY_ALPHA: 2, features.ColorSpace.RGB: 3, features.ColorSpace.RGB_ALPHA: 4, }[color_space] except KeyError as error: raise pytest.UsageError() from error shape = (*extra_dims, num_channels, *size) max_value = get_max_value(dtype) data = make_tensor(shape, low=0, high=max_value, dtype=dtype) if color_space in { features.ColorSpace.GRAY_ALPHA, features.ColorSpace.RGB_ALPHA } and constant_alpha: data[..., -1, :, :] = max_value return features.Image(data, color_space=color_space)
def _prepare_sample(self, data: Tuple[Any, Any]) -> Dict[str, Any]: image, target = data # They're both numpy arrays at this point return { "image": features.Image(image), "label": Label(target.item()), }
def _transform(self, input: Any, params: Dict[str, Any]) -> Any: if isinstance(input, (features.Image, PIL.Image.Image, np.ndarray)) or is_simple_tensor(input): output = F.to_image_tensor(input, copy=self.copy) return features.Image(output) else: return input
def _prepare_sample(self, data: Tuple[Any, Any]) -> Dict[str, Any]: image, target = data # They're both numpy arrays at this point return { "image": features.Image(image.transpose(2, 0, 1)), "label": Label(target.item(), categories=self._categories), }
def make_image(size=None, *, color_space, extra_dims=(), dtype=torch.float32): size = size or torch.randint(16, 33, (2, )).tolist() num_channels = { features.ColorSpace.GRAYSCALE: 1, features.ColorSpace.RGB: 3, }[color_space] shape = (*extra_dims, num_channels, *size) if dtype.is_floating_point: data = torch.rand(shape, dtype=dtype) else: data = torch.randint(0, torch.iinfo(dtype).max, shape, dtype=dtype) return features.Image(data, color_space=color_space)
def _transform(self, input: Any, params: Dict[str, Any]) -> Any: if type(input) is features.EncodedImage: output = F.decode_image_with_pil(input) return features.Image(output) else: return input
def pil(buffer: io.IOBase) -> features.Image: return features.Image(pil_to_tensor(PIL.Image.open(buffer)))
def make_image(**kwargs): data = make_tensor((3, *torch.randint(16, 33, (2, )).tolist())) return features.Image(data, **kwargs)