Sample from `arr` a patch centered at `center` with a different pixel pitch and number of pixels. Args: arr: (*b, c, h_in, w_in) or (*b, h_in, w_in, c), see `format`. the array to be sampled from. patch_center: (*b, 2) the center of each patch on arr
(
arr: torch.Tensor, # (*b, c, h_in, w_in)
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,)
mode: str = 'bilinear',
padding_mode: str = 'zeros',
format: str = 'chw', # 'hwc'
)
| 2722 | |
| 2723 | |
| 2724 | def sample_patch( |
| 2725 | arr: torch.Tensor, # (*b, c, h_in, w_in) |
| 2726 | patch_center: torch.Tensor, # (*b, 2) h, w |
| 2727 | patch_width_px: int, |
| 2728 | patch_width_pitch_scale: T.Union[float, torch.Tensor] = 1., # (*b,) |
| 2729 | patch_height_px: int = None, # (*b,) |
| 2730 | patch_height_pitch_scale: T.Union[float, torch.Tensor] = None, # (*b,) |
| 2731 | mode: str = 'bilinear', |
| 2732 | padding_mode: str = 'zeros', |
| 2733 | format: str = 'chw', # 'hwc' |
| 2734 | ): |
| 2735 | """ |
| 2736 | Sample from `arr` a patch centered at `center` with a different pixel pitch and number of pixels. |
| 2737 | |
| 2738 | Args: |
| 2739 | arr: |
| 2740 | (*b, c, h_in, w_in) or (*b, h_in, w_in, c), see `format`. the array to be sampled from. |
| 2741 | patch_center: |
| 2742 | (*b, 2) the center of each patch on arr. u: first dim [0, w_in], v: second dim [0, h_in] |
| 2743 | patch_width_px: |
| 2744 | number of pixels in the patch in width |
| 2745 | patch_width_pitch_scale: |
| 2746 | (*b,) the pitch of the patch (new_pitch / old_pitch) |
| 2747 | patch_height_px: |
| 2748 | if None, the same as `patch_width_px` |
| 2749 | patch_height_pitch_scale: |
| 2750 | if None, the same as `patch_width_pitch_scale` |
| 2751 | format: |
| 2752 | 'chw': arr is (b, c, h, w) |
| 2753 | 'hwc': arr is (b, h, w, c) |
| 2754 | |
| 2755 | Returns: |
| 2756 | (*b, c, patch_height_px, patch_width_px) or (*b, patch_height_px, patch_width_px, c) |
| 2757 | |
| 2758 | Note: |
| 2759 | coordinate system: |
| 2760 | The origin of the coordinate is at the top-left corner of `arr`. |
| 2761 | Each pixel in `arr` is 1 unit in width and height. |
| 2762 | The first dimension (u) is toward right and second dimension (v) is toward down. |
| 2763 | The first pixel center is `arr` is at (0.5, 0.5). |
| 2764 | This function should be compared with `uv_sampling`, which uses a different coordinate system. |
| 2765 | """ |
| 2766 | |
| 2767 | if format == 'chw': |
| 2768 | *b_shape, c, h, w = arr.shape |
| 2769 | arr = arr.reshape(-1, c, h, w) # (b, c, h, w) |
| 2770 | elif format == 'hwc': |
| 2771 | *b_shape, h, w, c = arr.shape |
| 2772 | arr = arr.reshape(-1, h, w, c).permute(0, 3, 1, 2) # (b, c, h, w) |
| 2773 | else: |
| 2774 | raise NotImplementedError |
| 2775 | |
| 2776 | b = np.prod(b_shape) |
| 2777 | device = arr.device |
| 2778 | |
| 2779 | uv = generate_patch_uv( |
| 2780 | patch_center=patch_center, # (*b, 2) |
| 2781 | patch_width_px=patch_width_px, |
nothing calls this directly
no test coverage detected