Generate uv coordinates ([0, w), [0, h)) of the patches centered at patch_center. Args: patch_center: (*b, 2) the center of each patch on arr. u: first dim [0, w_in], v: second dim [0, h_in] patch_width_px: number of pixels in the patch in width
(
patch_center: torch.Tensor, # (*b, 2) h, w
patch_width_px: int,
patch_width_pitch_scale: T.Union[float, torch.Tensor] = 1., # (*b,)
patch_height_px: int = None, # (*b,)
patch_height_pitch_scale: T.Union[float, torch.Tensor] = None, # (*b,)
device: torch.device = torch.device('cpu'),
)
| 2812 | |
| 2813 | |
| 2814 | def generate_patch_uv( |
| 2815 | patch_center: torch.Tensor, # (*b, 2) h, w |
| 2816 | patch_width_px: int, |
| 2817 | patch_width_pitch_scale: T.Union[float, torch.Tensor] = 1., # (*b,) |
| 2818 | patch_height_px: int = None, # (*b,) |
| 2819 | patch_height_pitch_scale: T.Union[float, torch.Tensor] = None, # (*b,) |
| 2820 | device: torch.device = torch.device('cpu'), |
| 2821 | ) -> torch.Tensor: |
| 2822 | """ |
| 2823 | Generate uv coordinates ([0, w), [0, h)) of the patches centered at patch_center. |
| 2824 | |
| 2825 | Args: |
| 2826 | patch_center: |
| 2827 | (*b, 2) the center of each patch on arr. u: first dim [0, w_in], v: second dim [0, h_in] |
| 2828 | patch_width_px: |
| 2829 | number of pixels in the patch in width |
| 2830 | patch_width_pitch_scale: |
| 2831 | (*b,) the pitch of the patch (new_pitch / old_pitch) |
| 2832 | patch_height_px: |
| 2833 | if None, the same as `patch_width_px` |
| 2834 | patch_height_pitch_scale: |
| 2835 | if None, the same as `patch_width_pitch_scale` |
| 2836 | int_only: |
| 2837 | whether the center is always at an integer index |
| 2838 | |
| 2839 | Returns: |
| 2840 | (*b, patch_height_px, patch_width_px, 2) the u (first dimension) and v (second dimension) |
| 2841 | Note the returned uv can go out of bound. |
| 2842 | """ |
| 2843 | |
| 2844 | *b_shape, _2 = patch_center.shape |
| 2845 | b = np.prod(b_shape) |
| 2846 | patch_center = patch_center.reshape(b, 2) # (b, 2) |
| 2847 | if isinstance(patch_width_pitch_scale, (int, float)): |
| 2848 | patch_width_pitch_scale = torch.ones(b, dtype=torch.float, device=device) * patch_width_pitch_scale |
| 2849 | if isinstance(patch_width_pitch_scale, torch.Tensor): |
| 2850 | patch_width_pitch_scale = patch_width_pitch_scale.reshape(b).to(device=device) # (b,) |
| 2851 | |
| 2852 | if patch_height_px is None: |
| 2853 | patch_height_px = patch_width_px |
| 2854 | |
| 2855 | if patch_height_pitch_scale is None: |
| 2856 | patch_height_pitch_scale = patch_width_pitch_scale |
| 2857 | if isinstance(patch_height_pitch_scale, (int, float)): |
| 2858 | patch_height_pitch_scale = torch.ones(b, dtype=torch.float, device=device) * patch_height_pitch_scale |
| 2859 | if isinstance(patch_height_pitch_scale, torch.Tensor): |
| 2860 | patch_height_pitch_scale = patch_height_pitch_scale.reshape(b).to(device=device) # (b,) |
| 2861 | |
| 2862 | # generate the canonical grid for the patch |
| 2863 | patch_half_width_px = patch_width_px / 2 |
| 2864 | patch_half_height_px = patch_height_px / 2 |
| 2865 | |
| 2866 | u, v = torch.meshgrid( |
| 2867 | torch.arange(patch_width_px, dtype=torch.float, device=device), |
| 2868 | torch.arange(patch_height_px, dtype=torch.float, device=device), |
| 2869 | indexing='xy', |
| 2870 | ) # u: (hp, wp), [0, w-1], v: (hp, wp) [0, h-1] top-left (0,0) |
| 2871 | u = u + (0.5 - patch_half_width_px) |
no test coverage detected