jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def calc_updown(self, orig_weight):
        oft_blocks = self.oft_blocks.to(orig_weight.device)
        eye = torch.eye(self.block_size, device=oft_blocks.device)

        if not self.is_R:
            block_Q = oft_blocks - oft_blocks.transpose(-1, -2) # ensure skew-symmetric orthogonal matrix
            if self.constraint != 0:
                norm_Q = torch.norm(block_Q.flatten())
                new_norm_Q = torch.clamp(norm_Q, max=self.constraint.to(oft_blocks.device))
                block_Q = block_Q * ((new_norm_Q + 1e-8) / (norm_Q + 1e-8))
            oft_blocks = torch.matmul(eye + block_Q, (eye - block_Q).float().inverse())

        R = oft_blocks.to(orig_weight.device)

        if not self.is_boft:
            # This errors out for MultiheadAttention, might need to be handled up-stream
            merged_weight = rearrange(orig_weight, '(k n) ... -> k n ...', k=self.num_blocks, n=self.block_size)
            merged_weight = torch.einsum(
                'k n m, k n ... -> k m ...',
                R,
                merged_weight
            )
            merged_weight = rearrange(merged_weight, 'k m ... -> (k m) ...')
        else:
            # TODO: determine correct value for scale
            scale = 1.0
            m = self.boft_m
            b = self.boft_b
            r_b = b // 2
            inp = orig_weight
            for i in range(m):
                bi = R[i] # b_num, b_size, b_size
                if i == 0:
                    # Apply multiplier/scale and rescale into first weight
                    bi = bi * scale + (1 - scale) * eye
                inp = rearrange(inp, "(c g k) ... -> (c k g) ...", g=2, k=2**i * r_b)
                inp = rearrange(inp, "(d b) ... -> d b ...", b=b)
                inp = torch.einsum("b i j, b j ... -> b i ...", bi, inp)
                inp = rearrange(inp, "d b ... -> (d b) ...")
                inp = rearrange(inp, "(c k g) ... -> (c g k) ...", g=2, k=2**i * r_b)
            merged_weight = inp

        # Rescale mechanism
        if self.rescale is not None:
            merged_weight = self.rescale.to(merged_weight) * merged_weight

        updown = merged_weight.to(orig_weight.device) - orig_weight.to(merged_weight.dtype)
        output_shape = orig_weight.shape
        return self.finalize_updown(updown, orig_weight, output_shape)
//  [python]
def calc_updown(self, orig_weight):
        oft_blocks = self.oft_blocks.to(orig_weight.device)
        eye = torch.eye(self.block_size, device=oft_blocks.device)

        if not self.is_R:
            block_Q = oft_blocks - oft_blocks.transpose(-1, -2) # ensure skew-symmetric orthogonal matrix
            if self.constraint != 0:
                norm_Q = torch.linalg.norm(block_Q.flatten())
                new_norm_Q = torch.clamp(norm_Q, max=self.constraint.to(oft_blocks.device))
                block_Q = block_Q * ((new_norm_Q + 1e-8) / (norm_Q + 1e-8))
            oft_blocks = torch.matmul(eye + block_Q, (eye - block_Q).float().inverse())

        R = oft_blocks.to(orig_weight.device)

        if not self.is_boft:
            # This errors out for MultiheadAttention, might need to be handled up-stream
            merged_weight = rearrange(orig_weight, '(k n) ... -> k n ...', k=self.num_blocks, n=self.block_size)
            merged_weight = torch.einsum(
                'k n m, k n ... -> k m ...',
                R,
                merged_weight
            )
            merged_weight = rearrange(merged_weight, 'k m ... -> (k m) ...')
        else:
            # TODO: determine correct value for scale
            scale = 1.0
            m = self.boft_m
            b = self.boft_b
            r_b = b // 2
            inp = orig_weight
            for i in range(m):
                bi = R[i] # b_num, b_size, b_size
                if i == 0:
                    # Apply multiplier/scale and rescale into first weight
                    bi = bi * scale + (1 - scale) * eye
                inp = rearrange(inp, "(c g k) ... -> (c k g) ...", g=2, k=2**i * r_b)
                inp = rearrange(inp, "(d b) ... -> d b ...", b=b)
                inp = torch.einsum("b i j, b j ... -> b i ...", bi, inp)
                inp = rearrange(inp, "d b ... -> (d b) ...")
                inp = rearrange(inp, "(c k g) ... -> (c g k) ...", g=2, k=2**i * r_b)
            merged_weight = inp

        # Rescale mechanism
        if self.rescale is not None:
            merged_weight = self.rescale.to(merged_weight) * merged_weight

        updown = merged_weight.to(orig_weight.device) - orig_weight.to(merged_weight.dtype)
        output_shape = orig_weight.shape
        return self.finalize_updown(updown, orig_weight, output_shape)
//  [python]
def _l1_dist(edges):
    # formula 2
    ed = th.norm(edges.src["nd"] - edges.dst["nd"], 1, 1)
    return {"ed": ed}
//  [python]
def _l1_dist(edges):
    ed = th.linalg.norm(edges.src["nd"] - edges.dst["nd"], 1, 1)
    return {"ed": ed}
//  [python]
def forward(self, input_data):
        batch, _, n_frames = input_data.shape
        if self.take_log:
            input_data = torch.log(input_data + self.EPS)
        out = self.rnn(input_data.permute(0, 2, 1))
        out = self.dropout(out)
        proj = self.embedding_layer(out)
        proj = self.embedding_act(proj)
        proj = proj.view(batch, n_frames, -1, self.embedding_dim).transpose(1, 2)
        proj = proj.reshape(batch, -1, self.embedding_dim)
        proj_norm = torch.norm(proj, p=2, dim=-1, keepdim=True)
        projection_final = proj / (proj_norm + self.EPS)
        mask_out = self.mask_layer(out).view(batch, n_frames, self.n_src, self.input_dim)
        mask_out = mask_out.permute(0, 2, 3, 1)
        mask_out = self.mask_act(mask_out)
        return projection_final, mask_out
//  [python]
def forward(self, input_data):
        batch, _, n_frames = input_data.shape
        if self.take_log:
            input_data = torch.log(input_data + self.EPS)
        out = self.rnn(input_data.permute(0, 2, 1))
        out = self.dropout(out)
        proj = self.embedding_layer(out)
        proj = self.embedding_act(proj)
        proj = proj.view(batch, n_frames, -1, self.embedding_dim).transpose(1, 2)
        proj = proj.reshape(batch, -1, self.embedding_dim)
        proj_norm = torch.linalg.norm(proj, ord=2, dim=-1, keepdim=True)
        projection_final = proj / (proj_norm + self.EPS)
        mask_out = self.mask_layer(out).view(batch, n_frames, self.n_src, self.input_dim)
        mask_out = mask_out.permute(0, 2, 3, 1)
        mask_out = self.mask_act(mask_out)
        return projection_final, mask_out
