def gauss_smoothen_image(cfg, img, sigma_rel):
fsz = cfg.pc_gauss_kernel_size
kernel = gauss_kernel_1d(fsz, sigma_rel)
in_channels = img.shape[-1]
k1 = torch.repeat(kernel.reshape(1, fsz, 1, 1), (1, 1, in_channels, 1))
k2 = torch.repeat(kernel.reshape((fsz, 1, 1, 1)), (1, 1, in_channels, 1))
img_tmp = img
img_tmp = tf.nn.depthwise_conv2d(img_tmp, k1, [1, 1, 1, 1], padding="SAME")
img_tmp = tf.nn.depthwise_conv2d(img_tmp, k2, [1, 1, 1, 1], padding="SAME")
return img_tmp
def feature_distance(img1, img2, device):
if (features_model is None):
model = models.vgg19(pretrained=True).to(device=device)
model.eval()
layer = model.features
if (img1.shape[1] == 1):
img1 = torch.repeat(img1, 3, axis=1)
if (img2.shape[1] == 1):
img2 = torch.repeat(img2, 3, axis=1)
img1_feature_vector = layer(img1_tensor)
img2_feature_vector = layer(img2_tensor)
return ((img1_feature_vector - img2_feature_vector)**2).mean()
def batch_global_rigid_transformation(Rs,
Js,
parent,
rotate_base=False,
opts=None):
"""
Computes absolute joint locations given pose.
rotate_base: if True, rotates the global rotation by 90 deg in x axis.
if False, this is the original SMPL coordinate.
Args:
Rs: N x 24 x 3 x 3 rotation vector of K joints
Js: N x 24 x 3, joint locations before posing
parent: 24 holding the parent id for each index
Returns
new_J : `Tensor`: N x 24 x 3 location of absolute joints
A : `Tensor`: N x 24 4 x 4 relative joint transformations for LBS.
"""
if rotate_base:
print('Flipping the SMPL coordinate frame!!!!')
rot_x = torch.Tensor([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
rot_x = torch.reshape(torch.repeat(rot_x, [N, 1]),
[N, 3, 3]) # In tf it was tile
root_rotation = torch.matmul(Rs[:, 0, :, :], rot_x)
else:
root_rotation = Rs[:, 0, :, :]
# Now Js is N x 24 x 3 x 1
Js = Js.unsqueeze(-1)
N = Rs.shape[0]
def make_A(R, t):
# Rs is N x 3 x 3, ts is N x 3 x 1
R_homo = torch.nn.functional.pad(R, (0, 0, 0, 1, 0, 0))
t_homo = torch.cat([t, torch.ones([N, 1, 1]).to(device=Rs.device)], 1)
return torch.cat([R_homo, t_homo], 2)
A0 = make_A(root_rotation, Js[:, 0])
results = [A0]
for i in range(1, parent.shape[0]):
j_here = Js[:, i] - Js[:, parent[i]]
A_here = make_A(Rs[:, i], j_here)
res_here = torch.matmul(results[parent[i]], A_here)
results.append(res_here)
# 10 x 24 x 4 x 4
results = torch.stack(results, dim=1)
new_J = results[:, :, :3, 3]
# --- Compute relative A: Skinning is based on
# how much the bone moved (not the final location of the bone)
# but (final_bone - init_bone)
# ---
Js_w0 = torch.cat([Js, torch.zeros([N, 35, 1, 1]).to(device=Rs.device)], 2)
init_bone = torch.matmul(results, Js_w0)
# Append empty 4 x 3:
init_bone = torch.nn.functional.pad(init_bone, (3, 0, 0, 0, 0, 0, 0, 0))
A = results - init_bone
return new_J, A
def torch_place(arr, mask, vals):
n = mask.sum()
vals = vals.flatten()
if vals.numel() < n:
reps = torch.ceil(n / vals.numel())
vals = torch.repeat(vals, reps)
arr[mask] = vals[:n]
def forward(self, input, labels, step=0, alpha=-1):
for i in range(step, -1, -1):
index = self.n_layer - i - 1
label_embeddings = label_emb(labels) # N x K
label_embeddings = torch.repeat(1, 1, input.size(2), input.size(3))
input = torch.cat([input, label_embeddings], dim=1)
if i == step:
out = self.from_rgb[index](input)
if i == 0:
out_std = torch.sqrt(out.var(0, unbiased=False) + 1e-8)
mean_std = out_std.mean()
mean_std = mean_std.expand(out.size(0), 1, 4, 4)
out = torch.cat([out, mean_std], 1)
out = self.progression[index](out)
if i > 0:
if i == step and 0 <= alpha < 1:
skip_rgb = F.avg_pool2d(input, 2)
skip_rgb = self.from_rgb[index + 1](skip_rgb)
out = (1 - alpha) * skip_rgb + alpha * out
out = out.squeeze(2).squeeze(2)
# print(input.size(), out.size(), step)
out = self.linear(out)
return out
def with_denormalize(self, policy):
assert (torch.max(policy) <= 1) or (torch.min(policy) >=
-1), "denormalized"
if len(policy.shape) == 2:
assert (policy.shape[1] == self.output_size
), "action error, shape is {} and not {}".format(
policy.shape[1], self.output_size)
upper = torch.repeat(self.upper_bounds.reshape(1, -1),
policy.shape[0],
axis=0)
lower = torch.repeat(self.lower_bounds.reshape(1, -1),
policy.shape[0],
axis=0)
else:
upper = self.upper_bounds
lower = self.lower_bounds
policy = 0.5 * (policy + 1) * (upper - lower) + lower
return policy
def tfm_padtrim_signal(sig, width, pad_mode="zeros"):
'''Pad signal to specified width, using specified pad mode'''
c, x = sig.shape
if (x == width): return sig
elif (x > width): return sig[:, :width]
elif pad_mode.lower() == "zeros":
padding = torch.zeros((c, width - x))
return torch.cat((sig, padding), 1)
elif pad_mode.lower() == "repeat":
repeats = width // x + 1
return torch.repeat(1, repeats)[:, :width]
else:
raise ValueError(
f"pad_mode {pad_mode} not currently supported, only 'zeros', or 'repeat'"
)
def get_instance_seg_v1_net(self, point_cloud, one_hot_vec):
'''
@author： Qiao
实例分割网络
notice：tensorflow是 NHWC pytorch为 NCHW 需要调整
Input:
point_cloud: TF tensor in shape (B,4,N)
frustum point clouds with XYZ and intensity in point channels
XYZs are in frustum coordinate
one_hot_vec: TF tensor in shape (B,3)
length-3 vectors indicating predicted object type
end_points: dict
Output:
logits: TF tensor in shape (B,2,N), scores for bkg/clutter and object
end_points: dict
'''
batch_size = point_cloud.get_shape()[0].value
num_point = point_cloud.get_shape()[2].value
# net = tf.expand_dims(point_cloud, 2)
net = torch.unsqueeze(point_cloud, 3)
net = self.get_instance_seg_1(net)
net = self.get_instance_seg_2(net)
point_feat = self.get_instance_seg_3(net)
net = self.get_instance_seg_4(point_feat)
net = self.get_instance_seg_5(net)
global_feat = self.get_instance_seg_pool_1(net)
# 把通道数拼起来 pytorch中为第二个
global_feat = torch.cat([global_feat, torch.unsqueeze(torch.unsqueeze(one_hot_vec, 1), 1)], 1)
global_feat_expand = torch.repeat(global_feat, [1, num_point, 1, 1])
concat_feat = torch.cat([point_feat, global_feat_expand],3)
net = self.get_instance_seg_6(concat_feat)
net = self.get_instance_seg_7(net)
net = self.get_instance_seg_8(net)
net = self.get_instance_seg_9(net)
net = self.get_instance_seg_dp_1(net)
logits = self.get_instance_seg_10(net)
logits = torch.squeeze(logits, 2) # BxCxN
return logits
def tfm_padtrim_signal(sig, width, pad_mode="zeros"):
'''Pad signal to specified width, using specified pad mode'''
c, x = sig.shape
pad_m = pad_mode.lower()
if (x == width): return sig
elif (x > width): return sig[:, :width]
elif pad_m in ["zeros", "zeros-after"]:
zeros_front = random.randint(0, width - x) if pad_m == "zeros" else 0
pad_front = torch.zeros((c, zeros_front))
pad_back = torch.zeros((c, width - x - zeros_front))
return torch.cat((pad_front, sig, pad_back), 1)
elif pad_m == "repeat":
repeats = width // x + 1
return torch.repeat(1, repeats)[:, :width]
else:
raise ValueError(
f"pad_mode {pad_m} not currently supported, only 'zeros', 'zeros-after', or 'repeat'"
)
def forward(self, pic_set):
hyperparams = self.hyperparams
simh = []
simpics = []
N = len(pic_set)
BS = pic_set[0].shape[0]
# forward pass
pic_set = torch.cat(pic_set, dim=0)
h, aux = self.encoder_c(pic_set)
if self._share_enc:
h = h.reshape(N, BS, -1)
h_mean = torch.mean(h, dim=0)
h = torch.repeat(h, N, dim=0)
input_i = [h, aux]
dec_pics = self.decoder(input_i)
# reshaping
pic_set = pic_set.reshape(N, BS, -1)
dec_pics = dec_pics.reshape(N, BS, -1)
# h = h.reshape(N,BS,-1)
# compute loss for each step
loss_pairs = 0
for i in range(N):
if i == 2: break
i_noisy = i
i_dec = (i + 1) % N
pic_pair = []
pic_pair.append(pic_set[i_noisy])
pic_pair.append(dec_pics[i_dec])
loss_pairs += self.compute_img_loss(pic_pair, proj=False)
loss_pairs = loss_pairs / len(pic_set)
# compute losses
# dec_pics = [dec_pic for dec_pic in dec_pics]
# loss_x = self.compute_img_loss(dec_pics,proj=False)
# h = [h_i for h_i in h]
# loss_h = self.compute_enc_loss(h,proj=False)
# loss = loss_pairs + hyperparams.x * loss_x + hyperparams.h * loss_h
loss = loss_pairs
return loss
def parse(self, det, tag, adjust=True, refine=True):
ans = self.match_torch(**self.top_k_torch(det, tag))
if adjust:
ans = self.adjust_torch(ans, det)
scores = [i[:, 2].mean() for i in ans[0]]
if refine:
# for each image
for i in range(len(ans)):
# for each detected person
for j in range(len(ans[i])):
det_ = det[i]
tag_ = tag[i]
ans_ = ans[i][j]
if not self.tag_per_joint:
tag_ = torch.repeat(tag_, (self.num_joints, 1, 1, 1))
ans[i][j] = self.refine_torch(det_, tag_, ans_)
# after this refinement step, there may be multiple detections with almost identical keypoints...
# an attempt to aggregate them is done afterwards in SimpleHigherHRNet
return ans, scores
def batch_global_rigid_transformation(Rs, Js, parent, rotate_base = False, betas_logscale=None, opts=None):
"""
Computes absolute joint locations given pose.
rotate_base: if True, rotates the global rotation by 90 deg in x axis.
if False, this is the original SMPL coordinate.
Args:
Rs: N x 24 x 3 x 3 rotation vector of K joints
Js: N x 24 x 3, joint locations before posing
parent: 24 holding the parent id for each index
Returns
new_J : `Tensor`: N x 24 x 3 location of absolute joints
A : `Tensor`: N x 24 4 x 4 relative joint transformations for LBS.
"""
if rotate_base:
print('Flipping the SMPL coordinate frame!!!!')
rot_x = torch.Tensor([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
rot_x = torch.reshape(torch.repeat(rot_x, [N, 1]), [N, 3, 3]) # In tf it was tile
root_rotation = torch.matmul(Rs[:, 0, :, :], rot_x)
else:
root_rotation = Rs[:, 0, :, :]
# Now Js is N x 24 x 3 x 1
Js = Js.unsqueeze(-1)
N = Rs.shape[0]
Js_orig = Js.clone()
scaling_factors = torch.ones(N, parent.shape[0], 3).to(Rs.device)
if betas_logscale is not None:
leg_joints = list(range(7,11)) + list(range(11,15)) + list(range(17,21)) + list(range(21,25))
tail_joints = list(range(25, 32))
ear_joints = [33, 34]
beta_scale_mask = torch.zeros(35, 3, 6).to(betas_logscale.device)
beta_scale_mask[leg_joints, [2], [0]] = 1.0 # Leg lengthening
beta_scale_mask[leg_joints, [0], [1]] = 1.0 # Leg fatness
beta_scale_mask[leg_joints, [1], [1]] = 1.0 # Leg fatness
beta_scale_mask[tail_joints, [0], [2]] = 1.0 # Tail lengthening
beta_scale_mask[tail_joints, [1], [3]] = 1.0 # Tail fatness
beta_scale_mask[tail_joints, [2], [3]] = 1.0 # Tail fatness
beta_scale_mask[ear_joints, [1], [4]] = 1.0 # Ear y
beta_scale_mask[ear_joints, [2], [5]] = 1.0 # Ear z
beta_scale_mask = torch.transpose(
beta_scale_mask.reshape(35*3, 6), 0, 1)
betas_scale = torch.exp(betas_logscale @ beta_scale_mask)
scaling_factors = betas_scale.reshape(-1, 35, 3)
scale_factors_3x3 = torch.diag_embed(scaling_factors, dim1=-2, dim2=-1)
def make_A(R, t):
# Rs is N x 3 x 3, ts is N x 3 x 1
R_homo = torch.nn.functional.pad(R, (0,0,0,1,0,0))
t_homo = torch.cat([t, torch.ones([N, 1, 1]).to(Rs.device)], 1)
return torch.cat([R_homo, t_homo], 2)
A0 = make_A(root_rotation, Js[:, 0])
results = [A0]
for i in range(1, parent.shape[0]):
j_here = Js[:, i] - Js[:, parent[i]]
s_par_inv = torch.inverse(scale_factors_3x3[:, parent[i]])
rot = Rs[:, i]
s = scale_factors_3x3[:, i]
rot_new = s_par_inv @ rot @ s
A_here = make_A(rot_new, j_here)
res_here = torch.matmul(
results[parent[i]], A_here)
results.append(res_here)
# 10 x 24 x 4 x 4
results = torch.stack(results, dim=1)
# scale updates
new_J = results[:, :, :3, 3]
# --- Compute relative A: Skinning is based on
# how much the bone moved (not the final location of the bone)
# but (final_bone - init_bone)
# ---
Js_w0 = torch.cat([Js_orig, torch.zeros([N, 35, 1, 1]).to(Rs.device)], 2)
init_bone = torch.matmul(results, Js_w0)
# Append empty 4 x 3:
init_bone = torch.nn.functional.pad(init_bone, (3,0,0,0,0,0,0,0))
A = results - init_bone
return new_J, A
def batch_global_rigid_transformation(Rs,
Js,
parent,
rotate_base=False,
betas_extra=None,
device="cuda"):
"""
Computes absolute joint locations given pose.
rotate_base: if True, rotates the global rotation by 90 deg in x axis.
if False, this is the original SMPL coordinate.
Args:
Rs: N x 24 x 3 x 3 rotation vector of K joints
Js: N x 24 x 3, joint locations before posing
parent: 24 holding the parent id for each index
Returns
new_J : `Tensor`: N x 24 x 3 location of absolute joints
A : `Tensor`: N x 24 4 x 4 relative joint transformations for LBS.
"""
# Now Js is N x 24 x 3 x 1
Js = Js.unsqueeze(-1)
N = Rs.shape[0]
if rotate_base:
print('Flipping the SMPL coordinate frame!!!!')
rot_x = torch.Tensor([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
rot_x = torch.reshape(torch.repeat(rot_x, [N, 1]),
[N, 3, 3]) # In tf it was tile
root_rotation = torch.matmul(Rs[:, 0, :, :], rot_x)
else:
root_rotation = Rs[:, 0, :, :]
Js_orig = Js.clone()
scaling_factors = torch.ones(N, parent.shape[0], 3).to(device)
if betas_extra is not None:
scaling_factors = betas_extra.reshape(-1, 35, 3)
# debug_only
# scaling_factors[:, 25:32, 0] = 0.2
# scaling_factors[:, 7, 2] = 2.0
scale_factors_3x3 = torch.diag_embed(scaling_factors, dim1=-2, dim2=-1)
def make_A(R, t):
# Rs is N x 3 x 3, ts is N x 3 x 1
R_homo = torch.nn.functional.pad(R, (0, 0, 0, 1, 0, 0))
t_homo = torch.cat([t, torch.ones([N, 1, 1]).to(device)], 1)
return torch.cat([R_homo, t_homo], 2)
A0 = make_A(root_rotation, Js[:, 0])
results = [A0]
for i in range(1, parent.shape[0]):
j_here = Js[:, i] - Js[:, parent[i]]
s_par_inv = torch.inverse(scale_factors_3x3[:, parent[i]])
rot = Rs[:, i]
s = scale_factors_3x3[:, i]
rot_new = s_par_inv @ rot @ s
A_here = make_A(rot_new, j_here)
res_here = torch.matmul(results[parent[i]], A_here)
results.append(res_here)
# 10 x 24 x 4 x 4
results = torch.stack(results, dim=1)
# scale updates
new_J = results[:, :, :3, 3]
# --- Compute relative A: Skinning is based on
# how much the bone moved (not the final location of the bone)
# but (final_bone - init_bone)
# ---
Js_w0 = torch.cat([Js_orig, torch.zeros([N, 35, 1, 1]).to(device)], 2)
init_bone = torch.matmul(results, Js_w0)
# Append empty 4 x 3:
init_bone = torch.nn.functional.pad(init_bone, (3, 0, 0, 0, 0, 0, 0, 0))
A = results - init_bone
return new_J, A
def repeat(*args, **kwargs):
return torch.repeat(*args, **kwargs)