(B, Y, X, stack=False, norm=False, device="cuda")
| 1069 | |
| 1070 | |
| 1071 | def meshgrid2d(B, Y, X, stack=False, norm=False, device="cuda"): |
| 1072 | # returns a meshgrid sized B x Y x X |
| 1073 | |
| 1074 | grid_y = torch.linspace(0.0, Y - 1, Y, device=torch.device(device)) |
| 1075 | grid_y = torch.reshape(grid_y, [1, Y, 1]) |
| 1076 | grid_y = grid_y.repeat(B, 1, X) |
| 1077 | |
| 1078 | grid_x = torch.linspace(0.0, X - 1, X, device=torch.device(device)) |
| 1079 | grid_x = torch.reshape(grid_x, [1, 1, X]) |
| 1080 | grid_x = grid_x.repeat(B, Y, 1) |
| 1081 | |
| 1082 | if stack: |
| 1083 | # note we stack in xy order |
| 1084 | # (see https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.grid_sample) |
| 1085 | grid = torch.stack([grid_x, grid_y], dim=-1) |
| 1086 | return grid |
| 1087 | else: |
| 1088 | return grid_y, grid_x |
| 1089 | |
| 1090 | def get_points_on_a_grid(grid_size, interp_shape, |
| 1091 | grid_center=(0, 0), device="cuda"): |
no test coverage detected