//  [python]
def batch_l2_loss(real, predict):
    loss_batch = torch.norm(real - predict, p=2, dim=1)
    return loss_batch.mean()
//  [python]
def batch_l2_loss(real, predict):
    loss_batch = torch.linalg.norm(real - predict, ord=2, dim=1)
    return loss_batch.mean()
//  [python]
def magnitude(self) -> Tuple[torch.Tensor]:
        '''Magnitude of the SE(3) transformation.
        Returns:
            dr: rotation angle in degrees.
            dt: translation distance in meters.
        '''
        trace = torch.diagonal(self.R, dim1=-1, dim2=-2).sum(-1)
        cos = torch.clamp((trace - 1) / 2, -1, 1)
        dr = torch.acos(cos).abs() / math.pi * 180
        dt = torch.norm(self.t, dim=-1)
        return dr, dt
//  [python]
def magnitude(self) -> Tuple[torch.Tensor]:
        '''Magnitude of the SE(3) transformation.
        Returns:
            dr: rotation angle in degrees.
            dt: translation distance in meters.
        '''
        trace = torch.diagonal(self.R, dim1=-1, dim2=-2).sum(-1)
        cos = torch.clamp((trace - 1) / 2, -1, 1)
        dr = torch.acos(cos).abs() / math.pi * 180
        dt = torch.linalg.norm(self.t, dim=-1)
        return dr, dt
//  [python]
def batch_smpl_pose_l2_error(real, predict):
    batch_size = real.shape[0]
    real = batch_rodrigues(real.reshape(-1, 3)).contiguous()
    predict = batch_rodrigues(predict.reshape(-1, 3)).contiguous()
    loss = torch.norm((real - predict).view(-1, 9), p=2, dim=-1)
    loss = loss.reshape(batch_size, -1).mean(-1)
    return loss
//  [python]
def batch_smpl_pose_l2_error(real, predict):
    batch_size = real.shape[0]
    real = batch_rodrigues(real.reshape(-1, 3)).contiguous()
    predict = batch_rodrigues(predict.reshape(-1, 3)).contiguous()
    loss = torch.linalg.norm((real - predict).view(-1, 9), ord=2, dim=-1)
    loss = loss.reshape(batch_size, -1).mean(-1)
    return loss
//  [python]
def batch_rodrigues(axisang):
    axisang_norm = torch.norm(axisang + 1e-08, p=2, dim=1)
    angle = torch.unsqueeze(axisang_norm, -1)
    axisang_normalized = torch.div(axisang, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * axisang_normalized], dim=1)
    rot_mat = quat2mat(quat)
    rot_mat = rot_mat.view(rot_mat.shape[0], 9)
    return rot_mat
//  [python]
def batch_rodrigues(axisang):
    axisang_norm = torch.linalg.norm(axisang + 1e-08, ord=2, dim=1)
    angle = torch.unsqueeze(axisang_norm, -1)
    axisang_normalized = torch.div(axisang, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * axisang_normalized], dim=1)
    rot_mat = quat2mat(quat)
    rot_mat = rot_mat.view(rot_mat.shape[0], 9)
    return rot_mat
//  [python]
def _calc_matched_PCKh_(real, pred, kp2d_mask, error_thresh=0.143):
    PCKs = torch.ones(len(kp2d_mask)).float() * -1.0
    if kp2d_mask.sum() > 0:
        vis = (real > -1.0).sum(-1) == real.shape[-1]
        error = torch.norm(real - pred, p=2, dim=-1)
        for ind, (e, v) in enumerate(zip(error, vis)):
            if v.sum() < 2:
                continue
            real_valid = real[ind, v]
            person_scales = torch.sqrt((real_valid[:, 0].max(-1).values - real_valid[:, 0].min(-1).values) ** 2 + (real_valid[:, 1].max(-1).values - real_valid[:, 1].min(-1).values) ** 2)
            error_valid = e[v]
            correct_kp_mask = (error_valid / person_scales < error_thresh).float()
            PCKs[ind] = correct_kp_mask.sum() / len(correct_kp_mask)
    return PCKs
//  [python]
def _calc_matched_PCKh_(real, pred, kp2d_mask, error_thresh=0.143):
    PCKs = torch.ones(len(kp2d_mask)).float() * -1.0
    if kp2d_mask.sum() > 0:
        vis = (real > -1.0).sum(-1) == real.shape[-1]
        error = torch.linalg.norm(real - pred, p=2, dim=-1)
        for ind, (e, v) in enumerate(zip(error, vis)):
            if v.sum() < 2:
                continue
            real_valid = real[ind, v]
            person_scales = torch.sqrt((real_valid[:, 0].max(-1).values - real_valid[:, 0].min(-1).values) ** 2 + (real_valid[:, 1].max(-1).values - real_valid[:, 1].min(-1).values) ** 2)
            error_valid = e[v]
            correct_kp_mask = (error_valid / person_scales < error_thresh).float()
            PCKs[ind] = correct_kp_mask.sum() / len(correct_kp_mask)
    return PCKs
//  [python]
def batch_kp_2d_l2_loss(real, pred):
    """ 
    Directly supervise the 2D coordinates of global joints, like torso
    While supervise the relative 2D coordinates of part joints, like joints on face, feets
    """
    vis_mask = ((real > -1.99).sum(-1) == real.shape[-1]).float()
    for parent_joint, leaf_joints in constants.joint2D_tree.items():
        parent_id = constants.SMPL_ALL_54[parent_joint]
        leaf_ids = np.array([constants.SMPL_ALL_54[leaf_joint] for leaf_joint in leaf_joints])
        vis_mask[:, leaf_ids] = vis_mask[:, [parent_id]] * vis_mask[:, leaf_ids]
        real[:, leaf_ids] -= real[:, [parent_id]]
        pred[:, leaf_ids] -= pred[:, [parent_id]]
    bv_mask = torch.logical_and(vis_mask.sum(-1) > 0, (real - pred).sum(-1).sum(-1) != 0)
    vis_mask = vis_mask[bv_mask]
    loss = 0
    if vis_mask.sum() > 0:
        diff = torch.norm(real[bv_mask] - pred[bv_mask], p=2, dim=-1)
        loss = (diff * vis_mask).sum(-1) / (vis_mask.sum(-1) + 0.0001)
        if torch.isnan(loss).sum() > 0 or (loss > 1000).sum() > 0:
            return 0
            None
            non_position = torch.isnan(loss)
            None
            return 0
    return loss
