Update EdgeCape/models/detectors/EdgeCape.py

EdgeCape/models/detectors/EdgeCape.py

@@ -12,6 +12,7 @@ from mmpose.models.builder import POSENETS
 from mmpose.models.detectors.base import BasePose
 from EdgeCape.models.backbones.adapter import DPT
 from EdgeCape.models.backbones.dino import DINO
+import json
 
 
 @POSENETS.register_module()
@@ -62,14 +63,7 @@ class EdgeCape(BasePose):
         self.keypoint_head_module.init_weights()
 
     def forward(self,
-                img_s,
-                img_q,
-                target_s=None,
-                target_weight_s=None,
-                target_q=None,
-                target_weight_q=None,
-                img_metas=None,
-                return_loss=True,
+                input_str,
                 **kwargs):
         """Calls either forward_train or forward_test depending on whether
         return_loss=True. Note this setting will change the expected inputs.
@@ -78,6 +72,30 @@ class EdgeCape(BasePose):
         should be double nested (i.e.  List[Tensor], List[List[dict]]), with
         the outer list indicating test time augmentations.
         """
+
+        str_dict = json.loads(input_str)
+        str_dict["img_s"] = [torch.tensor(str_dict["img_s"], dtype=torch.float32).cuda()]
+        str_dict["img_q"] = torch.tensor(str_dict["img_q"], dtype=torch.float32).cuda()
+        str_dict["target_weight_s"] = [torch.tensor(str_dict["target_weight_s"], dtype=torch.float32).cuda()]
+        str_dict["target_s"] = [torch.tensor(str_dict["target_s"], dtype=torch.float32).cuda()]
+        
+        str_dict['img_metas'][0]['sample_joints_3d'][0] = torch.tensor(str_dict['img_metas'][0]['sample_joints_3d'][0])
+        str_dict['img_metas'][0]['query_joints_3d'] = torch.tensor(str_dict['img_metas'][0]['query_joints_3d'])
+        str_dict['img_metas'][0]['sample_center'][0] = torch.tensor(str_dict['img_metas'][0]['sample_center'][0])
+        str_dict['img_metas'][0]['query_center'] = torch.tensor(str_dict['img_metas'][0]['query_center'])
+        str_dict['img_metas'][0]['sample_scale'][0] = torch.tensor(str_dict['img_metas'][0]['sample_scale'][0])
+        str_dict['img_metas'][0]['query_scale'] = torch.tensor(str_dict['img_metas'][0]['query_scale'])
+        
+        img_s = str_dict["img_s"]
+        img_q = str_dict["img_q"]
+        target_s = str_dict["target_s"]
+        target_weight_s = str_dict["target_weight_s"]
+        target_q = str_dict["target_q"]
+        target_weight_q = str_dict["target_weight_q"]
+        return_loss = str_dict["return_loss"]
+        img_metas = str_dict["img_metas"]
+        kwargs = {}
+        
         if return_loss:
             return self.forward_train(img_s, target_s, target_weight_s, img_q,
                                       target_q, target_weight_q, img_metas,