def test_shape(self): laf = torch.rand(5, 3, 2, 3) img = torch.rand(5, 3, 10, 10) PS = 3 from kornia.feature.laf import generate_patch_grid_from_normalized_LAF grid = generate_patch_grid_from_normalized_LAF(img, laf, PS) assert grid.shape == (15, 3, 3, 2)
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