Update difpoint/inference.py

difpoint/inference.py

@@ -55,7 +55,16 @@ import os
 import datetime
 import platform
 from omegaconf import OmegaConf
-from difpoint.src.pipelines.faster_live_portrait_pipeline import FasterLivePortraitPipeline
+#from difpoint.src.pipelines.faster_live_portrait_pipeline import FasterLivePortraitPipeline
+from difpoint.src.live_portrait_pipeline import LivePortraitPipeline
+from difpointsrc.config.argument_config import ArgumentConfig
+from difpoint.src.config.inference_config import InferenceConfig
+from difpoint.src.config.crop_config import CropConfig
+from difpoint.src.live_portrait_pipeline import LivePortraitPipeline
+from difpoint.src.utils.retargeting_utils import calc_eye_close_ratio, calc_lip_close_ratio
+from difpoint.src.utils.camera import get_rotation_matrix
+from difpoint.src.utils.video import images2video, co
+
 
 FFMPEG = "ffmpeg"
 
@@ -178,13 +187,12 @@ class Inferencer(object):
         self.wav2lip_model.cuda()
         self.wav2lip_model.eval()
 
-        # specify configs for inference
-        self.inf_cfg = OmegaConf.load("difpoint/configs/trt_mp_infer.yaml")
-        self.inf_cfg.infer_params.flag_pasteback = False
+        args = tyro.cli(ArgumentConfig)
 
-        self.live_portrait_pipeline = FasterLivePortraitPipeline(cfg=self.inf_cfg, is_animal=False)
-        #ret = self.live_portrait_pipeline.prepare_source(source_image)
+        self.inf_cfg = partial_fields(InferenceConfig, args.__dict__)  # use attribute of args to initial InferenceConfig
+        self.crop_cfg = partial_fields(CropConfig, args.__dict__)  # use attribute of args to initial CropConfig
 
+        self.live_portrait_pipeline = LivePortraitPipeline(inference_cfg=self.inf_cfg, crop_cfg=self.crop_cfg)
         print('#'*25+f'End initialization, cost time {time.time()-st}'+'#'*25)
 
     def _norm(self, data_dict):
@@ -286,24 +294,15 @@ class Inferencer(object):
         else:
             input_image = image[0]
 
-        I_s = (torch.FloatTensor(input_image.transpose((2, 0, 1))).unsqueeze(0).cuda() / 255).cpu().numpy()
-        pitch, yaw, roll, t, exp, scale, kp = self.live_portrait_pipeline.model_dict["motion_extractor"].predict(
-                        I_s)
-        x_s_info = {
-                        "pitch": pitch,
-                        "yaw": yaw,
-                        "roll": roll,
-                        "t": t,
-                        "exp": exp,
-                        "scale": scale,
-                        "kp": kp
-                    }
-        x_c_s = kp.reshape(1, 21, -1)
+        I_s = torch.FloatTensor(input_image.transpose((2, 0, 1))).unsqueeze(0).cuda() / 255
+
+        x_s_info = self.live_portrait_pipeline.live_portrait_wrapper.get_kp_info(I_s)
+        x_c_s = x_s_info['kp'].reshape(1, 21, -1)
         R_s = get_rotation_matrix(x_s_info['pitch'], x_s_info['yaw'], x_s_info['roll'])
-        f_s = self.live_portrait_pipeline.model_dict["app_feat_extractor"].predict(I_s)
-        x_s = transform_keypoint(pitch, yaw, roll, t, exp, scale, kp)
+        f_s = self.live_portrait_pipeline.live_portrait_wrapper.extract_feature_3d(I_s)
+        x_s = self.live_portrait_pipeline.live_portrait_wrapper.transform_keypoint(x_s_info)
 
-        flag_lip_zero = self.inf_cfg.infer_params.flag_normalize_lip 
+        flag_lip_zero = self.inf_cfg.flag_lip_zero  # not overwrite 
     
        
 
@@ -440,7 +439,7 @@ class Inferencer(object):
                 pass
             elif self.inf_cfg.infer_params.flag_stitching and not self.inf_cfg.infer_params.flag_eye_retargeting and not self.inf_cfg.infer_params.flag_lip_retargeting:
                 # with stitching and without retargeting
-                x_d_i_new = self.live_portrait_pipeline.stitching(x_s, x_d_i_new)
+                x_d_i_new = self.live_portrait_pipeline.live_portrait_wrapper.stitching(x_s, x_d_i_new)
             else:
                 eyes_delta, lip_delta = None, None
                 if self.inf_cfg.infer_params.flag_eye_retargeting:
@@ -458,10 +457,11 @@ class Inferencer(object):
                                 (lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)
 
                 if self.inf_cfg.infer_params.flag_stitching:
-                    x_d_i_new = self.live_portrait_pipeline.stitching(x_s, x_d_i_new)
+                    x_d_i_new = self.live_portrait_pipeline.live_portrait_wrapper.stitching(x_s, x_d_i_new)
 
-            out = self.live_portrait_pipeline.model_dict["warping_spade"].predict(f_s, x_s, x_d_i_new).cpu().numpy().astype(np.uint32)
-            I_p_lst.append(out)
+            out = self.live_portrait_pipeline.live_portrait_wrapper.warp_decode(f_s, x_s, x_d_i_new)
+            I_p_i = self.live_portrait_pipeline.live_portrait_wrapper.parse_output(out['out'])[0]
+            I_p_lst.append(I_p_i)
 
         video_name = os.path.basename(save_path)
         video_save_dir = os.path.dirname(save_path)