def body_fitting_loss(body_pose,
betas,
model_joints,
camera_t,
camera_center,
joints_2d,
joints_conf,
pose_prior,
focal_length=5000,
sigma=100,
pose_prior_weight=4.78,
shape_prior_weight=5,
angle_prior_weight=15.2,
output='sum'):
"""
Loss function for body fitting
"""
# pose_prior_weight = 1.
# shape_prior_weight = 1.
# angle_prior_weight = 1.
# sigma = 10.
batch_size = body_pose.shape[0]
rotation = torch.eye(3, device=body_pose.device).unsqueeze(0).expand(
batch_size, -1, -1)
projected_joints = perspective_projection(model_joints, rotation, camera_t,
focal_length, camera_center)
# Weighted robust reprojection error
reprojection_error = gmof(projected_joints - joints_2d, sigma)
reprojection_loss = (joints_conf**2) * reprojection_error.sum(dim=-1)
# Pose prior loss
pose_prior_loss = (pose_prior_weight**2) * pose_prior(body_pose, betas)
# Angle prior for knees and elbows
angle_prior_loss = (angle_prior_weight**
2) * angle_prior(body_pose).sum(dim=-1)
# Regularizer to prevent betas from taking large values
shape_prior_loss = (shape_prior_weight**2) * (betas**2).sum(dim=-1)
total_loss = reprojection_loss.sum(
dim=-1) + pose_prior_loss + angle_prior_loss + shape_prior_loss
print(f'joints: {reprojection_loss[0].sum().item():.2f}, '
f'pose_prior: {pose_prior_loss[0].item():.2f}, '
f'angle_prior: {angle_prior_loss[0].item():.2f}, '
f'shape_prior: {shape_prior_loss[0].item():.2f}')
if output == 'sum':
return total_loss.sum()
elif output == 'reprojection':
return reprojection_loss
def camera_fitting_loss(model_joints,
camera_t,
camera_t_est,
camera_center,
joints_2d,
joints_conf,
focal_length=5000,
depth_loss_weight=100):
"""
Loss function for camera optimization.
"""
# Project model joints
batch_size = model_joints.shape[0]
rotation = torch.eye(3, device=model_joints.device).unsqueeze(0).expand(
batch_size, -1, -1)
projected_joints = perspective_projection(model_joints, rotation, camera_t,
focal_length, camera_center)
op_joints = ['OP RHip', 'OP LHip', 'OP RShoulder', 'OP LShoulder']
op_joints_ind = [JOINT_IDS[joint] for joint in op_joints]
gt_joints = ['Right Hip', 'Left Hip', 'Right Shoulder', 'Left Shoulder']
gt_joints_ind = [JOINT_IDS[joint] for joint in gt_joints]
reprojection_error_op = (joints_2d[:, op_joints_ind] -
projected_joints[:, op_joints_ind])**2
reprojection_error_gt = (joints_2d[:, gt_joints_ind] -
projected_joints[:, gt_joints_ind])**2
# Check if for each example in the batch all 4 OpenPose detections are valid, otherwise use the GT detections
# OpenPose joints are more reliable for this task, so we prefer to use them if possible
is_valid = (joints_conf[:, op_joints_ind].min(dim=-1)[0][:, None, None] >
0).float()
reprojection_loss = (is_valid * reprojection_error_op +
(1 - is_valid) * reprojection_error_gt).sum(dim=(1,
2))
# Loss that penalizes deviation from depth estimate
depth_loss = (depth_loss_weight**
2) * (camera_t[:, 2] - camera_t_est[:, 2])**2
total_loss = reprojection_loss + depth_loss
return total_loss.sum()
def project_3d_to_2d(joints_3d, bbox, pred_cam):
joints_3d_tensor = torch.from_numpy(joints_3d).unsqueeze(0)
focal_length = 5000
camera_center = torch.zeros(2).unsqueeze(0)
rotation = torch.eye(3).unsqueeze(0)
pred_camera = torch.from_numpy(pred_cam).unsqueeze(0) # shape: (1, 3)
pred_cam_t = torch.stack([
pred_camera[:, 1], pred_camera[:, 2], 2 * 5000. /
(224. * pred_camera[:, 0] + 1e-9)
],
dim=-1)
joints_2d = perspective_projection(joints_3d_tensor, rotation, pred_cam_t,
focal_length,
camera_center).squeeze().numpy()
# scale back to the original image
scaled_joints_2d = joints_2d * bbox[2:] / 224 + bbox[:2]
scaled_joints_2d = scaled_joints_2d.astype(int)
return scaled_joints_2d
joints_3d_body_center = (joints_3d_shoulder_left + joints_3d_shoulder_right
+ joints_3d_hip_middle) / 3
body_orientation_starting_point = joints_3d_body_center
joints_3d_body_direction = body_orientation_starting_point - np.cross(
(joints_3d_shoulder_left - joints_3d_hip_middle),
(joints_3d_shoulder_right - joints_3d_hip_middle)).reshape(1, 3)
joints_3d = np.vstack(
(joints_3d, body_orientation_starting_point.reshape(1, 3),
joints_3d_body_direction))
joints_3d_tensor = torch.from_numpy(joints_3d).unsqueeze(0)
joints_2d = perspective_projection(joints_3d_tensor, rotation, pred_cam_t,
focal_length,
camera_center).squeeze().numpy()
# extraction of head orientation
mid_front = (joints_2d[15:16] + joints_2d[16:17]) / 2 # (Leye + Reye) / 2
mid_back = (joints_2d[42:43] + joints_2d[43:44]) / 2 # (Jaw + Head) / 2
joints_2d = np.vstack((joints_2d, mid_front, mid_back))
# scale back to the original image
scaled_joints_2d = joints_2d * bbox[2:] / 224 + bbox[:2]
scaled_joints_2d = scaled_joints_2d.astype(int)
for i in range(len(scaled_joints_2d)):
img = cv2.circle(img, (scaled_joints_2d[i, 0], scaled_joints_2d[i, 1]),
2, (0, 0, 255), -1)
def temporal_body_fitting_loss(body_pose,
betas,
model_joints,
camera_t,
camera_center,
joints_2d,
joints_conf,
pose_prior,
focal_length=5000,
sigma=100,
pose_prior_weight=4.78,
shape_prior_weight=5,
angle_prior_weight=15.2,
smooth_2d_weight=0.01,
smooth_3d_weight=1.0,
output='sum'):
"""
Loss function for body fitting
"""
# pose_prior_weight = 1.
# shape_prior_weight = 1.
# angle_prior_weight = 1.
# sigma = 10.
batch_size = body_pose.shape[0]
rotation = torch.eye(3, device=body_pose.device).unsqueeze(0).expand(
batch_size, -1, -1)
projected_joints = perspective_projection(model_joints, rotation, camera_t,
focal_length, camera_center)
# Weighted robust reprojection error
reprojection_error = gmof(projected_joints - joints_2d, sigma)
reprojection_loss = (joints_conf**2) * reprojection_error.sum(dim=-1)
# Pose prior loss
pose_prior_loss = (pose_prior_weight**2) * pose_prior(body_pose, betas)
# Angle prior for knees and elbows
angle_prior_loss = (angle_prior_weight**
2) * angle_prior(body_pose).sum(dim=-1)
# Regularizer to prevent betas from taking large values
shape_prior_loss = (shape_prior_weight**2) * (betas**2).sum(dim=-1)
total_loss = reprojection_loss.sum(
dim=-1) + pose_prior_loss + angle_prior_loss + shape_prior_loss
# Smooth 2d joint loss
joint_conf_diff = joints_conf[1:]
joints_2d_diff = projected_joints[1:] - projected_joints[:-1]
smooth_j2d_loss = (joint_conf_diff**2) * joints_2d_diff.abs().sum(dim=-1)
smooth_j2d_loss = torch.cat([
torch.zeros(1, smooth_j2d_loss.shape[1], device=body_pose.device),
smooth_j2d_loss
]).sum(dim=-1)
smooth_j2d_loss = (smooth_2d_weight**2) * smooth_j2d_loss
# Smooth 3d joint loss
joints_3d_diff = model_joints[1:] - model_joints[:-1]
# joints_3d_diff = joints_3d_diff * 100.
smooth_j3d_loss = (joint_conf_diff**2) * joints_3d_diff.abs().sum(dim=-1)
smooth_j3d_loss = torch.cat([
torch.zeros(1, smooth_j3d_loss.shape[1], device=body_pose.device),
smooth_j3d_loss
]).sum(dim=-1)
smooth_j3d_loss = (smooth_3d_weight**2) * smooth_j3d_loss
total_loss += smooth_j2d_loss + smooth_j3d_loss
# print(f'joints: {reprojection_loss[0].sum().item():.2f}, '
# f'pose_prior: {pose_prior_loss[0].item():.2f}, '
# f'angle_prior: {angle_prior_loss[0].item():.2f}, '
# f'shape_prior: {shape_prior_loss[0].item():.2f}, '
# f'smooth_j2d: {smooth_j2d_loss.sum().item()}, '
# f'smooth_j3d: {smooth_j3d_loss.sum().item()}')
if output == 'sum':
return total_loss.sum()
elif output == 'reprojection':
return reprojection_loss