//  [python]
def batch_kp_2d_l2_loss(real, pred):
    """ 
    Directly supervise the 2D coordinates of global joints, like torso
    While supervise the relative 2D coordinates of part joints, like joints on face, feets
    """
    vis_mask = ((real > -1.99).sum(-1) == real.shape[-1]).float()
    for parent_joint, leaf_joints in constants.joint2D_tree.items():
        parent_id = constants.SMPL_ALL_54[parent_joint]
        leaf_ids = np.array([constants.SMPL_ALL_54[leaf_joint] for leaf_joint in leaf_joints])
        vis_mask[:, leaf_ids] = vis_mask[:, [parent_id]] * vis_mask[:, leaf_ids]
        real[:, leaf_ids] -= real[:, [parent_id]]
        pred[:, leaf_ids] -= pred[:, [parent_id]]
    bv_mask = torch.logical_and(vis_mask.sum(-1) > 0, (real - pred).sum(-1).sum(-1) != 0)
    vis_mask = vis_mask[bv_mask]
    loss = 0
    if vis_mask.sum() > 0:
        diff = torch.linalg.norm(real[bv_mask] - pred[bv_mask], p=2, dim=-1)
        loss = (diff * vis_mask).sum(-1) / (vis_mask.sum(-1) + 0.0001)
        if torch.isnan(loss).sum() > 0 or (loss > 1000).sum() > 0:
            return 0
            None
            non_position = torch.isnan(loss)
            None
            return 0
    return loss
//  [python]
def remove_outlier(outputs, relative_scale_thresh=3, scale_thresh=0.25):
    cam_trans = outputs['cam_trans']
    N = len(cam_trans)
    if N < 3:
        return outputs
    trans_diff = cam_trans.unsqueeze(1).repeat(1, N, 1) - cam_trans.unsqueeze(0).repeat(N, 1, 1)
    trans_dist_mat = torch.norm(trans_diff, p=2, dim=-1)
    trans_dist_mat = torch.sort(trans_dist_mat).values[:, 1:-1]
    mean_dist = trans_dist_mat.mean(1)
    relative_scale = mean_dist / ((mean_dist.sum() - mean_dist) / (N - 1))
    outlier_mask = relative_scale > relative_scale_thresh
    outlier_mask *= outputs['cam'][:, 0] < scale_thresh
    removed_subj_inds = torch.where(outlier_mask)[0]
    if len(removed_subj_inds) > 0:
        outputs = remove_subjects(outputs, removed_subj_inds)
    return outputs
//  [python]
def remove_outlier(outputs, relative_scale_thresh=3, scale_thresh=0.25):
    cam_trans = outputs['cam_trans']
    N = len(cam_trans)
    if N < 3:
        return outputs
    trans_diff = cam_trans.unsqueeze(1).repeat(1, N, 1) - cam_trans.unsqueeze(0).repeat(N, 1, 1)
    trans_dist_mat = torch.linalg.norm(trans_diff, p=2, dim=-1)
    trans_dist_mat = torch.sort(trans_dist_mat).values[:, 1:-1]
    mean_dist = trans_dist_mat.mean(1)
    relative_scale = mean_dist / ((mean_dist.sum() - mean_dist) / (N - 1))
    outlier_mask = relative_scale > relative_scale_thresh
    outlier_mask *= outputs['cam'][:, 0] < scale_thresh
    removed_subj_inds = torch.where(outlier_mask)[0]
    if len(removed_subj_inds) > 0:
        outputs = remove_subjects(outputs, removed_subj_inds)
    return outputs
//  [python]
def _calc_joint_angle_error(pred_mat, gt_mat, return_axis_angle=False):
    """
    Compute the geodesic distance between the two input matrices.
    :param pred_mat: predicted rotation matrices. Shape: ( Seq, 9g, 3, 3)
    :param gt_mat: ground truth rotation matrices. Shape: ( Seq, 9, 3, 3)
    :return: Mean geodesic distance between input matrices.
    """
    r1 = pred_mat.reshape(-1, 3, 3)
    r2 = gt_mat.reshape(-1, 3, 3)
    r2t = r2.permute(0, 2, 1)
    r = torch.matmul(r1, r2t)
    axis_angles = rotation_matrix_to_angle_axis(r)
    angles = torch.norm(axis_angles, dim=-1) * (180.0 / np.pi)
    if return_axis_angle:
        return angles, axis_angles
    return angles
//  [python]
def _calc_joint_angle_error(pred_mat, gt_mat, return_axis_angle=False):
    """
    Compute the geodesic distance between the two input matrices.
    :param pred_mat: predicted rotation matrices. Shape: ( Seq, 9g, 3, 3)
    :param gt_mat: ground truth rotation matrices. Shape: ( Seq, 9, 3, 3)
    :return: Mean geodesic distance between input matrices.
    """
    r1 = pred_mat.reshape(-1, 3, 3)
    r2 = gt_mat.reshape(-1, 3, 3)
    r2t = r2.permute(0, 2, 1)
    r = torch.matmul(r1, r2t)
    axis_angles = rotation_matrix_to_angle_axis(r)
    angles = torch.linalg.norm(axis_angles, dim=-1) * (180.0 / np.pi)
    if return_axis_angle:
        return angles, axis_angles
    return angles
//  [python]
def suppressing_silimar_mesh_and_2D_center(self, params_preds, pred_batch_ids, pred_czyxs, top_score, center2D_thresh=5, pose_thresh=2.5):
        pose_params_preds = params_preds[:, args().cam_dim:args().cam_dim + 22 * args().rot_dim]
        N = len(pred_czyxs)
        center2D_similarity = torch.norm((pred_czyxs[:, 1:].unsqueeze(1).repeat(1, N, 1) - pred_czyxs[:, 1:].unsqueeze(0).repeat(N, 1, 1)).float(), p=2, dim=-1)
        same_batch_id_mask = pred_batch_ids.unsqueeze(1).repeat(1, N) == pred_batch_ids.unsqueeze(0).repeat(N, 1)
        center2D_similarity[~same_batch_id_mask] = center2D_thresh + 1
        similarity = center2D_similarity <= center2D_thresh
        center_similar_inds = torch.where(similarity.sum(-1) > 1)[0]
        for s_inds in center_similar_inds:
            pose_angulars = rot6D_to_angular(pose_params_preds[similarity[s_inds]])
            pose_angular_base = rot6D_to_angular(pose_params_preds[s_inds].unsqueeze(0)).repeat(len(pose_angulars), 1)
            pose_similarity = batch_smpl_pose_l2_error(pose_angulars, pose_angular_base)
            sim_past = similarity[s_inds].clone()
            similarity[s_inds, sim_past] = pose_similarity < pose_thresh
        score_map = similarity * top_score.unsqueeze(0).repeat(N, 1)
        nms_inds = torch.argmax(score_map, 1) == torch.arange(N)
        return [item[nms_inds] for item in [pred_batch_ids, pred_czyxs, top_score]], nms_inds
