|
|
"""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 |
|
|
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) |
|
|
|
