Spaces:
Runtime error
Runtime error
File size: 3,440 Bytes
5a256aa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import torch
import torch.nn as nn
import torch.nn.functional as F
def compute_tri_normal(geometry, tris):
tri_1 = tris[:, 0]
tri_2 = tris[:, 1]
tri_3 = tris[:, 2]
vert_1 = torch.index_select(geometry, 1, tri_1)
vert_2 = torch.index_select(geometry, 1, tri_2)
vert_3 = torch.index_select(geometry, 1, tri_3)
nnorm = torch.cross(vert_2 - vert_1, vert_3 - vert_1, 2)
normal = nn.functional.normalize(nnorm)
return normal
def euler2rot(euler_angle):
batch_size = euler_angle.shape[0]
theta = euler_angle[:, 0].reshape(-1, 1, 1)
phi = euler_angle[:, 1].reshape(-1, 1, 1)
psi = euler_angle[:, 2].reshape(-1, 1, 1)
one = torch.ones(batch_size, 1, 1).to(euler_angle.device)
zero = torch.zeros(batch_size, 1, 1).to(euler_angle.device)
rot_x = torch.cat(
(
torch.cat((one, zero, zero), 1),
torch.cat((zero, theta.cos(), theta.sin()), 1),
torch.cat((zero, -theta.sin(), theta.cos()), 1),
),
2,
)
rot_y = torch.cat(
(
torch.cat((phi.cos(), zero, -phi.sin()), 1),
torch.cat((zero, one, zero), 1),
torch.cat((phi.sin(), zero, phi.cos()), 1),
),
2,
)
rot_z = torch.cat(
(
torch.cat((psi.cos(), -psi.sin(), zero), 1),
torch.cat((psi.sin(), psi.cos(), zero), 1),
torch.cat((zero, zero, one), 1),
),
2,
)
return torch.bmm(rot_x, torch.bmm(rot_y, rot_z))
def rot_trans_pts(geometry, rot, trans):
rott_geo = torch.bmm(rot, geometry.permute(0, 2, 1)) + trans[:, :, None]
return rott_geo.permute(0, 2, 1)
def cal_lap_loss(tensor_list, weight_list):
lap_kernel = (
torch.Tensor((-0.5, 1.0, -0.5))
.unsqueeze(0)
.unsqueeze(0)
.float()
.to(tensor_list[0].device)
)
loss_lap = 0
for i in range(len(tensor_list)):
in_tensor = tensor_list[i]
in_tensor = in_tensor.view(-1, 1, in_tensor.shape[-1])
out_tensor = F.conv1d(in_tensor, lap_kernel)
loss_lap += torch.mean(out_tensor ** 2) * weight_list[i]
return loss_lap
def proj_pts(rott_geo, focal_length, cxy):
cx, cy = cxy[0], cxy[1]
X = rott_geo[:, :, 0]
Y = rott_geo[:, :, 1]
Z = rott_geo[:, :, 2]
fxX = focal_length * X
fyY = focal_length * Y
proj_x = -fxX / Z + cx
proj_y = fyY / Z + cy
return torch.cat((proj_x[:, :, None], proj_y[:, :, None], Z[:, :, None]), 2)
def forward_rott(geometry, euler_angle, trans):
rot = euler2rot(euler_angle)
rott_geo = rot_trans_pts(geometry, rot, trans)
return rott_geo
def forward_transform(geometry, euler_angle, trans, focal_length, cxy):
rot = euler2rot(euler_angle)
rott_geo = rot_trans_pts(geometry, rot, trans)
proj_geo = proj_pts(rott_geo, focal_length, cxy)
return proj_geo
def cal_lan_loss(proj_lan, gt_lan):
return torch.mean((proj_lan - gt_lan) ** 2)
def cal_col_loss(pred_img, gt_img, img_mask):
pred_img = pred_img.float()
# loss = torch.sqrt(torch.sum(torch.square(pred_img - gt_img), 3))*img_mask/255
loss = (torch.sum(torch.square(pred_img - gt_img), 3)) * img_mask / 255
loss = torch.sum(loss, dim=(1, 2)) / torch.sum(img_mask, dim=(1, 2))
loss = torch.mean(loss)
return loss
|