r"""grid_sample_2d for NCHW layout
(
data: np.ndarray,
grid: np.ndarray,
method="bilinear",
layout="NCHW",
padding_mode="zeros",
align_corners=True,
)
| 32 | |
| 33 | |
| 34 | def grid_sample_2d( |
| 35 | data: np.ndarray, |
| 36 | grid: np.ndarray, |
| 37 | method="bilinear", |
| 38 | layout="NCHW", |
| 39 | padding_mode="zeros", |
| 40 | align_corners=True, |
| 41 | ): |
| 42 | r"""grid_sample_2d for NCHW layout""" |
| 43 | |
| 44 | assert method in ("bilinear", "nearest", "bicubic"), f"{method} is not supported" |
| 45 | assert layout == "NCHW" |
| 46 | assert padding_mode in ("zeros", "border", "reflection"), f"{padding_mode} is not supported" |
| 47 | assert len(data.shape) == len(grid.shape) == 4 |
| 48 | |
| 49 | batch, channel = data.shape[:2] |
| 50 | in_height, in_width = data.shape[2:] |
| 51 | out_height, out_width = grid.shape[2:] |
| 52 | out_shape = [batch, channel, out_height, out_width] |
| 53 | out = np.zeros(out_shape) |
| 54 | |
| 55 | def _get_pixel(b, c, h, w): |
| 56 | if 0 <= h <= in_height - 1 and 0 <= w <= in_width - 1: |
| 57 | return data[b, c, h, w] |
| 58 | return 0 |
| 59 | |
| 60 | def _unnormalize(h, w): |
| 61 | if align_corners: |
| 62 | new_h = (h + 1) * (in_height - 1) / 2 |
| 63 | new_w = (w + 1) * (in_width - 1) / 2 |
| 64 | else: |
| 65 | new_h = -0.5 + (h + 1) * in_height / 2 |
| 66 | new_w = -0.5 + (w + 1) * in_width / 2 |
| 67 | return (new_h, new_w) |
| 68 | |
| 69 | def _clip_coordinates(x, size): |
| 70 | return min(max(x, 0), size - 1) |
| 71 | |
| 72 | def _reflect_coordinates(i, size): |
| 73 | def __refelection(i, size, corner_start): |
| 74 | def __reflect(index, size, corner_start): |
| 75 | index_align_corner = abs(corner_start - index) |
| 76 | size_times = index_align_corner // size |
| 77 | even = size_times % 2 == 0 |
| 78 | extra = index_align_corner - size_times * size |
| 79 | return extra + corner_start if even else size - extra + corner_start |
| 80 | |
| 81 | if corner_start <= i <= size + corner_start: |
| 82 | new_i = i |
| 83 | else: |
| 84 | new_i = __reflect(i, size, corner_start) |
| 85 | return new_i |
| 86 | |
| 87 | if align_corners: |
| 88 | x = __refelection(i, size - 1, 0) |
| 89 | else: |
| 90 | x = __refelection(i, size, -0.5) |
| 91 | return x |
nothing calls this directly
no test coverage detected
searching dependent graphs…