//  [python]
def suppressing_silimar_mesh_and_2D_center(self, params_preds, pred_batch_ids, pred_czyxs, top_score, center2D_thresh=5, pose_thresh=2.5):
        pose_params_preds = params_preds[:, args().cam_dim:args().cam_dim + 22 * args().rot_dim]
        N = len(pred_czyxs)
        center2D_similarity = torch.linalg.norm((pred_czyxs[:, 1:].unsqueeze(1).repeat(1, N, 1) - pred_czyxs[:, 1:].unsqueeze(0).repeat(N, 1, 1)).float(), ord=2, dim=-1)
        same_batch_id_mask = pred_batch_ids.unsqueeze(1).repeat(1, N) == pred_batch_ids.unsqueeze(0).repeat(N, 1)
        center2D_similarity[~same_batch_id_mask] = center2D_thresh + 1
        similarity = center2D_similarity <= center2D_thresh
        center_similar_inds = torch.where(similarity.sum(-1) > 1)[0]
        for s_inds in center_similar_inds:
            pose_angulars = rot6D_to_angular(pose_params_preds[similarity[s_inds]])
            pose_angular_base = rot6D_to_angular(pose_params_preds[s_inds].unsqueeze(0)).repeat(len(pose_angulars), 1)
            pose_similarity = batch_smpl_pose_l2_error(pose_angulars, pose_angular_base)
            sim_past = similarity[s_inds].clone()
            similarity[s_inds, sim_past] = pose_similarity < pose_thresh
        score_map = similarity * top_score.unsqueeze(0).repeat(N, 1)
        nms_inds = torch.argmax(score_map, 1) == torch.arange(N)
        return [item[nms_inds] for item in [pred_batch_ids, pred_czyxs, top_score]], nms_inds
//  [python]
def forward(self, joint_out, joint_gt):
        joint_out = joint_out.reshape(-1, self.joint_num, 2)
        joint_gt = joint_gt.reshape(-1, self.joint_num, 2)
        mask = joint_gt > 0
        joint_out = joint_out * mask
        joint_gt = joint_gt * mask
        J = torch.norm(joint_out[:,self.id_i,:] - joint_out[:,self.id_j,:], p=2, dim=-1, keepdim=False)
        Y = torch.norm(joint_gt[:,self.id_i,:] - joint_gt[:,self.id_j,:], p=2, dim=-1, keepdim=False)
        loss = torch.abs(J-Y)
        return loss.mean()
//  [python]
def forward(self, joint_out, joint_gt):
        joint_out = joint_out.reshape(-1, self.joint_num, 2)
        joint_gt = joint_gt.reshape(-1, self.joint_num, 2)
        mask = joint_gt > 0
        joint_out = joint_out * mask
        joint_gt = joint_gt * mask
        J = torch.linalg.norm(joint_out[:,self.id_i,:] - joint_out[:,self.id_j,:], ord=2, dim=-1, keepdim=False)
        Y = torch.linalg.norm(joint_gt[:,self.id_i,:] - joint_gt[:,self.id_j,:], ord=2, dim=-1, keepdim=False)
        loss = torch.abs(J-Y)
        return loss.mean()
//  [python]
def match_gt_pred(self, center_gts_info, center_preds_info, device, is_training):
        vgt_batch_ids, vgt_person_ids, vgt_centers = center_gts_info
        vpred_batch_ids, flat_inds, cyxs, top_score = center_preds_info
        mc = {key: [] for key in ['batch_ids', 'flat_inds', 'person_ids', 'conf']}
        if self.match_preds_to_gts_for_supervision:
            for match_ind in torch.arange(len(vgt_batch_ids)):
                batch_id, person_id, center_gt = vgt_batch_ids[match_ind], vgt_person_ids[match_ind], vgt_centers[match_ind]
                pids = torch.where(vpred_batch_ids == batch_id)[0]
                if len(pids) == 0:
                    continue
                closet_center_ind = pids[torch.argmin(torch.norm(cyxs[pids].float() - center_gt[None].float(), dim=-1))]
                center_matched = cyxs[closet_center_ind].long()
                cy, cx = torch.clamp(center_matched, 0, self.map_size - 1)
                flat_ind = cy * args().centermap_size + cx
                mc['batch_ids'].append(batch_id)
                mc['flat_inds'].append(flat_ind)
                mc['person_ids'].append(person_id)
                mc['conf'].append(top_score[closet_center_ind])
            keys_list = list(mc.keys())
            for key in keys_list:
                if key != 'conf':
                    mc[key] = torch.Tensor(mc[key]).long()
                if args().max_supervise_num != -1 and is_training:
                    mc[key] = mc[key][:args().max_supervise_num]
        else:
            mc['batch_ids'] = vgt_batch_ids.long()
            mc['flat_inds'] = flatten_inds(vgt_centers.long())
            mc['person_ids'] = vgt_person_ids.long()
            mc['conf'] = torch.zeros(len(vgt_person_ids))
        return mc
//  [python]
def match_gt_pred(self, center_gts_info, center_preds_info, device, is_training):
        vgt_batch_ids, vgt_person_ids, vgt_centers = center_gts_info
        vpred_batch_ids, flat_inds, cyxs, top_score = center_preds_info
        mc = {key: [] for key in ['batch_ids', 'flat_inds', 'person_ids', 'conf']}
        if self.match_preds_to_gts_for_supervision:
            for match_ind in torch.arange(len(vgt_batch_ids)):
                batch_id, person_id, center_gt = vgt_batch_ids[match_ind], vgt_person_ids[match_ind], vgt_centers[match_ind]
                pids = torch.where(vpred_batch_ids == batch_id)[0]
                if len(pids) == 0:
                    continue
                closet_center_ind = pids[torch.argmin(torch.linalg.norm(cyxs[pids].float() - center_gt[None].float(), dim=-1))]
                center_matched = cyxs[closet_center_ind].long()
                cy, cx = torch.clamp(center_matched, 0, self.map_size - 1)
                flat_ind = cy * args().centermap_size + cx
                mc['batch_ids'].append(batch_id)
                mc['flat_inds'].append(flat_ind)
                mc['person_ids'].append(person_id)
                mc['conf'].append(top_score[closet_center_ind])
            keys_list = list(mc.keys())
            for key in keys_list:
                if key != 'conf':
                    mc[key] = torch.Tensor(mc[key]).long()
                if args().max_supervise_num != -1 and is_training:
                    mc[key] = mc[key][:args().max_supervise_num]
        else:
            mc['batch_ids'] = vgt_batch_ids.long()
            mc['flat_inds'] = flatten_inds(vgt_centers.long())
            mc['person_ids'] = vgt_person_ids.long()
            mc['conf'] = torch.zeros(len(vgt_person_ids))
        return mc
