def test_AutoPadAvgPool2d(subtests, auto_pad_pool_params, input_image): image_size = extract_image_size(input_image) for params in auto_pad_pool_params: with subtests.test(**params): conv = utils.AutoPadAvgPool2d(**params) output_image = conv(input_image) actual = extract_image_size(output_image) expected = tuple(side_length // stride for side_length, stride in zip( image_size, to_2d_arg(params["stride"]))) assert actual == expected
def test_AutoPadConvTranspose2d(subtests, auto_pad_conv_params, input_image): in_channels = out_channels = extract_num_channels(input_image) image_size = extract_image_size(input_image) for params in auto_pad_conv_params: with subtests.test(**params): conv = utils.AutoPadConvTranspose2d(in_channels, out_channels, **params) output_image = conv(input_image) actual = extract_image_size(output_image) expected = tuple(side_length * stride for side_length, stride in zip( image_size, to_2d_arg(params["stride"]))) assert actual == expected
def test_li_wand_2016_nst_smoke(subtests, mocker, content_image, style_image): spy = mocker.patch( mocks.make_mock_target("li_wand_2016", "_nst", "misc", "get_input_image"), wraps=get_input_image, ) mock = mocker.patch( mocks.make_mock_target("li_wand_2016", "_nst", "optim", "pyramid_image_optimization")) hyper_parameters = paper.hyper_parameters() paper.nst(content_image, style_image) args, kwargs = mock.call_args input_image, criterion, pyramid = args get_optimizer = kwargs["get_optimizer"] preprocessor = kwargs["preprocessor"] postprocessor = kwargs["postprocessor"] initial_resize = pyramid[-1].resize_image with subtests.test("input_image"): args = utils.call_args_to_namespace(spy.call_args, get_input_image) assert args.starting_point == hyper_parameters.nst.starting_point assert extract_image_size(args.content_image) == extract_image_size( initial_resize(content_image)) with subtests.test("style_image"): desired_style_image = preprocessor(initial_resize(style_image)) for loss in criterion.style_loss.children(): ptu.assert_allclose(loss.target_image, desired_style_image) with subtests.test("criterion"): assert isinstance(criterion, type(paper.perceptual_loss())) with subtests.test("pyramid"): assert isinstance(pyramid, type(paper.image_pyramid())) with subtests.test("optimizer"): assert is_callable(get_optimizer) optimizer = get_optimizer(input_image) assert isinstance(optimizer, type(paper.optimizer(input_image))) with subtests.test("preprocessor"): assert isinstance(preprocessor, type(paper.preprocessor())) with subtests.test("postprocessor"): assert isinstance(postprocessor, type(paper.postprocessor()))
def get_size( input: Union[Tuple[int, int], torch.Tensor]) -> Tuple[int, int]: if isinstance(input, torch.Tensor): return extract_image_size(input) else: return input
def test_iter_resize(self): class TestOperator(ops.PixelComparisonOperator): def target_image_to_repr(self, image): return image, None def input_image_to_repr(self, image, ctx): pass def calculate_score(self, input_repr, target_repr, ctx): pass initial_image_size = (5, 4) edge_sizes = (2, 4) torch.manual_seed(0) target_guide = torch.rand((1, 3, *initial_image_size)) target_image = torch.rand((1, 3, *initial_image_size)) input_guide = torch.rand((1, 3, *initial_image_size)) aspect_ratio = calculate_aspect_ratio(initial_image_size) image_sizes = [ edge_to_image_size(edge_size, aspect_ratio) for edge_size in edge_sizes ] op = TestOperator() op.set_target_guide(target_guide) op.set_target_image(target_image) op.set_input_guide(input_guide) image_pyramid = pyramid.ImagePyramid(edge_sizes, 1, resize_targets=(op,)) for pyramid_level, image_size in zip(image_pyramid, image_sizes): for attr in ("target_guide", "target_image", "input_guide"): actual = extract_image_size(getattr(op, attr)) desired = image_size assert actual == desired
def test_extract_image_size(): height = 2 width = 3 image = torch.empty(1, 1, height, width) actual = image_.extract_image_size(image) desired = (height, width) assert actual == desired
def test_ValidRandomCrop_identity(self): image = self.load_image() size = extract_image_size(image) transform = transforms.ValidRandomCrop(size) actual = transform(image) desired = image self.assertImagesAlmostEqual(actual, desired)
def forward(self, image: torch.Tensor) -> torch.Tensor: old_height, old_width = extract_image_size(image) if old_height > old_width: new_height = self.edge_size new_width = int(new_height / old_height * old_width) else: new_width = self.edge_size new_height = int(new_width / old_width * old_height) return cast(torch.Tensor, F.resize(image, [new_height, new_width]))
def forward(self, image: torch.Tensor) -> torch.Tensor: return cast( torch.Tensor, F.resize( image, [ round(length * self.factor) for length in extract_image_size(image) ], ), )
def default_image_pyramid_optim_loop( input_image: torch.Tensor, criterion: nn.Module, pyramid: ImagePyramid, get_optimizer: Optional[Callable[[torch.Tensor], Optimizer]] = None, preprocessor: Optional[nn.Module] = None, postprocessor: Optional[nn.Module] = None, quiet: bool = False, logger: Optional[OptimLogger] = None, get_pyramid_level_header: Optional[Callable[ [int, PyramidLevel, Tuple[int, int]], str]] = None, log_fn: Optional[Callable[[int, Union[torch.Tensor, pystiche.LossDict]], None]] = None, ) -> torch.Tensor: aspect_ratio = extract_aspect_ratio(input_image) if get_optimizer is None: get_optimizer = default_image_optimizer if logger is None: logger = OptimLogger() if get_pyramid_level_header is None: get_pyramid_level_header = default_pyramid_level_header output_image = input_image for num, level in enumerate(pyramid, 1): def image_optim_loop(input_image: torch.Tensor) -> torch.Tensor: return default_image_optim_loop( input_image, criterion, get_optimizer=get_optimizer, num_steps=iter(level), preprocessor=preprocessor, postprocessor=postprocessor, quiet=quiet, logger=logger, log_fn=log_fn, ) with torch.no_grad(): input_image = level.resize_image(output_image, aspect_ratio=aspect_ratio) if quiet: output_image = image_optim_loop(input_image) else: input_image_size = extract_image_size(input_image) header = get_pyramid_level_header(num, level, input_image_size) with logger.environment(header): output_image = image_optim_loop(input_image) return output_image
def forward(self, image: torch.Tensor) -> torch.Tensor: top, left = self.get_params(extract_image_size(image), self.size) height, width = self.size return cast( torch.Tensor, F.crop( image, top=top, left=left, height=height, width=width, ), )
def create_guides(img): height, width = image.extract_image_size(img) top_height = height // 2 bottom_height = height - top_height top_mask = torch.cat( ( torch.ones([1, 1, top_height, width], dtype=torch.bool), torch.zeros([1, 1, bottom_height, width], dtype=torch.bool), ), 2, ) bottom_mask = ~top_mask return {"top": top_mask.float(), "bottom": bottom_mask.float()}
def forward(self, image: torch.Tensor) -> torch.Tensor: height, width = extract_image_size(image) alpha = math.radians(self.angle) if not self.clockwise: alpha *= -1 bounding_box = _computeBB(width, height, alpha) top = bounding_box[1] left = bounding_box[0] height = bounding_box[3] - top width = bounding_box[2] - left return cast( torch.Tensor, F.crop(image, top=top, left=left, height=height, width=width))
def make_input_image(starting_point: str, *, content_image: torch.Tensor) -> torch.Tensor: if starting_point == "content": return content_image.clone() elif starting_point == "random": return torch.rand_like(content_image) else: input_image = load_image( starting_point, size=None, device=content_image.device, ) input_image_size = extract_image_size(input_image) content_image_size = extract_image_size(content_image) if extract_image_size(input_image) != content_image_size: warnings.warn( f"Image size of starting point and content image mismatches: " f"{input_image_size} != {content_image_size}. " f"Resizing to {content_image_size} to move forward.") input_image = resize(input_image, list(content_image_size)) return input_image
def calculate_size(self, image: torch.Tensor) -> Tuple[int, int]: old_height, old_width = extract_image_size(image) new_height = old_height - old_height % self.multiple new_width = old_width - old_width % self.multiple return new_height, new_width
def default_image_pyramid_optim_loop( input_image: torch.Tensor, criterion: nn.Module, pyramid: ImagePyramid, get_optimizer: Optional[Callable[[torch.Tensor], Optimizer]] = None, preprocessor: Optional[nn.Module] = None, postprocessor: Optional[nn.Module] = None, quiet: bool = False, logger: Optional[OptimLogger] = None, get_pyramid_level_header: Optional[Callable[ [int, PyramidLevel, Tuple[int, int]], str]] = None, log_fn: Optional[Callable[[int, Union[torch.Tensor, pystiche.LossDict]], None]] = None, ) -> torch.Tensor: r"""Perform a image optimization for :class:`pystiche.pyramid.ImagePyramid` s with integrated logging. Args: input_image: Image to be optimized. criterion: Optimization criterion. pyramid: Image pyramid. get_optimizer: Optional getter for the optimizer. If ``None``, :func:`default_image_optimizer` is used. Defaults to ``None``. preprocessor: Optional preprocessor that is called with the ``input_image`` before the optimization. postprocessor: Optional preprocessor that is called with the ``input_image`` after the optimization. quiet: If ``True``, not information is logged during the optimization. Defaults to ``False``. logger: Optional custom logger. If ``None``, :class:`pystiche.optim.OptimLogger` is used. Defaults to ``None``. get_pyramid_level_header: Optional custom getter for the logged pyramid level header. It is called before each level with the current level number, the :class:`pystiche.pyramid.PyramidLevel`, and the size of the ``input_image``. If ``None`` :func:`pystiche.optim.default_pyramid_level_header` is used. Defaults to ``None``. log_fn: Optional custom logging function. It is called in every optimization step with the current step and loss. If ``None``, :func:`pystiche.optim.default_image_optim_log_fn` is used. Defaults to ``None``. """ aspect_ratio = extract_aspect_ratio(input_image) if get_optimizer is None: get_optimizer = default_image_optimizer if logger is None: logger = OptimLogger() if get_pyramid_level_header is None: get_pyramid_level_header = default_pyramid_level_header output_image = input_image for num, level in enumerate(pyramid, 1): def image_optim_loop(input_image: torch.Tensor) -> torch.Tensor: return default_image_optim_loop( input_image, criterion, get_optimizer=get_optimizer, num_steps=iter(level), preprocessor=preprocessor, postprocessor=postprocessor, quiet=quiet, logger=logger, log_fn=log_fn, ) with torch.no_grad(): input_image = level.resize_image(output_image, aspect_ratio=aspect_ratio) if quiet: output_image = image_optim_loop(input_image) else: input_image_size = extract_image_size(input_image) header = get_pyramid_level_header(num, level, input_image_size) with logger.environment(header): output_image = image_optim_loop(input_image) return output_image
def forward(self, input_image: torch.Tensor) -> torch.Tensor: old_height, old_width = image.extract_image_size(input_image) new_height = old_height - old_height % self.multiple new_width = old_width - old_width % self.multiple return input_image[..., :new_height, :new_width]
def rescale(image, factor=2.0): return F.resize( image, [int(length * factor) for length in extract_image_size(image)])
def forward(self, image: torch.Tensor) -> torch.Tensor: origin = self.get_random_origin(extract_image_size(image), self.size) return cast(torch.Tensor, F.crop(image, origin, self.size))
def get_eval_transform(image): eval_transform = Resize(extract_image_size(image)) + RGBToGrayscale() return eval_transform.to(image.device)
def extract_image_sizes(images): return zip(*[extract_image_size(image) for image in images])
def test_ValidRandomCrop_identity(test_image): size = image_.extract_image_size(test_image) transform = transforms.ValidRandomCrop(size) assert_is_identity_transform(transform, test_image)