| 22 | |
| 23 | class SetCriterion(nn.Module): |
| 24 | def __init__(self, |
| 25 | num_classes, |
| 26 | matcher, |
| 27 | weight_dict, |
| 28 | focal_alpha, |
| 29 | losses, |
| 30 | num_box_decoder_layers=2, |
| 31 | num_hand_face_decoder_layers=4, |
| 32 | num_body_points=17, |
| 33 | num_hand_points=6, |
| 34 | num_face_points=6, |
| 35 | smpl_loss_config=None, |
| 36 | convention='smplx_137'): |
| 37 | super().__init__() |
| 38 | self.num_classes = num_classes |
| 39 | self.matcher = matcher |
| 40 | self.weight_dict = weight_dict |
| 41 | self.losses = losses |
| 42 | self.focal_alpha = focal_alpha |
| 43 | self.vis = 0.1 |
| 44 | self.abs = 1 |
| 45 | self.num_body_points = num_body_points |
| 46 | self.num_hand_points = num_hand_points |
| 47 | self.num_face_points = num_face_points |
| 48 | self.num_box_decoder_layers = num_box_decoder_layers |
| 49 | self.num_hand_face_decoder_layers = num_hand_face_decoder_layers |
| 50 | self.convention = convention |
| 51 | self.body_oks = OKSLoss(linear=True, |
| 52 | num_keypoints=num_body_points, |
| 53 | eps=1e-6, |
| 54 | reduction='mean', |
| 55 | loss_weight=1.0) |
| 56 | self.hand_oks = OKSLoss(linear=True, |
| 57 | num_keypoints=num_hand_points, |
| 58 | eps=1e-6, |
| 59 | reduction='mean', |
| 60 | loss_weight=1.0) |
| 61 | self.face_oks = OKSLoss(linear=True, |
| 62 | num_keypoints=num_face_points, |
| 63 | eps=1e-6, |
| 64 | reduction='mean', |
| 65 | loss_weight=1.0) |
| 66 | |
| 67 | def loss_labels(self, |
| 68 | outputs, |