Spaces:
Running
Running
| import numpy as np | |
| class Heatmaps: | |
| def __init__(self, num_maps, map_size, sigma, stride=1, norm=True): | |
| self.num_maps = num_maps | |
| self.sigma = sigma | |
| self.double_sigma_pw2 = 2*sigma*sigma | |
| self.doublepi_sigma_pw2 = self.double_sigma_pw2 * np.pi | |
| self.stride = stride | |
| self.norm = norm | |
| if isinstance(map_size, (tuple, list)): | |
| self.width = map_size[0] | |
| self.height = map_size[1] | |
| else: | |
| self.width = map_size | |
| self.height = map_size | |
| grid_x = np.arange(self.width) * stride + stride / 2 - 0.5 | |
| self.grid_x = np.repeat(grid_x.reshape(1, self.width), self.num_maps, axis=0) | |
| grid_y = np.arange(self.height) * stride + stride / 2 - 0.5 | |
| self.grid_y = np.repeat(grid_y.reshape(1, self.height), self.num_maps, axis=0) | |
| def __call__(self, sample): | |
| landmarks = sample['landmarks'] | |
| landmarks = landmarks[-self.num_maps:] | |
| # Heatmap generation | |
| exp_x = np.exp(-(self.grid_x - landmarks[:, 0].reshape(-1, 1)) ** 2 / self.double_sigma_pw2) | |
| exp_y = np.exp(-(self.grid_y - landmarks[:, 1].reshape(-1, 1)) ** 2 / self.double_sigma_pw2) | |
| heatmaps = np.matmul(exp_y.reshape(self.num_maps, self.height, 1), exp_x.reshape(self.num_maps, 1, self.width)) | |
| if self.norm: | |
| heatmaps = heatmaps/self.doublepi_sigma_pw2 | |
| sample['heatmap2D'] = heatmaps | |
| return sample | |