def get_shape_chamfer_loss(self, pts, quat1, quat2, center1, center2,
part_cnt):
num_point = pts.shape[1]
part_pcs1 = qrot(quat1.unsqueeze(1).repeat(1, num_point, 1),
pts) + center1.unsqueeze(1).repeat(1, num_point, 1)
part_pcs2 = qrot(quat2.unsqueeze(1).repeat(1, num_point, 1),
pts) + center2.unsqueeze(1).repeat(1, num_point, 1)
t = 0
shape_pcs1 = []
shape_pcs2 = []
for cnt in part_cnt:
cur_shape_pc1 = part_pcs1[t:t + cnt].view(1, -1, 3)
cur_shape_pc2 = part_pcs2[t:t + cnt].view(1, -1, 3)
with torch.no_grad():
idx1 = furthest_point_sample(cur_shape_pc1, 2048).long()[0]
idx2 = furthest_point_sample(cur_shape_pc2, 2048).long()[0]
shape_pcs1.append(cur_shape_pc1[:, idx1])
shape_pcs2.append(cur_shape_pc2[:, idx2])
t += cnt
shape_pcs1 = torch.cat(shape_pcs1, dim=0) # numshapes x 2048 x 3
shape_pcs2 = torch.cat(shape_pcs2, dim=0) # numshapes x 2048 x 3
dist1, dist2 = chamfer_distance(shape_pcs1,
shape_pcs2,
transpose=False)
loss_per_data = torch.mean(dist1, dim=1) + torch.mean(dist2, dim=1)
return loss_per_data
def get_quat_loss(self, pts, quat1, quat2):
num_point = pts.shape[1]
pts1 = qrot(quat1.unsqueeze(1).repeat(1, num_point, 1), pts)
pts2 = qrot(quat2.unsqueeze(1).repeat(1, num_point, 1), pts)
dist1, dist2 = chamfer_distance(pts1, pts2, transpose=False)
loss_per_data = torch.mean(dist1, dim=1) + torch.mean(dist2, dim=1)
return loss_per_data
def get_l2_rotation_loss(self, pts, quat1, quat2):
num_point = pts.shape[1]
pts1 = qrot(quat1.unsqueeze(1).repeat(1, num_point, 1), pts)
pts2 = qrot(quat2.unsqueeze(1).repeat(1, num_point, 1), pts)
loss_per_data = (pts1 - pts2).pow(2).sum(dim=2).mean(dim=1)
return loss_per_data
def forward_kinematics(self, rotations, root_positions):
"""
Perform forward kinematics using the given trajectory and local rotations.
Arguments (where N = batch size, L = sequence length, J = number of joints):
-- rotations: (N, L, J, 4) tensor of unit quaternions describing the local rotations of each joint.
-- root_positions: (N, L, 3) tensor describing the root joint positions.
"""
assert len(rotations.shape) == 4
assert rotations.shape[-1] == 4
positions_world = []
rotations_world = []
expanded_offsets = self._offsets.expand(rotations.shape[0],
rotations.shape[1],
self._offsets.shape[0],
self._offsets.shape[1])
# Parallelize along the batch and time dimensions
for i in range(self._offsets.shape[0]):
if self._parents[i] == -1:
positions_world.append(root_positions)
rotations_world.append(rotations[:, :, 0])
else:
positions_world.append(qrot(rotations_world[self._parents[i]], expanded_offsets[:, :, i]) \
+ positions_world[self._parents[i]])
if self._has_children[i]:
rotations_world.append(
qmul(rotations_world[self._parents[i]], rotations[:, :,
i]))
else:
# This joint is a terminal node -> it would be useless to compute the transformation
rotations_world.append(None)
return torch.stack(positions_world, dim=3).permute(0, 1, 3, 2)
def render_boxes(out_fn, boxes):
boxes = torch.from_numpy(boxes)
cmd = 'cp %s %s' % (os.path.join(BASE_DIR, 'cube.mtl'), out_fn + '.mtl')
call(cmd, shell=True)
with open(out_fn + '.obj', 'w') as fout:
fout.write('mtllib %s\n' % (out_fn.split('/')[-1] + '.mtl'))
for j in range(boxes.shape[0]):
v = (qrot(boxes[j, 6:].unsqueeze(dim=0).repeat(8, 1),
cube_v_torch * boxes[j, 3:6].unsqueeze(dim=0)) +
boxes[j, :3].unsqueeze(dim=0)).numpy()
for i in range(8):
fout.write('v %f %f %f\n' % (v[i, 0], v[i, 1], v[i, 2]))
for i in range(6):
fout.write('usemtl f%d\n' % i)
fout.write('f %d %d %d\n' %
(cube_f[2 * i, 0] + 8 * j, cube_f[2 * i, 1] + 8 * j,
cube_f[2 * i, 2] + 8 * j))
fout.write(
'f %d %d %d\n' %
(cube_f[2 * i + 1, 0] + 8 * j, cube_f[2 * i + 1, 1] +
8 * j, cube_f[2 * i + 1, 2] + 8 * j))
cmd = 'cd %s && blender -noaudio --background camera_centered.blend --python render_blender.py %s %s > /dev/null' \
% (os.path.join(os.path.dirname(os.path.abspath(__file__))), \
out_fn+'.obj', out_fn)
call(cmd, shell=True)
def transform_pc_batch(pc, feat, anchor=False):
batch_size = feat.size(0)
num_point = pc.size(0)
pc = pc.repeat(batch_size, 1, 1)
center = feat[:, :3].unsqueeze(dim=1).repeat(1, num_point, 1)
shape = feat[:, 3:6].unsqueeze(dim=1).repeat(1, num_point, 1)
quat = feat[:, 6:].unsqueeze(dim=1).repeat(1, num_point, 1)
if not anchor:
pc = pc * shape
pc = qrot(quat.view(-1, 4), pc.view(-1, 3)).view(batch_size, num_point, 3)
if not anchor:
pc = pc + center
return pc
def sensorFromWorld(self, pts_world):
X1 = pts_world
quat = torch.cat((torch.ones(1), self.R1))
quat = quat / quat.norm()
quat = quat.expand(X1.shape[0], 4)
X1 = quaternion.qrot(quat, X1)
#X1_p = torch.cat((X1,torch.zeros(X1.shape[0],1)),dim=1)
#X2 = quaternion.qmul(quat,X1_p)
#qp = quat * torch.tensor([1.0,-1.0,-1.0,-1.0])
#X3 = quaternion.qmul(X2,qp)
#print pts_world, X1, quat
X1.add_(self.T1)
ab1 = X1.div(X1[:, 2].view(X1.shape[0], 1))
ab1 = ab1[:, :2]
r1 = ab1.mul(ab1)
r1 = r1.sum(dim=1, keepdim=True)
r1.sqrt_()
theta1 = r1.atan()
theta1_sq = theta1.mul(theta1)
theta1_d = torch.zeros(theta1.shape)
theta1_d.add_(self.K1[3])
theta1_d.mul_(theta1_sq)
theta1_d.add_(self.K1[2])
theta1_d.mul_(theta1_sq)
theta1_d.add_(self.K1[1])
theta1_d.mul_(theta1_sq)
theta1_d.add_(self.K1[0])
theta1_d.mul_(theta1_sq)
theta1_d.add_(1)
theta1_d.mul_(theta1)
#theta1_scale = theta1_d.div(r1)
theta1_scale = theta1_d.div(theta1)
x1_p = ab1.mul(theta1_scale)
uv_1 = x1_p.mul(self.F1)
uv_1.add_(self.C1)
return uv_1
def get_adj_loss(self, pts, quat, center, adjs):
num_point = pts.shape[1]
part_pcs = qrot(quat.unsqueeze(1).repeat(1, num_point, 1),
pts) + center.unsqueeze(1).repeat(1, num_point, 1)
loss = []
for cur_shape_adj in adjs:
cur_shape_loss = []
for adj in cur_shape_adj:
idx1, idx2 = adj
dist1, dist2 = chamfer_distance(part_pcs[idx1].unsqueeze(0),
part_pcs[idx2].unsqueeze(0),
transpose=False)
cur_loss = torch.min(dist1, dim=1)[0] + torch.min(dist2,
dim=1)[0]
cur_shape_loss.append(cur_loss)
loss.append(torch.stack(cur_shape_loss).mean())
return loss
def other_forward(t, r, t0_, r0_, t1_, r1_):
t0 = t0_.view(-1, 3)
r0 = r0_.view(-1, 4)
t1 = t1_.view(-1, 3)
r1 = r1_.view(-1, 4)
r0_inv = qinv(r0)
t_diff = qrot(r0_inv, (t1 - t0))
r_diff = qmul(r0_inv.clone(), r1)
self_t = t.repeat(t0.shape[0], 1)
self_r = r.repeat(t0.shape[0], 1)
t_loss = t_diff - self_t
self_r_inv = qinv(self_r)
r_loss = qmul(self_r_inv, r_diff)
norm_r_loss = r_loss / qlen(r_loss)
cut_norm_r_loss = norm_r_loss[:, 1:]
return torch.cat([t_loss, cut_norm_r_loss], dim=1)
def linear_assignment(self, mask1, mask2, similar_cnt, pts, centers1,
quats1, centers2, quats2):
'''
mask1, mask 2:
# part_cnt x 224 x 224
similar cnt
# shape: max_mask x 2
# first index is the index of parts without similar parts 1, 2, 3, 4, 4, 5, 5, 6, 6, 6, 6, ....
# second index is the number of similar part count: 1, 1, 1, 2, 2, 2, 2, 4, 4, 4, 4, ....
'''
bids = []
ids1 = []
ids2 = []
inds1 = []
inds2 = []
img_size = mask1.shape[-1]
t = 0
num_point = pts.shape[1]
max_num_part = centers1.shape[1]
# part_cnt = [item for sublist in part_cnt for item in sublist]
with torch.no_grad():
while t < similar_cnt.shape[0]:
cnt = similar_cnt[t].item()
bids = [t] * cnt
cur_mask1 = mask1[t:t + cnt].unsqueeze(1).repeat(
1, cnt, 1, 1).view(-1, img_size, img_size)
cur_mask2 = mask2[t:t + cnt].unsqueeze(0).repeat(
cnt, 1, 1, 1).view(-1, img_size, img_size)
dist_mat_mask = self.get_mask_loss(cur_mask1,
cur_mask2).view(cnt, cnt)
dist_mat_mask = torch.clamp(dist_mat_mask, max=-0.1)
cur_pts = pts[t:t + cnt]
cur_quats1 = quats1[t:t + cnt].unsqueeze(1).repeat(
1, num_point, 1)
cur_centers1 = centers1[t:t + cnt].unsqueeze(1).repeat(
1, num_point, 1)
cur_pts1 = qrot(cur_quats1, cur_pts) + cur_centers1
cur_quats2 = quats2[t:t + cnt].unsqueeze(1).repeat(
1, num_point, 1)
cur_centers2 = centers2[t:t + cnt].unsqueeze(1).repeat(
1, num_point, 1)
cur_pts2 = qrot(cur_quats2, cur_pts) + cur_centers2
cur_pts1 = cur_pts1.unsqueeze(1).repeat(1, cnt, 1, 1).view(
-1, num_point, 3)
cur_pts2 = cur_pts2.unsqueeze(0).repeat(cnt, 1, 1, 1).view(
-1, num_point, 3)
dist1, dist2 = chamfer_distance(cur_pts1,
cur_pts2,
transpose=False)
dist_mat_pts = (dist1.mean(1) + dist2.mean(1)).view(cnt, cnt)
dist_mat_pts = torch.clamp(dist_mat_pts, max=1) * 0.1
dist_mat = torch.add(dist_mat_mask, dist_mat_pts)
t += cnt
rind, cind = linear_sum_assignment(dist_mat.cpu().numpy())
ids1 = list(rind)
ids2 = list(cind)
inds1 += [bids[i] + ids1[i] for i in range(len(ids1))]
inds2 += [bids[i] + ids2[i] for i in range(len(ids2))]
return inds1, inds2
def camera_to_world(X, R, t): Q = np.tile(R, (*X.shape[:-1], 1)) V = X return qrot(Q, V) + t
def world_to_camera(X, R, t): Rt = qinverse(R) Q = np.tile(Rt, (*X.shape[:-1], 1)) V = X - t return qrot(Q, V)
def camera_to_world(X, R, t):
Q = np.tile(R, (*X.shape[:-1], 1))
V = X
return qrot(Q, V) + t
def world_to_camera(X, R, t):
Rt = qinverse(R)
Q = np.tile(Rt, (*X.shape[:-1], 1))
V = X - t
return qrot(Q, V)
a_hist = np.zeros((max_iteration, 3))
p_hist[0,0] = p_c[0,0]
p_hist[0,1] = p_c[1,0]
p_hist[0,2] = p_c[2,0]
a_hist[0,0] = pitch_c
a_hist[0,1] = yaw_c
a_hist[0,2] = roll_c
#qr = qn.angle_to_q(roll_c, pitch_c, yaw_c)
qr = qn.angle_to_q(roll, pitch, yaw)
qa = np.array([0, 1, 0, 0])
qb = qn.qrot(qr, qa)
Rq = qn.q_to_r(qr)
print("original R")
print(R)
print("quaternion to R")
print(Rq)
for i in range(max_iteration):
d_pitch = p_gain * (pitch_t - pitch_c)
d_yaw = p_gain * (yaw_t - yaw_c)
d_roll = p_gain * (roll_t - roll_c)
#pitch_c += d_pitch
#yaw_c += d_yaw
for j in range(cnt):
cur_part_id = cur_parts[j]
cur_part_pose = np.hstack([
matched_pred_center2_all[t + j].cpu().numpy(),
matched_pred_quat2_all[t + j].cpu().numpy()
])
to_save[cur_part_id] = cur_part_pose
fn = input_shape_id[t] + '-' + input_view_id[t] + '.npy'
np.save(os.path.join(test_res_matched_dir, fn), to_save)
print('saved', fn)
t += cnt
# get metric stats
pred_pts = qrot(
matched_pred_quat2_all.unsqueeze(1).repeat(1, num_point, 1),
input_pts) + matched_pred_center2_all.unsqueeze(1).repeat(
1, num_point, 1)
gt_pts = qrot(
matched_gt_quat_all.unsqueeze(1).repeat(1, num_point, 1),
input_pts) + matched_gt_center_all.unsqueeze(1).repeat(
1, num_point, 1)
dist1, dist2 = chamfer_distance(gt_pts,
pred_pts,
transpose=False,
sqrt=True)
dist = torch.mean(dist1, dim=1).cpu().numpy() + torch.mean(
dist2, dim=1).cpu().numpy()
dist /= 2.0
batch_size = input_box_size.shape[0]
correct = 0
def __getitem__(self, index):
shape_id = self.data[index]
cur_data_fn = os.path.join(self.data_dir, self.category,
'%s.npy' % (shape_id))
cur_data = np.load(cur_data_fn, allow_pickle=True).item()
view_id = self.convert_view(np.random.choice(np.arange(24)))
if self.data_split == "val" or self.data_split == "vis" or self.data_split == "test":
view_id = "00"
cur_data_view_fn = os.path.join(self.data_dir, self.category, shape_id,
'%s.npy' % (view_id))
cur_data_view = np.load(cur_data_view_fn, allow_pickle=True).item()
part_id = np.random.choice(cur_data['all_parts'])
equiv_edge_indices = []
# sort part order to put the single part in front and similar parts in the back
single_parts = [
x for x in cur_data['all_parts']
if len(cur_data[x]['similar']) == 1
]
# similar_part_inds = [i+1 for i in range(len(single_parts))]
other_parts = [
x for x in cur_data['all_parts'] if len(cur_data[x]['similar']) > 1
]
part_order = sorted(single_parts)
t = len(part_order)
for part in other_parts:
if part not in part_order:
cur_cnt = len(cur_data[part]['similar'])
part_order += sorted(cur_data[part]['similar'])
for t1 in range(t, t + cur_cnt):
for t2 in range(t, t + cur_cnt):
if t1 != t2:
equiv_edge_indices.append([t1, t2])
t += cur_cnt
if len(part_order) != len(cur_data['all_parts']):
print('check data ', shape_id, view_id)
data_feats = ()
for feat in self.data_features:
if feat == 'img':
img_fn = os.path.join(self.img_path, shape_id,
'view-{}'.format(view_id),
'shape-rgb.png')
if not os.path.exists(img_fn):
return None
with Image.open(img_fn) as fimg:
out = np.array(fimg, dtype=np.float32) / 255
white_img = np.ones((self.img_size, self.img_size, 3),
dtype=np.float32)
mask = np.tile(out[:, :, 3:4], [1, 1, 3])
out = out[:, :, :3] * mask + white_img * (1 - mask)
out = torch.from_numpy(out).permute(2, 0, 1).unsqueeze(0)
data_feats = data_feats + (out, )
elif feat == 'mask':
out = np.zeros(
(self.max_num_mask, self.img_size, self.img_size),
dtype=np.float32)
if len(cur_data['all_parts']) > self.max_num_mask:
return None
if len(cur_data['all_parts']) <= 1:
return None
for i in range(len(part_order)):
part_id = part_order[i]
mask_ids = np.array(cur_data_view[part_id]['mask'])
if len(mask_ids) > 0:
out[i, mask_ids[:, 0],
mask_ids[:,
1]] = 1 # real mask starts from [:, 1:, 1:] of the tensor
out = torch.from_numpy(out).unsqueeze(0)
data_feats = data_feats + (
out, ) # output tenshor shape [1, max mask, 225, 225]
elif feat == 'pts':
out = torch.zeros(
self.max_num_mask, self.num_points,
3).float() # first-dim is 0/1, meaning exist or not
if len(cur_data['all_parts']) > self.max_num_mask:
return None
if len(cur_data['all_parts']) <= 1:
return None
for i in range(len(part_order)):
part_id = part_order[i]
if len(cur_data[part_id]
['similar']) > self.max_num_similar_parts:
return None
out[i, :, :] = torch.from_numpy(
cur_data[part_id]['pts']).float()
out = out.unsqueeze(0)
data_feats = data_feats + (out, )
elif feat == 'box_size':
out = torch.zeros(
self.max_num_mask,
3).float() # first-dim is 0/1, meaning exist or not
if len(cur_data['all_parts']) > self.max_num_mask:
return None
if len(cur_data['all_parts']) <= 1:
return None
for i in range(len(part_order)):
part_id = part_order[i]
if len(cur_data[part_id]
['similar']) > self.max_num_similar_parts:
return None
out[i, :] = torch.from_numpy(
cur_data[part_id]['bbox_size']).float()
out = out.unsqueeze(0)
data_feats = data_feats + (out, )
elif feat == 'anchor':
out = torch.zeros(
self.max_num_mask, 6,
3).float() # first-dim is 0/1, meaning exist or not
if len(cur_data['all_parts']) > self.max_num_mask:
return None
if len(cur_data['all_parts']) <= 1:
return None
for i in range(len(part_order)):
part_id = part_order[i]
if len(cur_data[part_id]
['similar']) > self.max_num_similar_parts:
return None
xyz = cur_data[part_id]['bbox_size']
neg_xyz = -1.0 * cur_data[part_id]['bbox_size']
out[i, 0, :] = torch.tensor([xyz[0], 0, 0]).float()
out[i, 1, :] = torch.tensor([0, xyz[1], 0]).float()
out[i, 2, :] = torch.tensor([0, 0, xyz[2]]).float()
out[i, 3, :] = torch.tensor([neg_xyz[0], 0, 0]).float()
out[i, 4, :] = torch.tensor([0, neg_xyz[1], 0]).float()
out[i, 5, :] = torch.tensor([0, 0, neg_xyz[2]]).float()
out = out.unsqueeze(0)
data_feats = data_feats + (out, )
elif feat == 'sem_one_hot':
out = torch.zeros(self.max_num_mask,
len(self.part_sems)).float()
if len(cur_data['all_parts']) > self.max_num_mask:
return None
if len(cur_data['all_parts']) <= 1:
return None
for i in range(len(part_order)):
part_id = part_order[i]
if len(cur_data[part_id]
['similar']) > self.max_num_similar_parts:
return None
out[i, self.part_sem2id[cur_data[part_id]['sem']]] = 1
out = out.unsqueeze(0)
data_feats = data_feats + (out, )
elif feat == 'ins_one_hot':
# instance one hot: shape : max parts allowed x max_similar number of parts
# out [i, :] is the instance one hot of the parts
# [1, 0 .... ], [1, 0, ..... ] [1, 0, ..... ] [0, 1, ..... ]
out = torch.zeros(self.max_num_mask,
self.max_num_similar_parts).float()
if len(cur_data['all_parts']) > self.max_num_mask:
return None
if len(cur_data['all_parts']) <= 1:
return None
for i in range(len(part_order)):
part_id = part_order[i]
if len(cur_data[part_id]
['similar']) > self.max_num_similar_parts:
return None
counter = 0
for part in sorted(cur_data[part_id]['similar']):
ind = part_order.index(part)
out[ind, counter] = 1
counter += 1
out = out.unsqueeze(0)
data_feats = data_feats + (out, )
elif feat == 'shape_id':
if len(cur_data['all_parts']) > self.max_num_mask:
return None
data_feats = data_feats + (shape_id, )
elif feat == 'part_id':
if len(cur_data['all_parts']) > self.max_num_mask:
return None
data_feats = data_feats + (part_order, )
elif feat == 'view_id':
if len(cur_data['all_parts']) > self.max_num_mask:
return None
data_feats = data_feats + (view_id, )
elif feat == 'adj':
adj_pairs = []
with open(
os.path.join(self.data_dir, self.category,
shape_id + '-adj.txt'), 'r') as f:
for line in f:
l = line.strip()
adj_pairs.append((l.split()[0], l.split()[1]))
out = []
for pair in adj_pairs:
out.append(
(part_order.index(pair[0]), part_order.index(pair[1])))
data_feats = data_feats + (out, )
elif feat == 'adj_pt':
adj_idx_pairs = []
with open(
os.path.join(self.data_dir, self.category,
shape_id + '-adj-ptid.txt'), 'r') as f:
for line in f:
l = line.strip()
adj_idx_pairs.append((l.split()[0], l.split()[1],
l.split()[2], l.split()[3]))
out = []
for i in range(len(adj_idx_pairs)):
pair = adj_idx_pairs[i]
pts1 = torch.from_numpy(cur_data[pair[0]]['pts']).float()
pts2 = torch.from_numpy(cur_data[pair[1]]['pts']).float()
t1, rot1 = cur_data_view[pair[0]][
'cam_dof'][:3], cur_data_view[pair[0]]['cam_dof'][3:]
t2, rot2 = cur_data_view[pair[1]][
'cam_dof'][:3], cur_data_view[pair[1]]['cam_dof'][3:]
t1 = torch.from_numpy(t1).float()
t2 = torch.from_numpy(t2).float()
rot1 = torch.from_numpy(rot1).float()
rot2 = torch.from_numpy(rot2).float()
num_point = 1000
pc1 = qrot(rot1.unsqueeze(0).repeat(num_point, 1),
pts1) + t1.unsqueeze(0).repeat(num_point, 1)
pc2 = qrot(rot2.unsqueeze(0).repeat(num_point, 1),
pts2) + t2.unsqueeze(0).repeat(num_point, 1)
xyz1 = pc1[int(pair[2])]
xyz2 = pc2[int(pair[3])]
xyz1 = xyz1.float()
xyz2 = xyz2.float()
cur_pair_xyz = torch.zeros(2, 3)
cur_pair_xyz[0] = xyz1
cur_pair_xyz[1] = xyz2
out.append(cur_pair_xyz)
out = torch.stack(out)
data_feats = data_feats + (out, )
elif feat == 'total_parts_cnt':
if len(cur_data['all_parts']) > self.max_num_mask:
return None
if len(cur_data['all_parts']) <= 1:
return None
out = len(cur_data['all_parts'])
data_feats = data_feats + (out, )
elif feat == 'similar_parts_cnt':
# shape: max_mask x 2
# first index is the index of parts without similar parts 1, 2, 3, 4, 4, 5, 5, 6, 6, 6, 6, ....
# second index is the number of similar part count: 1, 1, 1, 2, 2, 2, 2, 4, 4, 4, 4, ....
out = torch.zeros(self.max_num_mask, 1, dtype=torch.long)
if len(cur_data['all_parts']) > self.max_num_mask:
return None
if len(cur_data['all_parts']) <= 1:
return None
total_part_cnt = len(cur_data['all_parts'])
counter = 0
for i in range(len(part_order)):
part_id = part_order[i]
if len(cur_data[part_id]
['similar']) > self.max_num_similar_parts:
return None
out[i, 0] = len(cur_data[part_id]['similar'])
out = out.unsqueeze(0)
data_feats = data_feats + (out, )
elif feat == 'similar_parts_edge_indices':
out = torch.Tensor(equiv_edge_indices).long()
data_feats = data_feats + (out, )
elif feat == 'parts_cam_dof':
out = torch.zeros(self.max_num_mask, 7).float()
if len(cur_data[part_id]
['similar']) > self.max_num_similar_parts:
return None
if len(cur_data['all_parts']) <= 1:
return None
for i in range(len(part_order)):
part_id = part_order[i]
if len(cur_data[part_id]
['similar']) > self.max_num_similar_parts:
return None
out[i, :] = torch.from_numpy(
cur_data_view[part_id]['cam_dof'])
data_feats = data_feats + (out, )
else:
raise ValueError('ERROR: unknown feat type %s!' % feat)
return data_feats
n_diff = np.linalg.norm(diff)
print("|diff|: {0}".format(n_diff))
# angle estimation using quaternion
print("Quaternion")
qa = np.array([0, p_ref[0][0], p_ref[1][0], p_ref[2][0]])
qb = np.array([0, p_c[0][0], p_c[1][0], p_c[2][0]])
qr = qn.solve_qr_lm(qa, qb)
roll_q, pitch_q, yaw_q = qn.q_to_angle(qr)
print("pitch: {0} --> {1}".format(pitch, pitch_q))
print("yaw: {0} --> {1}".format(yaw, yaw_q))
print("roll: {0} --> {1}".format(roll, roll_q))
inv_qr = qn.qinv(qr)
qa_reproj = qn.qrot(inv_qr, qb)
p_reprojq = qa_reproj[1:]
print("inv_qr: {0}".format(inv_qr))
diffq = np.array(p_ref).flatten() - p_reprojq
n_diffq = np.linalg.norm(diffq)
print("|diffq|: {0}".format(n_diffq))
print("Euler")
roll_e, pitch_e, yaw_e = ea.solve_angle(np.array(p_ref).flatten(), np.array(p_c).flatten())
print("pitch: {0} --> {1}".format(pitch, pitch_e))
print("yaw: {0} --> {1}".format(yaw, yaw_e))
print("roll: {0} --> {1}".format(roll, roll_e))
def forward(batch, data_features, network, conf, \
is_val=False, step=None, epoch=None, batch_ind=0, num_batch=1, start_time=0, \
log_console=False, log_tb=False, tb_writer=None, lr=None):
# prepare input
batch_index = 1
if len(batch) == 0:
return None
cur_batch_size = len(batch[data_features.index('total_parts_cnt')])
total_part_cnt = batch[data_features.index('total_parts_cnt')][0]
input_total_part_cnt = batch[data_features.index('total_parts_cnt')][
0] # 1
input_img = batch[data_features.index('img')][0] # 3 x H x W
input_img = input_img.repeat(input_total_part_cnt, 1, 1,
1) # part_cnt 3 x H x W
input_pts = batch[data_features.index('pts')][0].squeeze(
0)[:input_total_part_cnt]
input_ins_one_hot = batch[data_features.index('ins_one_hot')][0].squeeze(
0)[:input_total_part_cnt] # part_cnt x max_similar_parts
input_similar_part_cnt = batch[data_features.index('similar_parts_cnt')][
0].squeeze(0)[:input_total_part_cnt] # part_cnt x 1
input_shape_id = [batch[data_features.index('shape_id')][0]
] * input_total_part_cnt
input_view_id = [batch[data_features.index('view_id')][0]
] * input_total_part_cnt
# prepare gt:
gt_cam_dof = batch[data_features.index('parts_cam_dof')][0].squeeze(
0)[:input_total_part_cnt]
gt_mask = [
batch[data_features.index('mask')][0].squeeze(0)
[:input_total_part_cnt].to(conf.device)
]
input_total_part_cnt = [batch[data_features.index('total_parts_cnt')][0]]
input_similar_parts_edge_indices = [
batch[data_features.index('similar_parts_edge_indices')][0].to(
conf.device)
]
while total_part_cnt < 70 and batch_index < cur_batch_size:
cur_input_cnt = batch[data_features.index(
'total_parts_cnt')][batch_index]
total_part_cnt += cur_input_cnt
if total_part_cnt > 90:
total_part_cnt -= cur_input_cnt
batch_index += 1
continue
cur_batch_img = batch[data_features.index('img')][batch_index].repeat(
cur_input_cnt, 1, 1, 1)
input_img = torch.cat((input_img, cur_batch_img), dim=0)
input_pts = torch.cat((input_pts, batch[data_features.index('pts')]
[batch_index].squeeze(0)[:cur_input_cnt]),
dim=0)
input_ins_one_hot = torch.cat(
(input_ins_one_hot, batch[data_features.index('ins_one_hot')]
[batch_index].squeeze(0)[:cur_input_cnt]),
dim=0) # B x max_parts x max_similar_parts
input_total_part_cnt.append(
batch[data_features.index('total_parts_cnt')][batch_index]) # 1
input_similar_part_cnt = torch.cat(
(input_similar_part_cnt, batch[data_features.index(
'similar_parts_cnt')][batch_index].squeeze(0)[:cur_input_cnt]),
dim=0) # B x max_parts x 2
input_shape_id += [
batch[data_features.index('shape_id')][batch_index]
] * cur_input_cnt
input_view_id += [batch[data_features.index('view_id')][batch_index]
] * cur_input_cnt
gt_cam_dof = torch.cat((gt_cam_dof, batch[data_features.index(
'parts_cam_dof')][batch_index].squeeze(0)[:cur_input_cnt]),
dim=0)
# prepare gt
gt_mask.append(batch[data_features.index('mask')][batch_index].squeeze(
0)[:cur_input_cnt].to(conf.device))
input_similar_parts_edge_indices.append(batch[data_features.index(
'similar_parts_edge_indices')][batch_index].to(conf.device))
batch_index += 1
input_img = input_img.to(conf.device)
input_pts = input_pts.to(conf.device)
input_similar_part_cnt = input_similar_part_cnt.to(conf.device)
input_ins_one_hot = input_ins_one_hot.to(conf.device)
gt_cam_dof = gt_cam_dof.to(conf.device)
# prepare gt
gt_center = gt_cam_dof[:, :3] # B x 3
gt_quat = gt_cam_dof[:, 3:] # B x 4
batch_size = input_img.shape[0]
num_point = input_pts.shape[1]
# forward through the network
pred_masks, pred_center, pred_quat, pred_center2, pred_quat2 = network(
input_img - 0.5, input_pts, input_ins_one_hot, input_total_part_cnt,
input_similar_parts_edge_indices)
mask_loss_per_data = []
t = 0
matched_pred_mask_all = []
matched_gt_mask_all = []
matched_mask_loss_per_data_all = []
matched_pred_center_all = []
matched_gt_center_all = []
matched_pred_center2_all = []
matched_pred_quat_all = []
matched_gt_quat_all = []
matched_pred_quat2_all = []
matched_ins_onehot_all = []
for i in range(len(input_total_part_cnt)):
total_cnt = input_total_part_cnt[i]
matched_gt_ids, matched_pred_ids = network.linear_assignment(gt_mask[i], pred_masks[i], \
input_similar_part_cnt[t:t+total_cnt], input_pts[t:t+total_cnt], gt_center[t:t+total_cnt], \
gt_quat[t:t+total_cnt], pred_center[t:t+total_cnt], pred_quat[t:t+total_cnt])
# select the matched data
matched_pred_mask = pred_masks[i][matched_pred_ids]
matched_gt_mask = gt_mask[i][matched_gt_ids]
matched_pred_center = pred_center[t:t + total_cnt][matched_pred_ids]
matched_pred_center2 = pred_center2[t:t + total_cnt][matched_pred_ids]
matched_gt_center = gt_center[t:t + total_cnt][matched_gt_ids]
matched_pred_quat = pred_quat[t:t + total_cnt][matched_pred_ids]
matched_pred_quat2 = pred_quat2[t:t + total_cnt][matched_pred_ids]
matched_gt_quat = gt_quat[t:t + total_cnt][matched_gt_ids]
matched_ins_onehot = input_ins_one_hot[t:t +
total_cnt][matched_pred_ids]
matched_ins_onehot_all.append(matched_ins_onehot)
matched_gt_mask_all.append(matched_gt_mask)
matched_pred_mask_all.append(matched_pred_mask)
matched_pred_center_all.append(matched_pred_center)
matched_pred_center2_all.append(matched_pred_center2)
matched_gt_center_all.append(matched_gt_center)
matched_pred_quat_all.append(matched_pred_quat)
matched_pred_quat2_all.append(matched_pred_quat2)
matched_gt_quat_all.append(matched_gt_quat)
# for computing mask soft iou loss
matched_mask_loss_per_data = network.get_mask_loss(
matched_pred_mask, matched_gt_mask)
matched_mask_loss_per_data_all.append(matched_mask_loss_per_data)
mask_loss_per_data.append(matched_mask_loss_per_data.mean())
t += total_cnt
matched_ins_onehot_all = torch.cat(matched_ins_onehot_all, dim=0)
matched_pred_mask_all = torch.cat(matched_pred_mask_all, dim=0)
matched_gt_mask_all = torch.cat(matched_gt_mask_all, dim=0)
matched_mask_loss_per_data_all = torch.cat(matched_mask_loss_per_data_all,
dim=0)
matched_pred_quat_all = torch.cat(matched_pred_quat_all, dim=0)
matched_pred_quat2_all = torch.cat(matched_pred_quat2_all, dim=0)
matched_gt_quat_all = torch.cat(matched_gt_quat_all, dim=0)
matched_pred_center_all = torch.cat(matched_pred_center_all, dim=0)
matched_pred_center2_all = torch.cat(matched_pred_center2_all, dim=0)
matched_gt_center_all = torch.cat(matched_gt_center_all, dim=0)
center_loss_per_data = network.get_center_loss(matched_pred_center_all, matched_gt_center_all) + \
network.get_center_loss(matched_pred_center2_all, matched_gt_center_all)
quat_loss_per_data = network.get_quat_loss(input_pts, matched_pred_quat_all, matched_gt_quat_all) + \
network.get_quat_loss(input_pts, matched_pred_quat2_all, matched_gt_quat_all)
l2_rot_loss_per_data = network.get_l2_rotation_loss(input_pts, matched_pred_quat_all, matched_gt_quat_all) + \
network.get_l2_rotation_loss(input_pts, matched_pred_quat2_all, matched_gt_quat_all)
whole_shape_cd_per_data = network.get_shape_chamfer_loss(input_pts, matched_pred_quat_all, matched_gt_quat_all, matched_pred_center_all, matched_gt_center_all, input_total_part_cnt) + \
network.get_shape_chamfer_loss(input_pts, matched_pred_quat2_all, matched_gt_quat_all, matched_pred_center2_all, matched_gt_center_all, input_total_part_cnt)
# for each type of loss, compute avg loss per batch
mask_loss_per_data = torch.stack(mask_loss_per_data)
center_loss = center_loss_per_data.mean()
quat_loss = quat_loss_per_data.mean()
mask_loss = mask_loss_per_data.mean()
l2_rot_loss = l2_rot_loss_per_data.mean()
shape_chamfer_loss = whole_shape_cd_per_data.mean()
# compute total loss
total_loss = \
center_loss * conf.loss_weight_center + \
quat_loss * conf.loss_weight_quat + \
l2_rot_loss * conf.loss_weight_l2_rot + \
shape_chamfer_loss * conf.loss_weight_shape_chamfer
# display information
data_split = 'train'
if is_val:
data_split = 'val'
with torch.no_grad():
# log to console
if log_console:
utils.printout(conf.flog, \
f'''{strftime("%H:%M:%S", time.gmtime(time.time()-start_time)):>9s} '''
f'''{epoch:>5.0f}/{conf.epochs:<5.0f} '''
f'''{data_split:^10s} '''
f'''{batch_ind:>5.0f}/{num_batch:<5.0f} '''
f'''{100. * (1+batch_ind+num_batch*epoch) / (num_batch*conf.epochs):>9.1f}% '''
f'''{lr:>5.2E} '''
f'''{mask_loss.item():>10.5f}'''
f'''{center_loss.item():>10.5f}'''
f'''{quat_loss.item():>10.5f}'''
f'''{l2_rot_loss.item():>10.5f}'''
f'''{shape_chamfer_loss.item():>10.5f}'''
f'''{total_loss.item():>10.5f}''')
conf.flog.flush()
# log to tensorboard
if log_tb and tb_writer is not None:
tb_writer.add_scalar('mask_loss', mask_loss.item(), step)
tb_writer.add_scalar('center_loss', center_loss.item(), step)
tb_writer.add_scalar('quat_loss', quat_loss.item(), step)
tb_writer.add_scalar('l2 rotation_loss', l2_rot_loss.item(), step)
tb_writer.add_scalar('shape_chamfer_loss',
shape_chamfer_loss.item(), step)
tb_writer.add_scalar('total_loss', total_loss.item(), step)
tb_writer.add_scalar('lr', lr, step)
# gen visu
if is_val and (
not conf.no_visu) and epoch % conf.num_epoch_every_visu == 0:
visu_dir = os.path.join(conf.exp_dir, 'val_visu')
out_dir = os.path.join(visu_dir, 'epoch-%04d' % epoch)
input_img_dir = os.path.join(out_dir, 'input_img')
input_pts_dir = os.path.join(out_dir, 'input_pts')
gt_mask_dir = os.path.join(out_dir, 'gt_mask')
pred_mask_dir = os.path.join(out_dir, 'pred_mask')
gt_dof_dir = os.path.join(out_dir, 'gt_dof')
pred_dof_dir = os.path.join(out_dir, 'pred_dof')
pred_dof2_dir = os.path.join(out_dir, 'pred_dof2')
info_dir = os.path.join(out_dir, 'info')
if batch_ind == 0:
# create folders
os.mkdir(out_dir)
os.mkdir(input_img_dir)
os.mkdir(input_pts_dir)
os.mkdir(gt_mask_dir)
os.mkdir(pred_mask_dir)
os.mkdir(gt_dof_dir)
os.mkdir(pred_dof_dir)
os.mkdir(pred_dof2_dir)
os.mkdir(info_dir)
if batch_ind < conf.num_batch_every_visu:
utils.printout(conf.flog, 'Visualizing ...')
# compute pred_pts and gt_pts
pred_pts = qrot(
matched_pred_quat_all.unsqueeze(1).repeat(1, num_point, 1),
input_pts) + matched_pred_center_all.unsqueeze(1).repeat(
1, num_point, 1)
pred2_pts = qrot(
matched_pred_quat2_all.unsqueeze(1).repeat(
1, num_point, 1),
input_pts) + matched_pred_center2_all.unsqueeze(1).repeat(
1, num_point, 1)
gt_pts = qrot(
matched_gt_quat_all.unsqueeze(1).repeat(1, num_point, 1),
input_pts) + matched_gt_center_all.unsqueeze(1).repeat(
1, num_point, 1)
t = 0
for i in range(batch_size):
fn = 'data-%03d.png' % (batch_ind * batch_size + i)
cur_input_img = (
input_img[i].permute(1, 2, 0).cpu().numpy() *
255).astype(np.uint8)
Image.fromarray(cur_input_img).save(
os.path.join(input_img_dir, fn))
cur_input_pts = input_pts[i].cpu().numpy()
render_utils.render_pts(os.path.join(
BASE_DIR, input_pts_dir, fn),
cur_input_pts,
blender_fn='object_centered.blend')
cur_gt_mask = (matched_gt_mask_all[i].cpu().numpy() >
0.5).astype(np.uint8) * 255
Image.fromarray(cur_gt_mask).save(
os.path.join(gt_mask_dir, fn))
cur_pred_mask = (matched_pred_mask_all[i].cpu().numpy() >
0.5).astype(np.uint8) * 255
Image.fromarray(cur_pred_mask).save(
os.path.join(pred_mask_dir, fn))
cur_pred_pts = pred_pts[i].cpu().numpy()
render_utils.render_pts(os.path.join(
BASE_DIR, pred_dof_dir, fn),
cur_pred_pts,
blender_fn='camera_centered.blend')
cur_pred_pts = pred2_pts[i].cpu().numpy()
render_utils.render_pts(os.path.join(
BASE_DIR, pred_dof2_dir, fn),
cur_pred_pts,
blender_fn='camera_centered.blend')
cur_gt_pts = gt_pts[i].cpu().numpy()
render_utils.render_pts(os.path.join(
BASE_DIR, gt_dof_dir, fn),
cur_gt_pts,
blender_fn='camera_centered.blend')
with open(
os.path.join(info_dir, fn.replace('.png', '.txt')),
'w') as fout:
fout.write('shape_id: %s, view_id: %s\n' % (\
input_shape_id[i],\
input_view_id[i]))
fout.write('ins onehot %s\n' %
matched_ins_onehot_all[i])
fout.write('mask_loss: %f\n' %
matched_mask_loss_per_data_all[i].item())
fout.write('center_loss: %f\n' %
center_loss_per_data[i].item())
fout.write('quat_loss: %f\n' %
quat_loss_per_data[i].item())
fout.write('l2_rot_loss: %f\n' %
l2_rot_loss_per_data[i].item())
fout.write('shape_chamfer_loss %f\n' %
shape_chamfer_loss.item())
return total_loss
def forward(self, img, pc, pred_masks, part_feat, part_cnt,
equiv_edge_indices):
# get image feature
img = img.clone()
img[:, 0] = (img[:, 0] - 0.485) / 0.229
img[:, 1] = (img[:, 1] - 0.456) / 0.224
img[:, 2] = (img[:, 2] - 0.406) / 0.225
img_feat = torch.relu(self.resnet_final_bn(self.resnet(img)))
# get mask feature
all_masks = (torch.cat(pred_masks, dim=0).unsqueeze(dim=1) >
0.5).float()
mask_feat = torch.relu(
self.mask_resnet_final_bn(self.mask_resnet(all_masks)))
# get global feature per shape
part_feat_5d = torch.cat((img_feat, mask_feat, part_feat), dim=1)
part_feat_5d = torch.relu(self.mlp3(part_feat_5d))
# perform graph-conv
t = 0
i = 0
output_part_feat_5d = []
for cnt in part_cnt:
child_feats = part_feat_5d[t:t + cnt]
iter_feats = [child_feats]
cur_equiv_edge_indices = equiv_edge_indices[i]
for j in range(self.num_iteration):
if cur_equiv_edge_indices.shape[0] > 0:
node_edge_feats = torch.cat([
child_feats[cur_equiv_edge_indices[:, 0], :],
child_feats[cur_equiv_edge_indices[:, 1], :]
],
dim=1)
node_edge_feats = torch.relu(
self.node_edge_op[j](node_edge_feats))
new_child_feats = child_feats.new_zeros(cnt, 256)
new_child_feats = torch_scatter.scatter_mean(
node_edge_feats,
cur_equiv_edge_indices[:, 0],
dim=0,
out=new_child_feats)
child_feats = new_child_feats
iter_feats.append(child_feats)
all_iter_feat = torch.cat(iter_feats, dim=1)
output_part_feat_5d.append(all_iter_feat)
t += cnt
i += 1
output_part_feat_5d = torch.cat(output_part_feat_5d, dim=0)
# decode pose
center, quat = self.pose_decoder(output_part_feat_5d)
# refine final poses
num_point = pc.shape[1]
with torch.no_grad():
pc2 = qrot(quat.unsqueeze(1).repeat(1, num_point, 1),
pc) + center.unsqueeze(1).repeat(1, num_point, 1)
feat = torch.cat([output_part_feat_5d, center, quat], dim=1)
center2, quat2 = self.refine_net(pc2, feat, part_cnt,
equiv_edge_indices)
return center, quat, center2, quat2
def forward(self, vert):
vert = quat.qrot(
(self.rotation / (self.rotation.norm())).repeat(vert.shape[0],
1), vert - 0.5)
vert = vert + self.translation.unsqueeze(0) + 0.5
return vert
gt_w[20:40, :] = [0.00, 0.05, 0.0]
gt_w[50:60, :] = [0.01, -0.05, 0.01]
gt_w[70:90, :] = [-0.01, -0.00, 0.0]
for i in range(2, seq_len):
gt_angle[i] = gt_angle[i - 1] + gt_w[i]
# generate accelerometer G.T.
gn = [0, 0, 0, 1]
gt_a = np.zeros((seq_len, 3)) # gx, gy, gz
# ya,t = -Rt gn + e_a,t
for i in range(seq_len):
qr = qt.angle_to_q(gt_angle[i, 0], gt_angle[i, 1], gt_angle[i, 2])
qa = qt.qrot(qr, gn)
gt_a[i, 0] = -qa[1]
gt_a[i, 1] = -qa[2]
gt_a[i, 2] = -qa[3]
# generate magnetometer G.T.
gt_m = np.zeros((seq_len, 3))
mn = np.array([0, 1, 0, 0])
for i in range(seq_len):
qr = qt.angle_to_q(gt_angle[i, 0], gt_angle[i, 1], gt_angle[i, 2])
qm = qt.qrot(qr, mn)
b = np.zeros((3, 1))
b[0, 0] = qm[1]
b[1, 0] = qm[2]