//  [python]
def suppressing_redundant_prediction_via_projection(outputs, img_shape, thresh=16, conf_based=False):
    pj2ds = outputs['pj2d']
    N = len(pj2ds)
    if N == 1:
        return outputs
    pj2d_diff = pj2ds.unsqueeze(1).repeat(1, N, 1, 1) - pj2ds.unsqueeze(0).repeat(N, 1, 1, 1)
    pj2d_dist_mat = torch.norm(pj2d_diff, p=2, dim=-1).mean(-1)
    person_scales = outputs['cam'][:, 0] * 2
    ps1, ps2 = person_scales.unsqueeze(1).repeat(1, N), person_scales.unsqueeze(0).repeat(N, 1)
    max_scale_mat = torch.where(ps1 > ps2, ps1, ps2)
    pj2d_dist_mat_normalized = pj2d_dist_mat / max_scale_mat
    triu_mask = torch.triu(torch.ones_like(pj2d_dist_mat), diagonal=1) < 0.5
    pj2d_dist_mat_normalized[triu_mask] = 10000.0
    max_length = max(img_shape)
    thresh = thresh * max_length / 640
    repeat_subj_inds = torch.where(pj2d_dist_mat_normalized < thresh)
    if len(repeat_subj_inds) > 0:
        if conf_based:
            center_confs = outputs['center_confs']
            removed_subj_inds = torch.where(center_confs[repeat_subj_inds[0]] < center_confs[repeat_subj_inds[1]], repeat_subj_inds[0], repeat_subj_inds[1])
        else:
            removed_subj_inds = torch.where(person_scales[repeat_subj_inds[0]] < person_scales[repeat_subj_inds[1]], repeat_subj_inds[0], repeat_subj_inds[1])
        outputs = remove_subjects(outputs, removed_subj_inds)
    return outputs
//  [python]
def suppressing_redundant_prediction_via_projection(outputs, img_shape, thresh=16, conf_based=False):
    pj2ds = outputs['pj2d']
    N = len(pj2ds)
    if N == 1:
        return outputs
    pj2d_diff = pj2ds.unsqueeze(1).repeat(1, N, 1, 1) - pj2ds.unsqueeze(0).repeat(N, 1, 1, 1)
    pj2d_dist_mat = torch.linalg.norm(pj2d_diff, ord=2, dim=-1).mean(-1)
    person_scales = outputs['cam'][:, 0] * 2
    ps1, ps2 = person_scales.unsqueeze(1).repeat(1, N), person_scales.unsqueeze(0).repeat(N, 1)
    max_scale_mat = torch.where(ps1 > ps2, ps1, ps2)
    pj2d_dist_mat_normalized = pj2d_dist_mat / max_scale_mat
    triu_mask = torch.triu(torch.ones_like(pj2d_dist_mat), diagonal=1) < 0.5
    pj2d_dist_mat_normalized[triu_mask] = 10000.0
    max_length = max(img_shape)
    thresh = thresh * max_length / 640
    repeat_subj_inds = torch.where(pj2d_dist_mat_normalized < thresh)
    if len(repeat_subj_inds) > 0:
        if conf_based:
            center_confs = outputs['center_confs']
            removed_subj_inds = torch.where(center_confs[repeat_subj_inds[0]] < center_confs[repeat_subj_inds[1]], repeat_subj_inds[0], repeat_subj_inds[1])
        else:
            removed_subj_inds = torch.where(person_scales[repeat_subj_inds[0]] < person_scales[repeat_subj_inds[1]], repeat_subj_inds[0], repeat_subj_inds[1])
        outputs = remove_subjects(outputs, removed_subj_inds)
    return outputs
//  [python]
def _calc_param_loss(self, outputs, meta_data, matched_mask):
        params_loss_dict = {'Pose': 0, 'Shape': 0, 'Cam': 0, 'Prior': 0}
        if args().learn_relative:
            params_loss_dict.update({'R_Age': 0, 'R_Depth': 0})
        if 'params' in outputs:
            _check_params_(meta_data['params'])
            device = outputs['params']['body_pose'].device
            grot_masks, smpl_pose_masks, smpl_shape_masks = meta_data['valid_masks'][:, 3], meta_data['valid_masks'][:, 4], meta_data['valid_masks'][:, 5]
            if grot_masks.sum() > 0:
                params_loss_dict['Pose'] += batch_smpl_pose_l2_error(meta_data['params'][grot_masks, :3].contiguous(), outputs['params']['global_orient'][grot_masks].contiguous()).mean()
            if smpl_pose_masks.sum() > 0:
                params_loss_dict['Pose'] += batch_smpl_pose_l2_error(meta_data['params'][smpl_pose_masks, 3:22 * 3].contiguous(), outputs['params']['body_pose'][smpl_pose_masks, :21 * 3].contiguous()).mean()
            if smpl_shape_masks.sum() > 0:
                smpl_shape_diff = meta_data['params'][smpl_shape_masks, -10:].contiguous() - outputs['params']['betas'][smpl_shape_masks, :10].contiguous()
                params_loss_dict['Shape'] += torch.norm(smpl_shape_diff * self.shape_pca_weight, p=2, dim=-1).mean() / 20.0
            if (~smpl_shape_masks).sum() > 0:
                params_loss_dict['Shape'] += (outputs['params']['betas'][~smpl_shape_masks, :10] ** 2).mean() / 20.0
            if args().supervise_cam_params:
                cam_mask, pred_cam_params = meta_data['cam_mask'], outputs['params']['cam']
                if cam_mask.sum() > 0:
                    params_loss_dict['Cam'] += batch_l2_loss(meta_data['cams'][cam_mask], pred_cam_params[cam_mask])
            if args().learn_relative:
                if args().learn_relative_age:
                    params_loss_dict['R_Age'] = relative_age_loss(outputs['kid_offsets_pred'], meta_data['depth_info'][:, 0], matched_mask=matched_mask) + kid_offset_loss(outputs['kid_offsets_pred'], meta_data['kid_shape_offsets'], matched_mask=matched_mask) * 2
                if args().learn_relative_depth:
                    params_loss_dict['R_Depth'] = relative_depth_loss(outputs['cam_trans'][:, 2], meta_data['depth_info'][:, 3], outputs['reorganize_idx'], matched_mask=matched_mask)
            gmm_prior_loss = self.gmm_prior(outputs['params']['body_pose']).mean() / 100.0
            valuable_prior_loss_thresh = 5.0
            gmm_prior_loss[gmm_prior_loss < valuable_prior_loss_thresh] = 0
            params_loss_dict['Prior'] = gmm_prior_loss
        return params_loss_dict
