added expression regularization

Author: ankarako

preproc_video.py

@@ -5,7 +5,7 @@
 from glob import glob
 from pathlib import Path
 
-import cv2
+import sys
 import numpy as np
 import torch
 import torch.backends.cudnn as cudnn
@@ -83,6 +83,8 @@ def render(self, verts, faces, cameras, flame_mask_tex=None):
 
         return debug_view, None
 
+def l2_loss(inputs: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    return torch.sqrt(((inputs - target) ** 2).sum(dim=-1)).mean(dim=1).mean()
 
 class OptimizationLoss(torch.nn.Module):
     def __init__(
@@ -124,17 +126,20 @@ def forward(
         fan_lmks_tgt: torch.Tensor,
         seg_mask: torch.Tensor,
         seg_mask_tgt: torch.Tensor,
-        expresion_vector: torch.Tensor,
+        expression_vector: torch.Tensor,
         iris_lmks: torch.Tensor=None,
         iris_lmks_tgt: torch.Tensor=None,
     ) -> torch.Tensor:
         """"""
-        mp_loss = self.wing_loss(mp_lmks, mp_lmks_tgt)
-        fan_loss = self.wing_loss(fan_lmks, fan_lmks_tgt)
-        iris_loss = self.wing_loss(iris_lmks, iris_lmks_tgt) if iris_lmks is not None else torch.zeros([1] ,device=mp_lmks.device, dtype=torch.float32)
+        mp_loss = l2_loss(mp_lmks, mp_lmks_tgt)
+        fan_loss = l2_loss(fan_lmks, fan_lmks_tgt)
+        iris_loss = l2_loss(iris_lmks, iris_lmks_tgt) if iris_lmks is not None else torch.zeros([1] ,device=mp_lmks.device, dtype=torch.float32)
 
         seg_mask_loss = torch.abs(seg_mask - seg_mask_tgt).mean()
-        output = mp_loss * self.w_mp + fan_loss + iris_loss + seg_mask_loss * self.w_seg + expresion_vector.abs().mean() * self.w_reg
+
+        expression_reg = torch.mean(torch.square(expression_vector)) * self.w_reg
+        expression_reg += torch.mean(torch.square(expression_vector[1:] - expression_vector[:-1])) * 1e-1
+        output = mp_loss * self.w_mp + fan_loss + iris_loss + seg_mask_loss * self.w_seg + expression_reg
         return output
 
 
@@ -463,6 +468,23 @@ def process(self, image: np.ndarray, image_size: int=224):
         lmk = self.mica.flame.compute_landmarks(meshes)
         return meshes[0].detach().cpu(), code, lmk
 
+
+class Matting:
+    def __init__(self, script_path: str, chkp_path: str):
+        self.script_path = script_path
+        self.chkp_path = chkp_path
+    
+    def convert(self, video_path: str, output_mask_path: str):
+        args = "--variant mobilenetv3 " 
+        args += f"--checkpoint {self.chkp_path} " 
+        args += f"--input-source {video_path} "
+        args += "--output-type png_sequence "
+        args += f"--output-alpha {output_mask_path} "
+        args += "--device cuda"
+        cmd = f"python {self.script_path} {args}"
+        os.system(cmd)
+
+
 class DataSaver:
     def __init__(self, output_base: str, save_id_mesh: bool=True):
         self.output_base = output_base
@@ -541,6 +563,7 @@ def save_state(self,
         'cuda:0'
     )
 
+    matting = Matting(**conf.matting_kwargs)
     # create dataset
     # dataset = nir.get_dataset("SingleVideoDataset", **conf.video_dataset_kwargs)
 
@@ -551,13 +574,20 @@ def save_state(self,
     for filename in filenames:
         if not filename.endswith('mp4'):
             continue
+            
+        
 
         filepath = os.path.join(conf.base_dir, filename)
         print(f"Processing file: {filename}")
         dataset = nir.get_dataset("SingleVideoDataset", filepath=filepath, preload=True)
         dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, pin_memory=True, num_workers=1, collate_fn=nir.collate_fn)
 
         data_saver.set_output_state(filename.split('.')[0])
+
+        # preprocess whole video with matting
+        matting_alpha_path = data_saver.current_output_dir
+        print("estimating alpha masks")
+        matting.convert(filepath, matting_alpha_path)
 
         for frame_idx, data in tqdm(enumerate(dataloader), total=len(dataloader), desc="video progress"):
             data_saver.set_frame_index(frame_idx)

preproc_video.yaml

@@ -11,7 +11,7 @@ video_dataset_kwargs:
 
 
 flame_pose_expression_optimization_kwargs:
-  optim_iters: 2000
+  optim_iters: 1000
   log_result_only: true
   cam_init_z_trans: 0.5
 
@@ -40,7 +40,7 @@ flame_pose_expression_optimization_kwargs:
   loss_kwargs:
     w_mp: 0.3
     w_seg: 0.6
-    w_reg: 0.7
+    w_reg: 0.8
     wing_loss_kwargs: 
       omega: 10.0
       eps: 2.0
@@ -55,7 +55,7 @@ flame_pose_expression_optimization_kwargs:
     betas: [0.9, 0.999]
 
   sched_kwargs:
-    milestones: [200, 1000, 1500]
+    milestones: [200, 800, 1500]
     gamma: 0.1
 
   logger_kwargs:
@@ -64,3 +64,7 @@ flame_pose_expression_optimization_kwargs:
     port: 8097
     experiment_id: "Single video face parsing"
     log_iters: 500
+
+matting_kwargs:
+  script_path: /media/perukas/Home/_dev/_phd/libraries/RobustVideoMatting/inference.py
+  chkp_path: /media/perukas/Home/_dev/_phd/libraries/RobustVideoMatting/weights/rvm_mobilenetv3.pth