| 33 | |
| 34 | |
| 35 | class AiOSSMPLX(nn.Module): |
| 36 | def __init__( |
| 37 | self, |
| 38 | backbone, |
| 39 | transformer, |
| 40 | num_classes, |
| 41 | num_queries, |
| 42 | aux_loss=False, |
| 43 | iter_update=True, |
| 44 | query_dim=4, |
| 45 | random_refpoints_xy=False, |
| 46 | fix_refpoints_hw=-1, |
| 47 | num_feature_levels=1, |
| 48 | nheads=8, |
| 49 | two_stage_type='no', |
| 50 | dec_pred_class_embed_share=False, |
| 51 | dec_pred_bbox_embed_share=False, |
| 52 | dec_pred_pose_embed_share=False, |
| 53 | two_stage_class_embed_share=True, |
| 54 | two_stage_bbox_embed_share=True, |
| 55 | dn_number=100, |
| 56 | dn_box_noise_scale=0.4, |
| 57 | dn_label_noise_ratio=0.5, |
| 58 | dn_batch_gt_fuse=False, |
| 59 | dn_labelbook_size=100, |
| 60 | dn_attn_mask_type_list=['group2group'], |
| 61 | cls_no_bias=False, |
| 62 | num_group=100, |
| 63 | num_body_points=17, |
| 64 | num_hand_points=10, |
| 65 | num_face_points=10, |
| 66 | num_box_decoder_layers=2, |
| 67 | num_hand_face_decoder_layers=4, |
| 68 | body_model=dict( |
| 69 | type='smplx', |
| 70 | keypoint_src='smplx', |
| 71 | num_expression_coeffs=10, |
| 72 | keypoint_dst='smplx_137', |
| 73 | model_path='data/body_models/smplx', |
| 74 | use_pca=False, |
| 75 | use_face_contour=True), |
| 76 | train=True, |
| 77 | inference=False, |
| 78 | focal_length=[5000., 5000.], |
| 79 | camera_3d_size=2.5 |
| 80 | ): |
| 81 | super().__init__() |
| 82 | self.num_queries = num_queries |
| 83 | self.transformer = transformer |
| 84 | self.num_classes = num_classes |
| 85 | self.hidden_dim = hidden_dim = transformer.d_model |
| 86 | self.num_feature_levels = num_feature_levels |
| 87 | self.nheads = nheads |
| 88 | self.label_enc = nn.Embedding(dn_labelbook_size + 1, hidden_dim) |
| 89 | self.num_body_points = num_body_points |
| 90 | self.num_hand_points = num_hand_points |
| 91 | self.num_face_points = num_face_points |
| 92 | self.num_whole_body_points = num_body_points + 2*num_hand_points + num_face_points |