def forward(self, laf: torch.Tensor,
img: torch.Tensor) -> torch.Tensor: # type: ignore
"""
Args:
laf: (torch.Tensor) shape [BxNx2x3]
img: (torch.Tensor) shape [Bx1xHxW]
Returns:
laf_out: (torch.Tensor) shape [BxNx2x3]"""
raise_error_if_laf_is_not_valid(laf)
img_message: str = "Invalid img shape, we expect BxCxHxW. Got: {}".format(
img.shape)
if not torch.is_tensor(img):
raise TypeError("img type is not a torch.Tensor. Got {}".format(
type(img)))
if len(img.shape) != 4:
raise ValueError(img_message)
if laf.size(0) != img.size(0):
raise ValueError(
"Batch size of laf and img should be the same. Got {}, {}".
format(img.size(0), laf.size(0)))
B, N = laf.shape[:2]
PS: int = self.patch_size
patches: torch.Tensor = extract_patches_from_pyramid(
img, make_upright(laf), PS, True).view(-1, 1, PS, PS)
ellipse_shape: torch.Tensor = self.affine_shape_detector(patches)
ellipses = torch.cat(
[laf.view(-1, 2, 3)[..., 2].unsqueeze(1), ellipse_shape],
dim=2).view(B, N, 5)
scale_orig = get_laf_scale(laf)
laf_out = ellipse_to_laf(ellipses)
ellipse_scale = get_laf_scale(laf_out)
laf_out = scale_laf(laf_out, scale_orig / ellipse_scale)
return laf_out
def forward(self, laf: torch.Tensor, img: torch.Tensor) -> torch.Tensor:
"""
Args:
laf: shape [BxNx2x3]
img: shape [Bx1xHxW]
Returns:
laf_out shape [BxNx2x3]
"""
raise_error_if_laf_is_not_valid(laf)
img_message: str = "Invalid img shape, we expect BxCxHxW. Got: {}".format(img.shape)
if not torch.is_tensor(img):
raise TypeError("img type is not a torch.Tensor. Got {}".format(type(img)))
if len(img.shape) != 4:
raise ValueError(img_message)
if laf.size(0) != img.size(0):
raise ValueError(
"Batch size of laf and img should be the same. Got {}, {}".format(img.size(0), laf.size(0))
)
B, N = laf.shape[:2]
PS: int = self.patch_size
patches: torch.Tensor = extract_patches_from_pyramid(img, make_upright(laf), PS, True).view(-1, 1, PS, PS)
xy = self.features(self._normalize_input(patches)).view(-1, 3)
a1 = torch.cat([1.0 + xy[:, 0].reshape(-1, 1, 1), 0 * xy[:, 0].reshape(-1, 1, 1)], dim=2)
a2 = torch.cat([xy[:, 1].reshape(-1, 1, 1), 1.0 + xy[:, 2].reshape(-1, 1, 1)], dim=2)
new_laf_no_center = torch.cat([a1, a2], dim=1).reshape(B, N, 2, 2)
new_laf = torch.cat([new_laf_no_center, laf[:, :, :, 2:3]], dim=3)
scale_orig = get_laf_scale(laf)
ellipse_scale = get_laf_scale(new_laf)
laf_out = scale_laf(make_upright(new_laf), scale_orig / ellipse_scale)
return laf_out
def extract_patches_from_pyramid(
img: torch.Tensor, laf: torch.Tensor, PS: int = 32,
normalize_lafs_before_extraction: bool = True
) -> torch.Tensor:
"""Extract patches defined by LAFs from image tensor.
Copied from kornia.feature.laf.extract_patches_from_pyramid with one minor
difference - highlighted below.
"""
raise_error_if_laf_is_not_valid(laf)
if normalize_lafs_before_extraction:
nlaf: torch.Tensor = normalize_laf(laf, img)
else:
nlaf = laf
B, N, _, _ = laf.size()
_, ch, h, w = img.size()
scale = 2.0 * get_laf_scale(denormalize_laf(nlaf, img)) / float(PS)
pyr_idx = scale.log2().relu().long() # diff: floor instead of round
cur_img = img
cur_pyr_level = 0
out = torch.zeros(B, N, ch, PS, PS).to(nlaf.dtype).to(nlaf.device)
while min(cur_img.size(2), cur_img.size(3)) >= PS:
_, ch, h, w = cur_img.size()
# for loop temporarily, to be refactored
for i in range(B):
scale_mask = (pyr_idx[i] == cur_pyr_level).squeeze()
if (scale_mask.float().sum()) == 0:
continue
scale_mask = (scale_mask > 0).view(-1)
grid = generate_patch_grid_from_normalized_LAF(
cur_img[i: i + 1], nlaf[i: i + 1, scale_mask, :, :], PS)
patches = F.grid_sample(
cur_img[i: i + 1].expand(grid.size(0), ch, h, w),
grid, # type: ignore
padding_mode="border",
align_corners=False,
)
out[i].masked_scatter_(scale_mask.view(-1, 1, 1, 1), patches)
cur_img = pyrdown(cur_img)
cur_pyr_level += 1
return out