Spaces:
Running
on
Zero
Running
on
Zero
| #!/usr/bin/env python | |
| # Copyright (c) Xuangeng Chu ([email protected]) | |
| # Modified based on code from Orest Kupyn (University of Oxford). | |
| import torch | |
| import torchvision | |
| def reproject_vertices(flame_model, vgg_results): | |
| # flame_model = FLAMEModel(n_shape=300, n_exp=100, scale=1.0) | |
| vertices, _ = flame_model( | |
| shape_params=vgg_results['shapecode'], | |
| expression_params=vgg_results['expcode'], | |
| pose_params=vgg_results['posecode'], | |
| verts_sclae=1.0 | |
| ) | |
| vertices[:, :, 2] += 0.05 # MESH_OFFSET_Z | |
| vgg_landmarks3d = flame_model._vertices2landmarks(vertices) | |
| vgg_transform_results = vgg_results['transform'] | |
| rotation_mat = rot_mat_from_6dof(vgg_transform_results['rotation_6d']).type(vertices.dtype) | |
| translation = vgg_transform_results['translation'][:, None, :] | |
| scale = torch.clamp(vgg_transform_results['scale'][:, None], 1e-8) | |
| rot_vertices = vertices.clone() | |
| rot_vertices = torch.matmul(rotation_mat.unsqueeze(1), rot_vertices.unsqueeze(-1))[..., 0] | |
| vgg_landmarks3d = torch.matmul(rotation_mat.unsqueeze(1), vgg_landmarks3d.unsqueeze(-1))[..., 0] | |
| proj_vertices = (rot_vertices * scale) + translation | |
| vgg_landmarks3d = (vgg_landmarks3d * scale) + translation | |
| trans_padding, trans_scale = vgg_results['normalize']['padding'], vgg_results['normalize']['scale'] | |
| proj_vertices[:, :, 0] -= trans_padding[:, 0, None] | |
| proj_vertices[:, :, 1] -= trans_padding[:, 1, None] | |
| proj_vertices = proj_vertices / trans_scale[:, None, None] | |
| vgg_landmarks3d[:, :, 0] -= trans_padding[:, 0, None] | |
| vgg_landmarks3d[:, :, 1] -= trans_padding[:, 1, None] | |
| vgg_landmarks3d = vgg_landmarks3d / trans_scale[:, None, None] | |
| return proj_vertices.float()[..., :2], vgg_landmarks3d.float()[..., :2] | |
| def rot_mat_from_6dof(v: torch.Tensor) -> torch.Tensor: | |
| assert v.shape[-1] == 6 | |
| v = v.view(-1, 6) | |
| vx, vy = v[..., :3].clone(), v[..., 3:].clone() | |
| b1 = torch.nn.functional.normalize(vx, dim=-1) | |
| b3 = torch.nn.functional.normalize(torch.cross(b1, vy, dim=-1), dim=-1) | |
| b2 = -torch.cross(b1, b3, dim=1) | |
| return torch.stack((b1, b2, b3), dim=-1) | |
| def nms(boxes_xyxy, scores, flame_params, | |
| confidence_threshold: float = 0.5, iou_threshold: float = 0.5, | |
| top_k: int = 1000, keep_top_k: int = 100 | |
| ): | |
| for pred_bboxes_xyxy, pred_bboxes_conf, pred_flame_params in zip( | |
| boxes_xyxy.detach().float(), | |
| scores.detach().float(), | |
| flame_params.detach().float(), | |
| ): | |
| pred_bboxes_conf = pred_bboxes_conf.squeeze(-1) # [Anchors] | |
| conf_mask = pred_bboxes_conf >= confidence_threshold | |
| pred_bboxes_conf = pred_bboxes_conf[conf_mask] | |
| pred_bboxes_xyxy = pred_bboxes_xyxy[conf_mask] | |
| pred_flame_params = pred_flame_params[conf_mask] | |
| # Filter all predictions by self.nms_top_k | |
| if pred_bboxes_conf.size(0) > top_k: | |
| topk_candidates = torch.topk(pred_bboxes_conf, k=top_k, largest=True, sorted=True) | |
| pred_bboxes_conf = pred_bboxes_conf[topk_candidates.indices] | |
| pred_bboxes_xyxy = pred_bboxes_xyxy[topk_candidates.indices] | |
| pred_flame_params = pred_flame_params[topk_candidates.indices] | |
| # NMS | |
| idx_to_keep = torchvision.ops.boxes.nms(boxes=pred_bboxes_xyxy, scores=pred_bboxes_conf, iou_threshold=iou_threshold) | |
| final_bboxes = pred_bboxes_xyxy[idx_to_keep][: keep_top_k] # [Instances, 4] | |
| final_scores = pred_bboxes_conf[idx_to_keep][: keep_top_k] # [Instances, 1] | |
| final_params = pred_flame_params[idx_to_keep][: keep_top_k] # [Instances, Flame Params] | |
| return final_bboxes, final_scores, final_params | |