Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,744 Bytes
c295391 |
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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import torch
import torch.nn.functional as F
def activate_pose(pred_pose_enc, trans_act="linear", quat_act="linear", fl_act="linear"):
"""
Activate pose parameters with specified activation functions.
Args:
pred_pose_enc: Tensor containing encoded pose parameters [translation, quaternion, focal length]
trans_act: Activation type for translation component
quat_act: Activation type for quaternion component
fl_act: Activation type for focal length component
Returns:
Activated pose parameters tensor
"""
T = pred_pose_enc[..., :3]
quat = pred_pose_enc[..., 3:7]
fl = pred_pose_enc[..., 7:] # or fov
T = base_pose_act(T, trans_act)
quat = base_pose_act(quat, quat_act)
fl = base_pose_act(fl, fl_act) # or fov
pred_pose_enc = torch.cat([T, quat, fl], dim=-1)
return pred_pose_enc
def base_pose_act(pose_enc, act_type="linear"):
"""
Apply basic activation function to pose parameters.
Args:
pose_enc: Tensor containing encoded pose parameters
act_type: Activation type ("linear", "inv_log", "exp", "relu")
Returns:
Activated pose parameters
"""
if act_type == "linear":
return pose_enc
elif act_type == "inv_log":
return inverse_log_transform(pose_enc)
elif act_type == "exp":
return torch.exp(pose_enc)
elif act_type == "relu":
return F.relu(pose_enc)
else:
raise ValueError(f"Unknown act_type: {act_type}")
def activate_head(out, activation="norm_exp", conf_activation="expp1"):
"""
Process network output to extract 3D points and confidence values.
Args:
out: Network output tensor (B, C, H, W)
activation: Activation type for 3D points
conf_activation: Activation type for confidence values
Returns:
Tuple of (3D points tensor, confidence tensor)
"""
# Move channels from last dim to the 4th dimension => (B, H, W, C)
fmap = out.permute(0, 2, 3, 1) # B,H,W,C expected
# Split into xyz (first C-1 channels) and confidence (last channel)
xyz = fmap[:, :, :, :-1]
conf = fmap[:, :, :, -1]
if activation == "norm_exp":
d = xyz.norm(dim=-1, keepdim=True).clamp(min=1e-8)
xyz_normed = xyz / d
pts3d = xyz_normed * torch.expm1(d)
elif activation == "norm":
pts3d = xyz / xyz.norm(dim=-1, keepdim=True)
elif activation == "exp":
pts3d = torch.exp(xyz)
elif activation == "relu":
pts3d = F.relu(xyz)
elif activation == "inv_log":
pts3d = inverse_log_transform(xyz)
elif activation == "xy_inv_log":
xy, z = xyz.split([2, 1], dim=-1)
z = inverse_log_transform(z)
pts3d = torch.cat([xy * z, z], dim=-1)
elif activation == "sigmoid":
pts3d = torch.sigmoid(xyz)
elif activation == "linear":
pts3d = xyz
else:
raise ValueError(f"Unknown activation: {activation}")
if conf_activation == "expp1":
conf_out = 1 + conf.exp()
elif conf_activation == "expp0":
conf_out = conf.exp()
elif conf_activation == "sigmoid":
conf_out = torch.sigmoid(conf)
else:
raise ValueError(f"Unknown conf_activation: {conf_activation}")
return pts3d, conf_out
def inverse_log_transform(y):
"""
Apply inverse log transform: sign(y) * (exp(|y|) - 1)
Args:
y: Input tensor
Returns:
Transformed tensor
"""
return torch.sign(y) * (torch.expm1(torch.abs(y)))
|