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)