//  [python]
def _calc_param_loss(self, outputs, meta_data, matched_mask):
        params_loss_dict = {'Pose': 0, 'Shape': 0, 'Cam': 0, 'Prior': 0}
        if args().learn_relative:
            params_loss_dict.update({'R_Age': 0, 'R_Depth': 0})
        if 'params' in outputs:
            _check_params_(meta_data['params'])
            device = outputs['params']['body_pose'].device
            grot_masks, smpl_pose_masks, smpl_shape_masks = meta_data['valid_masks'][:, 3], meta_data['valid_masks'][:, 4], meta_data['valid_masks'][:, 5]
            if grot_masks.sum() > 0:
                params_loss_dict['Pose'] += batch_smpl_pose_l2_error(meta_data['params'][grot_masks, :3].contiguous(), outputs['params']['global_orient'][grot_masks].contiguous()).mean()
            if smpl_pose_masks.sum() > 0:
                params_loss_dict['Pose'] += batch_smpl_pose_l2_error(meta_data['params'][smpl_pose_masks, 3:22 * 3].contiguous(), outputs['params']['body_pose'][smpl_pose_masks, :21 * 3].contiguous()).mean()
            if smpl_shape_masks.sum() > 0:
                smpl_shape_diff = meta_data['params'][smpl_shape_masks, -10:].contiguous() - outputs['params']['betas'][smpl_shape_masks, :10].contiguous()
                params_loss_dict['Shape'] += torch.linalg.norm(smpl_shape_diff * self.shape_pca_weight, ord=2, dim=-1).mean() / 20.0
            if (~smpl_shape_masks).sum() > 0:
                params_loss_dict['Shape'] += (outputs['params']['betas'][~smpl_shape_masks, :10] ** 2).mean() / 20.0
            if args().supervise_cam_params:
                cam_mask, pred_cam_params = meta_data['cam_mask'], outputs['params']['cam']
                if cam_mask.sum() > 0:
                    params_loss_dict['Cam'] += batch_l2_loss(meta_data['cams'][cam_mask], pred_cam_params[cam_mask])
            if args().learn_relative:
                if args().learn_relative_age:
                    params_loss_dict['R_Age'] = relative_age_loss(outputs['kid_offsets_pred'], meta_data['depth_info'][:, 0], matched_mask=matched_mask) + kid_offset_loss(outputs['kid_offsets_pred'], meta_data['kid_shape_offsets'], matched_mask=matched_mask) * 2
                if args().learn_relative_depth:
                    params_loss_dict['R_Depth'] = relative_depth_loss(outputs['cam_trans'][:, 2], meta_data['depth_info'][:, 3], outputs['reorganize_idx'], matched_mask=matched_mask)
            gmm_prior_loss = self.gmm_prior(outputs['params']['body_pose']).mean() / 100.0
            valuable_prior_loss_thresh = 5.0
            gmm_prior_loss[gmm_prior_loss < valuable_prior_loss_thresh] = 0
            params_loss_dict['Prior'] = gmm_prior_loss
        return params_loss_dict
//  [python]
def calc_rot_tsl(
    x1: torch.Tensor,
    x2: torch.Tensor,
    x3: torch.Tensor
) -> (torch.Tensor, torch.Tensor):
    """
    """

    eps = 1e-4

    v1 = x3 - x2
    v2 = x1 - x2
    e1 = v1 / (torch.norm(v1) + eps)
    u2 = v2 - torch.inner(e1, v2) * e1
    e2 = u2 / (torch.norm(u2) + eps)
    e3 = torch.cross(e1, e2)
    rot_mat = torch.stack([e1, e2, e3], dim=0).permute(1, 0)
    tsl_vec = x2

    return rot_mat, tsl_vec
//  [python]
def calc_rot_tsl(
    x1: torch.Tensor,
    x2: torch.Tensor,
    x3: torch.Tensor
) -> (torch.Tensor, torch.Tensor):
    """
    """

    eps = 1e-4

    v1 = x3 - x2
    v2 = x1 - x2
    e1 = v1 / (torch.linalg.norm(v1) + eps)
    u2 = v2 - torch.inner(e1, v2) * e1
    e2 = u2 / (torch.linalg.norm(u2) + eps)
    e3 = torch.cross(e1, e2)
    rot_mat = torch.stack([e1, e2, e3], dim=0).permute(1, 0)
    tsl_vec = x2

    return rot_mat, tsl_vec
//  [python]
def match_gt_pred_3d_2d(self, center_gts_info_2d, center_gts_info_3d, pred_batch_ids, pred_czyxs, top_score, cam_czyx, device, is_training, batch_size=1, with_2d_matching=True):
        vgt_batch_ids, vgt_person_ids, vgt_centers = center_gts_info_2d
        vgt_batch_ids_3d, vgt_person_ids_3d, vgt_czyxs = center_gts_info_3d
        mc = {key: [] for key in ['batch_ids', 'matched_ids', 'person_ids', 'conf']}
        for match_ind in torch.arange(len(vgt_batch_ids_3d)):
            batch_id, person_id, center_gt = vgt_batch_ids_3d[match_ind], vgt_person_ids_3d[match_ind], vgt_czyxs[match_ind]
            pids = torch.where(pred_batch_ids == batch_id)[0]
            if len(pids) == 0:
                continue
            center_dist_3d = torch.norm(pred_czyxs[pids].float() - center_gt[None].float(), dim=-1)
            matched_pred_id = pids[torch.argmin(center_dist_3d)]
            mc['batch_ids'].append(batch_id)
            mc['matched_ids'].append(matched_pred_id)
            mc['person_ids'].append(person_id)
            mc['conf'].append(top_score[matched_pred_id])
        for match_ind in torch.arange(len(vgt_batch_ids)):
            batch_id, person_id, center_gt = vgt_batch_ids[match_ind], vgt_person_ids[match_ind], vgt_centers[match_ind]
            pids = torch.where(pred_batch_ids == batch_id)[0]
            if len(pids) == 0:
                continue
            matched_pred_id = pids[torch.argmin(torch.norm(pred_czyxs[pids, 1:].float() - center_gt[None].float(), dim=-1))]
            center_matched = pred_czyxs[matched_pred_id].long()
            mc['batch_ids'].append(batch_id)
            mc['matched_ids'].append(matched_pred_id)
            mc['person_ids'].append(person_id)
            mc['conf'].append(top_score[matched_pred_id])
        if args().eval_2dpose:
            for inds, (batch_id, person_id, center_gt) in enumerate(zip(vgt_batch_ids, vgt_person_ids, vgt_centers)):
                if batch_id in pred_batch_ids:
                    center_pred = pred_czyxs[pred_batch_ids == batch_id]
                    matched_id = torch.argmin(torch.norm(center_pred[:, 1:].float() - center_gt[None].float(), dim=-1))
                    matched_pred_id = np.where((pred_batch_ids == batch_id).cpu())[0][matched_id]
                    mc['matched_ids'].append(matched_pred_id)
                    mc['batch_ids'].append(batch_id)
                    mc['person_ids'].append(person_id)
        if len(mc['matched_ids']) == 0:
            mc.update({'batch_ids': [0], 'matched_ids': [0], 'person_ids': [0], 'conf': [0]})
        keys_list = list(mc.keys())
        for key in keys_list:
            if key == 'conf':
                mc[key] = torch.Tensor(mc[key])
            else:
                mc[key] = torch.Tensor(mc[key]).long()
            if args().max_supervise_num != -1 and is_training:
                mc[key] = mc[key][:args().max_supervise_num]
        return mc
