import numpy as np
import torch
import torch.nn.functional as F
import pytorch3d.transforms as transforms
from helper_functions.sampling_helper import sample_pixels_uniformly, sample_pixels_random, sample_valid_pixels_random
from helper_functions.geometry_helper import pose_compose
class RandomOptimizer():
def __init__(self, cfg, mipsfusion):
self.cfg = cfg
self.slam = mipsfusion
self.dataset = self.slam.dataset
self.device = self.slam.device
# parameters related to particle swarm template
# 6D pose format: [qx, qy, qz, tx, ty, tz], Tensor(6, )
self.particle_size = self.cfg["tracking"]["RO"]["particle_size"] # size of particle swarm template, default: 2000
self.scaling_coefficient1 = self.cfg["tracking"]["RO"]["initial_scaling_factor"] # initial scaling factor of each axis, default: 0.02
self.scaling_coefficient2 = self.cfg["tracking"]["RO"]["rescaling_factor"] # coefficient for update search size, default: 0.5
self.sdf_weight = 1000.
self.trunc_value = self.cfg["training"]["trunc"]
mean, cov = torch.zeros(6), torch.eye(6)
# PST
self.pre_sampled_particle = np.random.multivariate_normal(mean, cov, self.particle_size).astype(np.float32) # pre-sampled PST, ndarray(particle_size, 6)
self.pre_sampled_particle = torch.from_numpy(self.pre_sampled_particle).to(self.device)
self.pre_sampled_particle[0, :] = 0
self.pre_sampled_particle = torch.clamp(self.pre_sampled_particle, -2., 2.) # TEST
self.no_rel_trans = torch.tensor([1., 0., 0., 0., 0., 0., 0.]).to(self.device) # pose representing no transformation happened
# pre-sampled pixels
self.iW = self.cfg['tracking']['ignore_edge_W']
self.iH = self.cfg['tracking']['ignore_edge_H']
self.rays_dir = self.dataset.rays_d # dir vector of each pixel (in Camera Frame), Tensor(H, W, 3)
# self.n_rays = self.cfg["tracking"]["n_rays_h"] * self.cfg["tracking"]["n_rays_w"] # number of pixels sampled for RO
self.row_indices, self.col_indices = sample_pixels_uniformly(self.dataset.H, self.dataset.W,
self.cfg["tracking"]["RO"]["n_rows"], self.cfg["tracking"]["RO"]["n_cols"])
# camera parameters
self.fx, self.fy, self.cx, self.cy = self.dataset.fx, self.dataset.fy, self.dataset.cx, self.dataset.cy
self.intrinsic = torch.tensor( [ [self.fx, 0., self.cx],
[ 0., self.fy, self.cy],
[ 0., 0., 1.] ] ).to(self.device)
# @brief: convert a batch of 6D poses to 7D poses(quaternion + translation vector);
# @param batch_pose: N 6D poses: [qx, qy, qz, tx, ty, tz], Tensor(N, 6);
#-@return: N 7D pose: [qw, qx, qy, qz, tx, ty, tz], Tensor(N, 7).
def pose_6D_to_7D(self, batch_pose):
imag_sq_sum = batch_pose[:, 0] ** 2 + batch_pose[:, 1] ** 2 + batch_pose[:, 2] ** 2 # Tensor(N, )
sample_qw = torch.where(imag_sq_sum <= 1., torch.sqrt(1 - imag_sq_sum), 0.).unsqueeze(1) # Tensor(N, 1)
batch_pose_7D = torch.cat([sample_qw, batch_pose], dim=-1) # 7D pose(real part first), Tensor(N, 7)
# batch_pose_7D[:, :4] /= batch_pose_7D[:, :4].norm(dim=-1, keepdim=True) # ensure union quaternions
return batch_pose_7D
# convert poses of each particle(relative pose) to absolute pose;
# @param ref_pose_rot: Tensor(3, 3);
# @param ref_pose_trans: Tensor(3, 1);
# @param particle_template: Tensor(N, 7);
#-@return abs_rot: absolute pose of each particle(rotation), Tensor(N, 3, 3);
#-@return abs_trans: absolute pose of each particle(rotation), Tensor(N, 3, 1).
def get_abs_pose(self, ref_pose_rot, ref_pose_trans, particle_template):
delta_R = transforms.quaternion_to_matrix(particle_template[:, :4])
abs_rot = ref_pose_rot @ delta_R # Tensor(N, 3, 3)
abs_trans = ref_pose_trans + particle_template[:, 4:, None] # Tensor(N, 3, 1)
return abs_rot, abs_trans
# @brief: transform a bunch of 3D points using a given transformation;
# @param points: Tensor(m, 3);
# @param pose_rot:Tensor(N, 3, 3);
# @param pose_trans: Tensor(N, 3, 1);
#-@return: Tensor(N, m, 3).
def batch_points_trans(self, points, pose_rot, pose_trans):
transed_pts = pose_rot @ torch.transpose(points, 0, 1)
transed_pts = transed_pts + pose_trans
transed_pts = torch.transpose(transed_pts, 1, 2)
return transed_pts
# @brief: judging whether N pixels are out of range;
def get_range_mask(self, pixel_coords, img_h, img_w):
x_coords = pixel_coords[:, 0] # Tensor(n, )
y_coords = pixel_coords[:, 1] # Tensor(n, )
x_mask1 = torch.where(x_coords > 0., torch.ones_like(x_coords), torch.zeros_like(x_coords))
x_mask2 = torch.where(x_coords < img_w, torch.ones_like(x_coords), torch.zeros_like(x_coords))
x_mask = x_mask1 * x_mask2
y_mask1 = torch.where(y_coords > 0., torch.ones_like(y_coords), torch.zeros_like(y_coords))
y_mask2 = torch.where(y_coords < img_h, torch.ones_like(y_coords), torch.zeros_like(y_coords))
y_mask = y_mask1 * y_mask2
final_mask = x_mask * y_mask
return final_mask
# @brief: evaluate fitness value of each particle;
# param abs_rot: Tensor(N, 3, 3);
# param abs_trans: Tensor(N, 3, 1);
# @param last_frame_pose: tracked pose of last frame(T_wl), Tensor(4, 4);
# @param target_d: depth of selected pixels, Tensor(N, 1);
# @param rays_d_cam: ray dir of selected pixels, Tensor(N, 3);
#-@return fitness_value: fitness value of each particle, Tensor(N, );
#-@return mean_masked_sdf: transformed by each particle(candidate pose),mean predicted SDF, Tensor(N, ).
def get_fitness(self, model, abs_rot, abs_trans, last_frame_pose, target_d, rays_d_cam):
# Step 1: pixel coordinates --> camera coordinates
# 1.1: compute corresponding 3D points in Camera Coordinate System
cam_coords = rays_d_cam * target_d # Tensor(N, 3)
# 1.2: compute depth mask and overlapping mask
valid_mask = torch.where(target_d > 0., torch.ones_like(target_d), torch.zeros_like(target_d)).squeeze(-1)[None, ...] # Tensor(N, 1)
# Step 2: for each particle (candidate pose), use it to transform sampled 3D points to World coordinates
world_coords = self.batch_points_trans(cam_coords, abs_rot, abs_trans) # Tensor(N, n_rays, 3)
# Step 3: infer these N * n_rays 3D points, get predicted SDF values (metrics: m)
pred_sdf = model.run_network(world_coords)[..., 3:4].squeeze(-1) * self.trunc_value # Tensor(N, n_rays)
# Step 4: for each candidate pose, compute mean predicted SDF value of all valid pixels
mean_masked_sdf = torch.mean(valid_mask * torch.abs(pred_sdf), dim=-1) # Tensor(N, )
fitness_value = mean_masked_sdf * self.sdf_weight # fitness value of each particle, Tensor(N, )
return fitness_value, mean_masked_sdf
# @brief: update pose (starting pose of next round);
# @param success_flag: whether get non-empty APS, Tensor(, ), dtype=bool;
# @param rot_cur: Tensor(3, 3);
# @param trans_cur: Tensor(3, 1);
# @param mean_transform: weighted mean of APS, Tensor(7, );
#-@return rot_updated: Tensor(3, 3);
#-@return trans_updated: Tensor(3, 1).
def update_cur_pose(self, rot_cur, trans_cur, mean_transform):
delta_R = transforms.quaternion_to_matrix(mean_transform[:4]) # optimal delta_R of this round, Tensor(3, 3)
delta_t = mean_transform[4:][..., None] # optimal delta_t of this round, Tensor(3, 1)
rot_updated = rot_cur @ delta_R # R_u = dR @ R_c, Tensor(3, 3)
trans_updated = trans_cur + delta_t # t_u = t_c + dt, Tensor(3, 1)
return rot_updated, trans_updated
# @brief: according to searching result of this round, update search_size(if APS is empty, mean_transform is all 0);
# @param mean_pred_sdf: Tensor(, );
# @param mean_transform_quat: weighted mean of APS( [qx, qy, qz, tx, ty, tz] ), Tensor(6, );
# -@return: Tensor(1, 6).
def update_search_size(self, mean_pred_sdf, mean_transform):
s = torch.abs(mean_transform) + 0.0001 # Tensor(6, )
search_size = self.scaling_coefficient2 * mean_pred_sdf * s / s.norm() + 0.0001
return search_size[None, ...]
# @param depth_img: gt depth image of current frame, Tensor(h, w);
# @param last_frame_pose: tracked pose of last frame(c2w, Local Coordinate System), Tensor(4, 4);
# @param n_iter: iter num;
#-@return: Tensor(4, 4).
@torch.no_grad()
def optimize(self, model, depth_img, initial_pose, last_frame_pose, n_iter=10):
if n_iter <= 0:
return initial_pose
rot_cur, trans_cur = initial_pose[:3, :3], initial_pose[:3, 3:] # Tensor(3, 3) / Tensor(3, 1)
search_size = self.scaling_coefficient1
# # pixel sampling
# # (1) random sampling
# indice = sample_valid_pixels_random(depth_img[self.iH: -self.iH, self.iW: -self.iW], self.cfg["tracking"]["RO"]["pixel_num"])
# indice_h, indice_w = torch.remainder(indice, self.dataset.H - self.iH * 2), torch.div(indice, self.dataset.H - self.iH * 2, rounding_mode="floor")
# target_d = depth_img[self.iH: -self.iH, self.iW: -self.iW][indice_h, indice_w].to(self.device).unsqueeze(-1) # Tensor(pixel_num, 1)
# rays_d_cam = self.rays_dir[self.iH: -self.iH, self.iW: -self.iW][indice_h, indice_w, :].to(self.device) # Tensor(pixel_num, 3)
# (2) uniform sampling
indice_h, indice_w = self.row_indices, self.col_indices
target_d = depth_img[indice_h, indice_w].to(self.device).unsqueeze(-1) # Tensor(pixel_num, 1)
rays_d_cam = self.rays_dir[indice_h, indice_w, :].to(self.device) # Tensor(pixel_num, 3)
for i in range(n_iter):
offset = i % 5
indice_h, indice_w = self.row_indices + offset, self.col_indices + offset
target_d = depth_img[indice_h, indice_w].to(self.device).unsqueeze(-1) # Tensor(pixel_num, 1)
rays_d_cam = self.rays_dir[indice_h, indice_w, :].to(self.device) # Tensor(pixel_num, 3)
# Step 1: recover absolute pose for each particle(pose) in template
# Step 1.1: get delta pose from pre-sampled particles
rescaled_pst = self.pre_sampled_particle * search_size # Tensor(N, 6)
rescaled_pst_7D = self.pose_6D_to_7D(rescaled_pst) # Tensor(N, 7)
# Step 1.2: recover to absolute pose
abs_rot_pst, abs_trans_pst = self.get_abs_pose(rot_cur, trans_cur, rescaled_pst_7D) # Tensor(N, 3, 3) / Tensor(N, 3, 1)
# Step 2: *** evaluate (compute fitness value) each particle
fitness_values, pred_mean_sdf = self.get_fitness(model, abs_rot_pst, abs_trans_pst, last_frame_pose, target_d, rays_d_cam) # Tensor(N, ) / Tensor(N, )
# Step 3: filter advanced particle swarm (APS)
original_fitness = fitness_values[0]
better_mask = torch.where(fitness_values < original_fitness, torch.ones_like(original_fitness), torch.zeros_like(original_fitness)) # Tensor(N, )
weights = (original_fitness - fitness_values) * better_mask # weight of each particle(with mask, 0 for non-advanced particle), Tensor(N, )
weight_sum = torch.sum(weights) + 0.00001 # Tensor(, )
success_flag = ( torch.count_nonzero(better_mask) > 0 )
if success_flag:
mean_sdf = torch.sum(weights * pred_mean_sdf) / weight_sum # mean pred SDF of APS, Tensor(, )
else:
mean_sdf = pred_mean_sdf[0]
# Step 4: update R, t
if success_flag:
mean_transform = torch.sum(rescaled_pst_7D * weights[:, None], dim=0) / weight_sum # weighted mean of APS, Tensor(7, )
mean_transform_quat = mean_transform[:4] / (mean_transform[:4].norm() + 1e-5) # [qw, qx, qy, qz], Tensor(4, )
mean_transform = torch.cat([mean_transform_quat, mean_transform[4:]], dim=0) # weighted mean of APS, Tensor(7, )
rot_cur, trans_cur = self.update_cur_pose(rot_cur, trans_cur, mean_transform) # update current pose(starting point of next round) Tensor(3, 3) / Tensor(3, 1)
else:
mean_transform = self.no_rel_trans
# Step 5: rescaling particle swarm template (update search_size)
search_size_temp = self.update_search_size(mean_sdf, mean_transform[1:]) # Tensor(1, 6)
search_size = torch.where(success_flag, search_size_temp, search_size_temp * 2).to(self.device)
tracked_pose = pose_compose(rot_cur, trans_cur) # Tensor(4, 4)
return tracked_pose