def test_extract_aspect_ratio(): height = 2 width = 3 image = torch.empty(1, 1, height, width) actual = image_.extract_aspect_ratio(image) desired = width / height assert actual == pytest.approx(desired)
def demo_images(): cache_path = pathlib.Path(pystiche.home()) graphics_path = GRAPHICS / "demo_images" graphics_path.mkdir(exist_ok=True) api_path = HERE / "api" images = pystiche.demo.images() images.download() entries = {} for name, image in images: entries[name] = (image.file, extract_aspect_ratio(image.read())) if not (graphics_path / image.file).exists(): (graphics_path / image.file).symlink_to(cache_path / image.file) field_len = max( max(len(name) for name in entries.keys()) + 2, len("images")) def sep(char): return "+" + char * (field_len + 2) + "+" + char * (field_len + 2) + "+" def row(name, image): key = f"{name:{field_len}}" value = image or f"|{name}|" value += " " * (field_len - len(image or name)) return f"| {key} | {value} |" images_table = [ sep("-"), row("name", "image"), sep("="), *itertools.chain(*[(row(name, f"|{name}|"), sep("-")) for name in sorted(entries.keys())]), ] width = 300 aliases = [ f".. |{name}|\n" f" image:: ../graphics/demo_images/{file}\n" f" :width: {width}px\n" f" :height: {width / aspect_ratio:.0f}px\n" for name, (file, aspect_ratio) in entries.items() ] loader = jinja2.FileSystemLoader(searchpath=api_path) env = jinja2.Environment(loader=loader) template = env.get_template("pystiche.demo.rst.template") with open(api_path / "pystiche.demo.rst", "w") as fh: fh.write( template.render( images_table="\n".join(images_table), aliases="\n".join(aliases), )) return None, None
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 _resize( self, image: torch.Tensor, aspect_ratio: Optional[float], interpolation_mode: str, ) -> torch.Tensor: if aspect_ratio is None: aspect_ratio = extract_aspect_ratio(image) image_size = edge_to_image_size(self.edge_size, aspect_ratio, edge=self.edge) with torch.no_grad(): image = resize(image, image_size, interpolation=interpolation_mode) return image.detach()
def pyramid_image_optimization( 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, ) -> 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``, no information is printed to STDOUT during the optimization. Defaults to ``False``. """ aspect_ratio = extract_aspect_ratio(input_image) if get_optimizer is None: get_optimizer = default_image_optimizer output_image = input_image for level in OptimProgressBar("Pyramid", pyramid, quiet=quiet): with torch.no_grad(): input_image = level.resize_image(output_image, aspect_ratio=aspect_ratio) output_image = image_optimization( input_image, criterion, optimizer=get_optimizer, num_steps=level.num_steps, preprocessor=preprocessor, postprocessor=postprocessor, quiet=quiet, ) return output_image
def _generate_default_image_pyramid_optim_loop_asset( file, input_image, criterion, pyramid, get_optimizer=None, preprocessor=None, postprocessor=None, ): if get_optimizer is None: get_optimizer = default_image_optimizer aspect_ratio = extract_aspect_ratio(input_image) output_image = input_image.clone() for level in pyramid: with torch.no_grad(): output_image = level.resize_image(output_image, aspect_ratio=aspect_ratio) if preprocessor is not None: output_image = preprocessor(output_image) optimizer = get_optimizer(output_image) for _ in level: def closure(): optimizer.zero_grad() loss = criterion(output_image) loss.backward() return loss optimizer.step(closure) output_image = output_image.detach() if postprocessor is not None: output_image = postprocessor(output_image) input = {"image": input_image, "criterion": criterion, "pyramid": pyramid} params = { "get_optimizer": get_optimizer, "preprocessor": preprocessor, "postprocessor": postprocessor, } output = {"image": output_image} store_asset(input, params, output, file)
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