//  [python]
def match_gt_pred_3d_2d(self, center_gts_info_2d, center_gts_info_3d, pred_batch_ids, pred_czyxs, top_score, cam_czyx, device, is_training, batch_size=1, with_2d_matching=True):
        vgt_batch_ids, vgt_person_ids, vgt_centers = center_gts_info_2d
        vgt_batch_ids_3d, vgt_person_ids_3d, vgt_czyxs = center_gts_info_3d
        mc = {key: [] for key in ['batch_ids', 'matched_ids', 'person_ids', 'conf']}
        for match_ind in torch.arange(len(vgt_batch_ids_3d)):
            batch_id, person_id, center_gt = vgt_batch_ids_3d[match_ind], vgt_person_ids_3d[match_ind], vgt_czyxs[match_ind]
            pids = torch.where(pred_batch_ids == batch_id)[0]
            if len(pids) == 0:
                continue
            center_dist_3d = torch.linalg.norm(pred_czyxs[pids].float() - center_gt[None].float(), dim=-1)
            matched_pred_id = pids[torch.argmin(center_dist_3d)]
            mc['batch_ids'].append(batch_id)
            mc['matched_ids'].append(matched_pred_id)
            mc['person_ids'].append(person_id)
            mc['conf'].append(top_score[matched_pred_id])
        for match_ind in torch.arange(len(vgt_batch_ids)):
            batch_id, person_id, center_gt = vgt_batch_ids[match_ind], vgt_person_ids[match_ind], vgt_centers[match_ind]
            pids = torch.where(pred_batch_ids == batch_id)[0]
            if len(pids) == 0:
                continue
            matched_pred_id = pids[torch.argmin(torch.linalg.norm(pred_czyxs[pids, 1:].float() - center_gt[None].float(), dim=-1))]
            center_matched = pred_czyxs[matched_pred_id].long()
            mc['batch_ids'].append(batch_id)
            mc['matched_ids'].append(matched_pred_id)
            mc['person_ids'].append(person_id)
            mc['conf'].append(top_score[matched_pred_id])
        if args().eval_2dpose:
            for inds, (batch_id, person_id, center_gt) in enumerate(zip(vgt_batch_ids, vgt_person_ids, vgt_centers)):
                if batch_id in pred_batch_ids:
                    center_pred = pred_czyxs[pred_batch_ids == batch_id]
                    matched_id = torch.argmin(torch.linalg.norm(center_pred[:, 1:].float() - center_gt[None].float(), dim=-1))
                    matched_pred_id = np.where((pred_batch_ids == batch_id).cpu())[0][matched_id]
                    mc['matched_ids'].append(matched_pred_id)
                    mc['batch_ids'].append(batch_id)
                    mc['person_ids'].append(person_id)
        if len(mc['matched_ids']) == 0:
            mc.update({'batch_ids': [0], 'matched_ids': [0], 'person_ids': [0], 'conf': [0]})
        keys_list = list(mc.keys())
        for key in keys_list:
            if key == 'conf':
                mc[key] = torch.Tensor(mc[key])
            else:
                mc[key] = torch.Tensor(mc[key]).long()
            if args().max_supervise_num != -1 and is_training:
                mc[key] = mc[key][:args().max_supervise_num]
        return mc
//  [python]
def forward(self, x):
        if self.kernel == 'linear':
            x = x.view(-1, 1, self.n_features)
            k = self.sv * x
            k = k.sum(2)
        elif self.kernel == 'rbf':
            x = x.view(-1, 1, self.n_features)
            k = torch.pow(self.sv - x, 2)
            k = k.sum(2)
            k = torch.exp(-self.gamma * k)
        elif self.kernel == 'poly':
            k = torch.pow(self.gamma * torch.mm(x, self.sv.t()) + self.coef0, self.degree)
        elif self.kernel == 'sigmoid':
            k = torch.tanh(self.gamma * torch.mm(x, self.sv.t()) + self.coef0)
        elif self.kernel == 'cosine':
            norm_x = torch.norm(x, keepdim=True, dim=1)
            norm_sv = torch.norm(self.sv, keepdim=True, dim=1)
            norm = torch.mm(norm_x, norm_sv.t())
            k = torch.mm(x, self.sv.t())
            k = torch.div(k, norm)
        elif self.kernel == 'precomputed':
            k = x
        else:
            raise NotImplementedError('Hummingbird does not currently support {} kernel for KernelPCA. The supported kernels are linear, poly, rbf, sigmoid, cosine, and precomputed.'.format(self.kernel))
        k_pred_cols = (torch.sum(k, 1) / self.n_samples).view(-1, 1)
        k -= self.k_fit_rows
        k -= k_pred_cols
        k += self.k_fit_all
        return torch.mm(k, self.scaled_alphas)
//  [python]
def forward(self, x):
        if self.kernel == 'linear':
            x = x.view(-1, 1, self.n_features)
            k = self.sv * x
            k = k.sum(2)
        elif self.kernel == 'rbf':
            x = x.view(-1, 1, self.n_features)
            k = torch.pow(self.sv - x, 2)
            k = k.sum(2)
            k = torch.exp(-self.gamma * k)
        elif self.kernel == 'poly':
            k = torch.pow(self.gamma * torch.mm(x, self.sv.t()) + self.coef0, self.degree)
        elif self.kernel == 'sigmoid':
            k = torch.tanh(self.gamma * torch.mm(x, self.sv.t()) + self.coef0)
        elif self.kernel == 'cosine':
            norm_x = torch.linalg.norm(x, keepdim=True, dim=1)
            norm_sv = torch.linalg.norm(self.sv, keepdim=True, dim=1)
            norm = torch.mm(norm_x, norm_sv.t())
            k = torch.mm(x, self.sv.t())
            k = torch.div(k, norm)
        elif self.kernel == 'precomputed':
            k = x
        else:
            raise NotImplementedError('Hummingbird does not currently support {} kernel for KernelPCA. The supported kernels are linear, poly, rbf, sigmoid, cosine, and precomputed.'.format(self.kernel))
        k_pred_cols = (torch.sum(k, 1) / self.n_samples).view(-1, 1)
        k -= self.k_fit_rows
        k -= k_pred_cols
        k += self.k_fit_all
        return torch.mm(k, self.scaled_alphas)
//  [python]
def _project_batch(self, oft_r, eps=1e-5):
        # scaling factor for each of the smaller block matrix
        eps = eps * 1 / torch.sqrt(torch.tensor(oft_r.shape[0]))
        I = (  # noqa: E741
            torch.zeros((oft_r.size(1), oft_r.size(1)), device=oft_r.device, dtype=oft_r.dtype)
            .unsqueeze(0)
            .expand_as(oft_r)
        )
        diff = oft_r - I
        norm_diff = torch.norm(oft_r - I, dim=(1, 2), keepdim=True)
        mask = (norm_diff <= eps).bool()
        out = torch.where(mask, oft_r, I + eps * (diff / norm_diff))
        return out
//  [python]
def _project_batch(self, oft_r, eps=1e-5):
        # scaling factor for each of the smaller block matrix
        eps = eps * 1 / torch.sqrt(torch.tensor(oft_r.shape[0]))
        I = (  # noqa: E741
            torch.zeros((oft_r.size(1), oft_r.size(1)), device=oft_r.device, dtype=oft_r.dtype)
            .unsqueeze(0)
            .expand_as(oft_r)
        )
        diff = oft_r - I
        norm_diff = torch.linalg.norm(oft_r - I, dim=(1, 2), keepdim=True)
        mask = (norm_diff <= eps).bool()
        out = torch.where(mask, oft_r, I + eps * (diff / norm_diff))
        return out
