| 694 | |
| 695 | |
| 696 | class TransformerDecoder(nn.Module): |
| 697 | def __init__( |
| 698 | self, |
| 699 | decoder_layer, |
| 700 | num_layers, |
| 701 | norm=None, |
| 702 | return_intermediate=False, |
| 703 | d_model=256, |
| 704 | query_dim=4, |
| 705 | modulate_hw_attn=False, |
| 706 | num_feature_levels=1, |
| 707 | deformable_decoder=False, |
| 708 | dec_layer_number=None, # number of queries each layer in decoder |
| 709 | dec_layer_share=False, |
| 710 | dec_layer_dropout_prob=None, |
| 711 | num_box_decoder_layers=2, |
| 712 | num_hand_face_decoder_layers=4, |
| 713 | num_body_points=17, |
| 714 | num_hand_points=10, |
| 715 | num_face_points=10, |
| 716 | num_dn=100, |
| 717 | num_group=100): |
| 718 | super().__init__() |
| 719 | # pdb.set_trace() |
| 720 | if num_layers > 0: |
| 721 | self.layers = _get_clones(decoder_layer, |
| 722 | num_layers, |
| 723 | layer_share=dec_layer_share) |
| 724 | else: |
| 725 | self.layers = [] |
| 726 | self.num_layers = num_layers |
| 727 | self.norm = norm |
| 728 | self.return_intermediate = return_intermediate # True |
| 729 | assert return_intermediate, 'support return_intermediate only' |
| 730 | self.query_dim = query_dim # 4 |
| 731 | assert query_dim in [ |
| 732 | 2, 4 |
| 733 | ], 'query_dim should be 2/4 but {}'.format(query_dim) |
| 734 | self.num_feature_levels = num_feature_levels # 4 |
| 735 | |
| 736 | self.ref_point_head = MLP(query_dim // 2 * d_model, d_model, d_model, |
| 737 | 2) # 4//2 * 256, 256, 256, 2 |
| 738 | if not deformable_decoder: |
| 739 | self.query_pos_sine_scale = MLP(d_model, d_model, d_model, 2) |
| 740 | else: |
| 741 | self.query_pos_sine_scale = None |
| 742 | |
| 743 | self.num_body_points = num_body_points |
| 744 | self.num_hand_points = num_hand_points |
| 745 | self.num_face_points = num_face_points |
| 746 | self.query_scale = None |
| 747 | |
| 748 | # aios kp |
| 749 | self.bbox_embed = None |
| 750 | self.class_embed = None |
| 751 | self.pose_embed = None |
| 752 | self.pose_hw_embed = None |
| 753 | |