| """Model utilities for NLF.""" | |
| import torch | |
| def heatmap_to_image(coords: torch.Tensor, proc_side: int, stride: int, centered_stride: bool): | |
| last_image_pixel = proc_side - 1 | |
| last_receptive_center = last_image_pixel - (last_image_pixel % stride) | |
| coords_out = coords * last_receptive_center | |
| if centered_stride: | |
| coords_out = coords_out + stride // 2 | |
| return coords_out | |
| def heatmap_to_metric( | |
| coords: torch.Tensor, proc_side: int, stride: int, centered_stride: bool, box_size_m: float | |
| ): | |
| xy = coords[..., :2] | |
| coords2d = heatmap_to_image(xy, proc_side, stride, centered_stride) * box_size_m / proc_side | |
| return torch.cat([coords2d, coords[..., 2:] * box_size_m], dim=-1) | |