//  [python]
def set_template(self):
        
        fn_template_data = join(self.input_folder, f'{self.template_scan_name}.npz')
        template_data = np.load(fn_template_data)
        
        verts, faces = igl.read_triangle_mesh(join(self.data_folder, self.template_mesh_fn))
        v_boundary_mask = template_data['v_boundary_mask']
        vert_colors = template_data['v_colors']


        self.template_verts = torch.from_numpy(verts).float().cuda()
        self.template_faces = torch.from_numpy(faces).long().cuda()
        self.template_edges = compute_edges(self.template_verts, self.template_faces)
        self.template_v_boundary_mask = torch.from_numpy(v_boundary_mask).cuda()
        self.template_v_colors = torch.from_numpy(vert_colors[:, :3].astype(np.float32)).float().cuda()
        verts_ori = torch.from_numpy(template_data['scan_v']).float().cuda()
        self.template_e_lengths = torch.norm(verts_ori[self.template_edges[:, 0]] - \
                                             verts_ori[self.template_edges[:, 1]], dim=-1).cuda()


        fn_template_eigenfuncs = join(self.eigenfuncs_folder, f'{self.template_scan_name}_eigenfuncs.npz')
        template_eigenfuncs_data = np.load(fn_template_eigenfuncs)

        self.template_lbsw = torch.from_numpy(np.load(join(self.data_folder, self.template_lbsw_fn))).float().cuda()
        self.template_eigenfuncs = torch.from_numpy(template_eigenfuncs_data['eigenfuncs']).float().cuda()
        self.template_eigenvals = torch.from_numpy(template_eigenfuncs_data['eigenvals']).float().cuda()
//  [python]
def set_template(self):
        
        fn_template_data = join(self.input_folder, f'{self.template_scan_name}.npz')
        template_data = np.load(fn_template_data)
        
        verts, faces = igl.read_triangle_mesh(join(self.data_folder, self.template_mesh_fn))
        v_boundary_mask = template_data['v_boundary_mask']
        vert_colors = template_data['v_colors']


        self.template_verts = torch.from_numpy(verts).float().cuda()
        self.template_faces = torch.from_numpy(faces).long().cuda()
        self.template_edges = compute_edges(self.template_verts, self.template_faces)
        self.template_v_boundary_mask = torch.from_numpy(v_boundary_mask).cuda()
        self.template_v_colors = torch.from_numpy(vert_colors[:, :3].astype(np.float32)).float().cuda()
        verts_ori = torch.from_numpy(template_data['scan_v']).float().cuda()
        self.template_e_lengths = torch.linalg.norm(verts_ori[self.template_edges[:, 0]] - \
                                             verts_ori[self.template_edges[:, 1]], dim=-1).cuda()


        fn_template_eigenfuncs = join(self.eigenfuncs_folder, f'{self.template_scan_name}_eigenfuncs.npz')
        template_eigenfuncs_data = np.load(fn_template_eigenfuncs)

        self.template_lbsw = torch.from_numpy(np.load(join(self.data_folder, self.template_lbsw_fn))).float().cuda()
        self.template_eigenfuncs = torch.from_numpy(template_eigenfuncs_data['eigenfuncs']).float().cuda()
        self.template_eigenvals = torch.from_numpy(template_eigenfuncs_data['eigenvals']).float().cuda()
//  [python]
def OracleCPD(T, R, iteration, device, A1=None, A2=None, A3=None):
    """
    OracleCPD model:
        we call the <optimize> for each least square sub-problem
    """

    # preparation
    loss_list, rec_list = [], []
    d1, d2, d3 = T.shape

    if A1 is None:
        A1 = torch.randn(d1, R).to(device)
        A2 = torch.randn(d2, R).to(device)
        A3 = torch.randn(d3, R).to(device)
    
    # ALS 
    tic = time.time()
    for i in range(iteration):

        # sub-iteration
        A1 = optimize((A3.T@A3)*(A2.T@A2), torch.einsum('ijk,jr,kr->ri',T,A2,A3)).T
        A2 = optimize((A1.T@A1)*(A3.T@A3), torch.einsum('ijk,ir,kr->rj',T,A1,A3)).T
        A3 = optimize((A1.T@A1)*(A2.T@A2), torch.einsum('ijk,ir,jr->rk',T,A1,A2)).T

        # loss 
        rec = torch.einsum('ir,jr,kr->ijk',A1,A2,A3)
        LOSS = torch.norm(rec - T)/ torch.norm(T)
        recLOSS = 1 - torch.norm(rec - T) / torch.norm(T)
        if i % 20 == 0:
            print ('{}/{}'.format(i, iteration), 'loss:', LOSS.item(), 'rec:', recLOSS.item(), 'time span:', time.time() - tic)

        # collect loss
        tic = time.time() 
        loss_list.append(LOSS)
        rec_list.append(recLOSS)

    return A1, A2, A3, loss_list, rec_list
//  [python]
def OracleCPD(T, R, iteration, device, A1=None, A2=None, A3=None):
    """
    OracleCPD model:
        we call the <optimize> for each least square sub-problem
    """

    # preparation
    loss_list, rec_list = [], []
    d1, d2, d3 = T.shape

    if A1 is None:
        A1 = torch.randn(d1, R).to(device)
        A2 = torch.randn(d2, R).to(device)
        A3 = torch.randn(d3, R).to(device)
    
    # ALS 
    tic = time.time()
    for i in range(iteration):

        # sub-iteration
        A1 = optimize((A3.T@A3)*(A2.T@A2), torch.einsum('ijk,jr,kr->ri',T,A2,A3)).T
        A2 = optimize((A1.T@A1)*(A3.T@A3), torch.einsum('ijk,ir,kr->rj',T,A1,A3)).T
        A3 = optimize((A1.T@A1)*(A2.T@A2), torch.einsum('ijk,ir,jr->rk',T,A1,A2)).T

        # loss 
        rec = torch.einsum('ir,jr,kr->ijk',A1,A2,A3)
        LOSS = torch.linalg.norm(rec - T)/ torch.linalg.norm(T)
        recLOSS = 1 - torch.linalg.norm(rec - T) / torch.linalg.norm(T)
        if i % 20 == 0:
            print ('{}/{}'.format(i, iteration), 'loss:', LOSS.item(), 'rec:', recLOSS.item(), 'time span:', time.time() - tic)

        # collect loss
        tic = time.time() 
        loss_list.append(LOSS)
        rec_list.append(recLOSS)

    return A1, A2, A3, loss_list, rec_list
//  [python]
def forward(self, x):
        if self.norm:
            weights = self.context_vector(x) / torch.norm(self.context_vector.weight, p=2)
            weights = torch.add(torch.div(weights, 2.), .5)
        else:
            x_tr = self.activation(self.transform(x))
            weights = self.context_vector(x_tr)
            weights = torch.sigmoid(weights)
        x = x * weights
        return x, weights