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
