Spaces:
Running on Zero
Running on Zero
| """Perspective field utilities. | |
| Adapted from https://github.com/jinlinyi/PerspectiveFields | |
| """ | |
| import torch | |
| from siclib.utils.conversions import deg2rad, rad2deg | |
| def encode_up_bin(vector_field: torch.Tensor, num_bin: int) -> torch.Tensor: | |
| """Encode vector field into classification bins. | |
| Args: | |
| vector_field (torch.Tensor): gravity field of shape (2, h, w), with channel 0 cos(theta) and | |
| 1 sin(theta) | |
| num_bin (int): number of classification bins | |
| Returns: | |
| torch.Tensor: encoded bin indices of shape (1, h, w) | |
| """ | |
| angle = ( | |
| torch.atan2(vector_field[1, :, :], vector_field[0, :, :]) / torch.pi * 180 + 180 | |
| ) % 360 # [0,360) | |
| angle_bin = torch.round(torch.div(angle, (360 / (num_bin - 1)))).long() | |
| angle_bin[angle_bin == num_bin - 1] = 0 | |
| invalid = (vector_field == 0).sum(0) == vector_field.size(0) | |
| angle_bin[invalid] = num_bin - 1 | |
| return deg2rad(angle_bin.type(torch.LongTensor)) | |
| def decode_up_bin(angle_bin: torch.Tensor, num_bin: int) -> torch.Tensor: | |
| """Decode classification bins into vector field. | |
| Args: | |
| angle_bin (torch.Tensor): bin indices of shape (1, h, w) | |
| num_bin (int): number of classification bins | |
| Returns: | |
| torch.Tensor: decoded vector field of shape (2, h, w) | |
| """ | |
| angle = (angle_bin * (360 / (num_bin - 1)) - 180) / 180 * torch.pi | |
| cos = torch.cos(angle) | |
| sin = torch.sin(angle) | |
| vector_field = torch.stack((cos, sin), dim=1) | |
| invalid = angle_bin == num_bin - 1 | |
| invalid = invalid.unsqueeze(1).repeat(1, 2, 1, 1) | |
| vector_field[invalid] = 0 | |
| return vector_field | |
| def encode_bin_latitude(latimap: torch.Tensor, num_classes: int) -> torch.Tensor: | |
| """Encode latitude map into classification bins. | |
| Args: | |
| latimap (torch.Tensor): latitude map of shape (h, w) with values in [-90, 90] | |
| num_classes (int): number of classes | |
| Returns: | |
| torch.Tensor: encoded latitude bin indices | |
| """ | |
| boundaries = torch.arange(-90, 90, 180 / num_classes)[1:] | |
| binmap = torch.bucketize(rad2deg(latimap), boundaries) | |
| return binmap.type(torch.LongTensor) | |
| def decode_bin_latitude(binmap: torch.Tensor, num_classes: int) -> torch.Tensor: | |
| """Decode classification bins to latitude map. | |
| Args: | |
| binmap (torch.Tensor): encoded classification bins | |
| num_classes (int): number of classes | |
| Returns: | |
| torch.Tensor: latitude map of shape (h, w) | |
| """ | |
| bin_size = 180 / num_classes | |
| bin_centers = torch.arange(-90, 90, bin_size) + bin_size / 2 | |
| bin_centers = bin_centers.to(binmap.device) | |
| latimap = bin_centers[binmap] | |
| return deg2rad(latimap) | |