upgrade FaceAnimator module. Now uses https://github.com/wyhsirius/LIA …

…model
iperov · Sep 16, 2022 · 02de563 · 02de563
1 parent 42e835d
commit 02de563
Show file tree

Hide file tree

Showing 14 changed files with 144 additions and 201 deletions.
diff --git a/apps/DeepFaceLive/backend/FaceAligner.py b/apps/DeepFaceLive/backend/FaceAligner.py
@@ -3,6 +3,7 @@
 
 import numpy as np
 from xlib import os as lib_os
+from xlib.face import FRect
 from xlib.mp import csw as lib_csw
 from xlib.python import all_is_not_None
 
@@ -14,9 +15,11 @@
 class AlignMode(IntEnum):
     FROM_RECT = 0
     FROM_POINTS = 1
+    FROM_STATIC_RECT = 2
 
 AlignModeNames = ['@FaceAligner.AlignMode.FROM_RECT',
                   '@FaceAligner.AlignMode.FROM_POINTS',
+                  '@FaceAligner.AlignMode.FROM_STATIC_RECT',
                  ]
 
 class FaceAligner(BackendHost):
@@ -57,7 +60,7 @@ def on_start(self, weak_heap : BackendWeakHeap, reemit_frame_signal : BackendSig
         cs.align_mode.select(state.align_mode if state.align_mode is not None else AlignMode.FROM_POINTS)
 
         cs.face_coverage.enable()
-        cs.face_coverage.set_config(lib_csw.Number.Config(min=0.1, max=4.0, step=0.1, decimals=1, allow_instant_update=True))
+        cs.face_coverage.set_config(lib_csw.Number.Config(min=0.1, max=8.0, step=0.1, decimals=1, allow_instant_update=True))
         cs.face_coverage.set_number(state.face_coverage if state.face_coverage is not None else 2.2)
 
         cs.resolution.enable()
@@ -74,11 +77,11 @@ def on_start(self, weak_heap : BackendWeakHeap, reemit_frame_signal : BackendSig
         cs.freeze_z_rotation.set_flag(state.freeze_z_rotation if state.freeze_z_rotation is not None else False)
 
         cs.x_offset.enable()
-        cs.x_offset.set_config(lib_csw.Number.Config(min=-1, max=1, step=0.01, decimals=2, allow_instant_update=True))
+        cs.x_offset.set_config(lib_csw.Number.Config(min=-10, max=10, step=0.01, decimals=2, allow_instant_update=True))
         cs.x_offset.set_number(state.x_offset if state.x_offset is not None else 0)
 
         cs.y_offset.enable()
-        cs.y_offset.set_config(lib_csw.Number.Config(min=-1, max=1, step=0.01, decimals=2, allow_instant_update=True))
+        cs.y_offset.set_config(lib_csw.Number.Config(min=-10, max=10, step=0.01, decimals=2, allow_instant_update=True))
         cs.y_offset.set_number(state.y_offset if state.y_offset is not None else 0)
 
     def on_cs_align_mode(self, idx, align_mode):
@@ -164,18 +167,22 @@ def on_tick(self):
                         if face_ulmrks is not None:
                             fsi.face_resolution = state.resolution
 
+                            H, W = frame_image.shape[:2]
                             if state.align_mode == AlignMode.FROM_RECT:
                                 face_align_img, uni_mat = fsi.face_urect.cut(frame_image, coverage= state.face_coverage, output_size=state.resolution,
                                                                              x_offset=state.x_offset, y_offset=state.y_offset)
 
                             elif state.align_mode == AlignMode.FROM_POINTS:
                                 face_align_img, uni_mat = face_ulmrks.cut(frame_image, state.face_coverage, state.resolution,
-                                                                        exclude_moving_parts=state.exclude_moving_parts,
-                                                                        head_yaw=head_yaw,
-                                                                        x_offset=state.x_offset,
-                                                                        y_offset=state.y_offset-0.08,
-                                                                        freeze_z_rotation=state.freeze_z_rotation)
-
+                                                                          exclude_moving_parts=state.exclude_moving_parts,
+                                                                          head_yaw=head_yaw,
+                                                                          x_offset=state.x_offset,
+                                                                          y_offset=state.y_offset-0.08,
+                                                                          freeze_z_rotation=state.freeze_z_rotation)
+                            elif state.align_mode == AlignMode.FROM_STATIC_RECT:
+                                rect = FRect.from_ltrb([ 0.5 - (fsi.face_resolution/W)/2, 0.5 - (fsi.face_resolution/H)/2, 0.5 + (fsi.face_resolution/W)/2, 0.5 + (fsi.face_resolution/H)/2,])
+                                face_align_img, uni_mat = rect.cut(frame_image, coverage= state.face_coverage, output_size=state.resolution,
+                                                                             x_offset=state.x_offset, y_offset=state.y_offset)
 
                             fsi.face_align_image_name = f'{frame_image_name}_{face_id}_aligned'
                             fsi.image_to_align_uni_mat = uni_mat

diff --git a/apps/DeepFaceLive/backend/FaceAnimator.py b/apps/DeepFaceLive/backend/FaceAnimator.py
@@ -1,10 +1,8 @@
-import re
 import time
 from pathlib import Path
 
-import cv2
 import numpy as np
-from modelhub.onnx import TPSMM
+from modelhub.onnx import LIA
 from xlib import cv as lib_cv2
 from xlib import os as lib_os
 from xlib import path as lib_path
@@ -29,7 +27,7 @@ def __init__(self, weak_heap : BackendWeakHeap, reemit_frame_signal : BackendSig
     def get_control_sheet(self) -> 'Sheet.Host': return super().get_control_sheet()
 
     def _get_name(self):
-        return super()._get_name()# + f'{self._id}'
+        return super()._get_name()
 
 class FaceAnimatorWorker(BackendWorker):
     def get_state(self) -> 'WorkerState': return super().get_state()
@@ -44,11 +42,10 @@ def on_start(self, weak_heap : BackendWeakHeap, reemit_frame_signal : BackendSig
 
         self.pending_bcd = None
 
-        self.tpsmm_model = None
+        self.lia_model : LIA = None
 
         self.animatable_img = None
-        self.driving_ref_kp = None
-        self.last_driving_kp = None
+        self.driving_ref_motion = None
 
         lib_os.set_timer_resolution(1)
 
@@ -58,14 +55,12 @@ def on_start(self, weak_heap : BackendWeakHeap, reemit_frame_signal : BackendSig
         cs.animatable.call_on_selected(self.on_cs_animatable)
 
         cs.animator_face_id.call_on_number(self.on_cs_animator_face_id)
-        cs.relative_mode.call_on_flag(self.on_cs_relative_mode)
         cs.relative_power.call_on_number(self.on_cs_relative_power)
         cs.update_animatables.call_on_signal(self.update_animatables)
         cs.reset_reference_pose.call_on_signal(self.on_cs_reset_reference_pose)
 
-
         cs.device.enable()
-        cs.device.set_choices( TPSMM.get_available_devices(), none_choice_name='@misc.menu_select')
+        cs.device.set_choices( LIA.get_available_devices(), none_choice_name='@misc.menu_select')
         cs.device.select(state.device)
 
     def update_animatables(self):
@@ -76,7 +71,7 @@ def update_animatables(self):
     def on_cs_device(self, idx, device):
         state, cs = self.get_state(), self.get_control_sheet()
         if device is not None and state.device == device:
-            self.tpsmm_model = TPSMM(device)
+            self.lia_model = LIA(device)
 
             cs.animatable.enable()
             self.update_animatables()
@@ -85,12 +80,9 @@ def on_cs_device(self, idx, device):
             cs.animator_face_id.enable()
             cs.animator_face_id.set_config(lib_csw.Number.Config(min=0, max=16, step=1, decimals=0, allow_instant_update=True))
             cs.animator_face_id.set_number(state.animator_face_id if state.animator_face_id is not None else 0)
-
-            cs.relative_mode.enable()
-            cs.relative_mode.set_flag(state.relative_mode if state.relative_mode is not None else True)
-
+
             cs.relative_power.enable()
-            cs.relative_power.set_config(lib_csw.Number.Config(min=0.0, max=1.0, step=0.01, decimals=2, allow_instant_update=True))
+            cs.relative_power.set_config(lib_csw.Number.Config(min=0.0, max=2.0, step=0.01, decimals=2, allow_instant_update=True))
             cs.relative_power.set_number(state.relative_power if state.relative_power is not None else 1.0)
 
             cs.update_animatables.enable()
@@ -105,20 +97,15 @@ def on_cs_animatable(self, idx, animatable):
 
         state.animatable = animatable
         self.animatable_img = None
-        self.animatable_kp = None
-        self.driving_ref_kp = None
+        self.driving_ref_motion = None
 
         if animatable is not None:
             try:
-                W,H = self.tpsmm_model.get_input_size()
+                W,H = self.lia_model.get_input_size()
                 ip = ImageProcessor(lib_cv2.imread(self.animatables_path / animatable))
                 ip.fit_in(TW=W, TH=H, pad_to_target=True, allow_upscale=True)
 
-                animatable_img = ip.get_image('HWC')
-                animatable_kp = self.tpsmm_model.extract_kp(animatable_img)
-
-                self.animatable_img = animatable_img
-                self.animatable_kp = animatable_kp
+                self.animatable_img = ip.get_image('HWC')
             except Exception as e:
                 cs.animatable.unselect()
 
@@ -133,13 +120,6 @@ def on_cs_animator_face_id(self, animator_face_id):
         cs.animator_face_id.set_number(animator_face_id)
         self.save_state()
         self.reemit_frame_signal.send()
-
-    def on_cs_relative_mode(self, relative_mode):
-        state, cs = self.get_state(), self.get_control_sheet()
-        state.relative_mode = relative_mode
-        self.save_state()
-        self.reemit_frame_signal.send()
-
     def on_cs_relative_power(self, relative_power):
         state, cs = self.get_state(), self.get_control_sheet()
         cfg = cs.relative_power.get_config()
@@ -149,7 +129,7 @@ def on_cs_relative_power(self, relative_power):
         self.reemit_frame_signal.send()
 
     def on_cs_reset_reference_pose(self):
-        self.driving_ref_kp = self.last_driving_kp
+        self.driving_ref_motion = None
         self.reemit_frame_signal.send()
 
     def on_tick(self):
@@ -162,8 +142,8 @@ def on_tick(self):
             if bcd is not None:
                 bcd.assign_weak_heap(self.weak_heap)
 
-                tpsmm_model = self.tpsmm_model
-                if tpsmm_model is not None and self.animatable_img is not None:
+                lia_model = self.lia_model
+                if lia_model is not None and self.animatable_img is not None:
 
                     for i, fsi in enumerate(bcd.get_face_swap_info_list()):
                         if state.animator_face_id == i:
@@ -172,14 +152,10 @@ def on_tick(self):
 
                                 _,H,W,_ = ImageProcessor(face_align_image).get_dims()
 
-                                driving_kp = self.last_driving_kp = tpsmm_model.extract_kp(face_align_image)
-
-                                if self.driving_ref_kp is None:
-                                    self.driving_ref_kp = driving_kp
+                                if self.driving_ref_motion is None:
+                                    self.driving_ref_motion = lia_model.extract_motion(face_align_image)
 
-                                anim_image = tpsmm_model.generate(self.animatable_img, self.animatable_kp, driving_kp,
-                                                                  self.driving_ref_kp if state.relative_mode else None,
-                                                                  relative_power=state.relative_power)
+                                anim_image = lia_model.generate(self.animatable_img, face_align_image, self.driving_ref_motion, power=state.relative_power)
                                 anim_image = ImageProcessor(anim_image).resize((W,H)).get_image('HWC')
 
                                 fsi.face_swap_image_name = f'{fsi.face_align_image_name}_swapped'
@@ -203,7 +179,6 @@ def __init__(self):
             self.device = lib_csw.DynamicSingleSwitch.Client()
             self.animatable = lib_csw.DynamicSingleSwitch.Client()
             self.animator_face_id = lib_csw.Number.Client()
-            self.relative_mode = lib_csw.Flag.Client()
             self.update_animatables = lib_csw.Signal.Client()
             self.reset_reference_pose = lib_csw.Signal.Client()
             self.relative_power = lib_csw.Number.Client()
@@ -214,7 +189,6 @@ def __init__(self):
             self.device = lib_csw.DynamicSingleSwitch.Host()
             self.animatable = lib_csw.DynamicSingleSwitch.Host()
             self.animator_face_id = lib_csw.Number.Host()
-            self.relative_mode = lib_csw.Flag.Host()
             self.update_animatables = lib_csw.Signal.Host()
             self.reset_reference_pose = lib_csw.Signal.Host()
             self.relative_power = lib_csw.Number.Host()
@@ -223,5 +197,4 @@ class WorkerState(BackendWorkerState):
     device = None
     animatable : str = None
     animator_face_id : int = None
-    relative_mode : bool = None
     relative_power : float = None
diff --git a/apps/DeepFaceLive/ui/QFaceAnimator.py b/apps/DeepFaceLive/ui/QFaceAnimator.py
@@ -31,10 +31,8 @@ def __init__(self, backend : FaceAnimator, animatables_path : Path):
 
         q_animator_face_id_label = QLabelPopupInfo(label=L('@QFaceAnimator.animator_face_id') )
         q_animator_face_id       = QSpinBoxCSWNumber(cs.animator_face_id, reflect_state_widgets=[q_animator_face_id_label])
-
-        q_relative_mode_label = QLabelPopupInfo(label=L('@QFaceAnimator.relative_mode') )
-        q_relative_mode = QCheckBoxCSWFlag(cs.relative_mode, reflect_state_widgets=[q_relative_mode_label])
-
+
+        q_relative_power_label = QLabelPopupInfo(label=L('@QFaceAnimator.relative_power') )
         q_relative_power = QSliderCSWNumber(cs.relative_power)
 
         q_update_animatables = QXPushButtonCSWSignal(cs.update_animatables, image=QXImageDB.reload_outline('light gray'), button_size=(24,22) )
@@ -52,9 +50,8 @@ def __init__(self, backend : FaceAnimator, animatables_path : Path):
         grid_l.addWidget(q_animator_face_id_label, row, 0, alignment=qtx.AlignRight | qtx.AlignVCenter  )
         grid_l.addWidget(q_animator_face_id, row, 1, alignment=qtx.AlignLeft )
         row += 1
-        grid_l.addWidget(q_relative_mode_label, row, 0, alignment=qtx.AlignRight | qtx.AlignVCenter  )
-        grid_l.addLayout(qtx.QXHBoxLayout([q_relative_mode,2,q_relative_power]), row, 1, alignment=qtx.AlignLeft )
-
+        grid_l.addWidget(q_relative_power_label, row, 0, alignment=qtx.AlignRight | qtx.AlignVCenter  )
+        grid_l.addWidget(q_relative_power, row, 1 )
 
         row += 1
         grid_l.addWidget(q_reset_reference_pose, row, 0, 1, 2  )

diff --git a/localization/localization.py b/localization/localization.py
@@ -656,13 +656,13 @@ def localize(s : Union[str, None]) -> Union[str, None]:
                 'it-IT' : 'Animatore Face ID',
                 'ja-JP' : '動かす顔のID番号'},
 
-    'QFaceAnimator.relative_mode':{
-                'en-US' : 'Relative mode',
-                'ru-RU' : 'Относительный режим',
-                'zh-CN' : '相对模式',
-                'es-ES' : 'Modo relativo',
-                'it-IT' : 'Modalità relativa',
-                'ja-JP' : '相対モード'},
+    'QFaceAnimator.relative_power':{
+                'en-US' : 'Relative power',
+                'ru-RU' : 'Относительная сила',
+                'zh-CN' : 'Relative power',
+                'es-ES' : 'Relative power',
+                'it-IT' : 'Relative power',
+                'ja-JP' : 'Relative power'},
 
     'QFaceAnimator.reset_reference_pose':{
                 'en-US' : 'Reset reference pose',
@@ -1143,7 +1143,15 @@ def localize(s : Union[str, None]) -> Union[str, None]:
                 'es-ES' : 'De los puntos',
                 'it-IT' : 'Da punti',
                 'ja-JP' : '点から'},
-
+
+    'FaceAligner.AlignMode.FROM_STATIC_RECT':{
+                'en-US' : 'From static rect',
+                'ru-RU' : 'Из статичного прямоугольника',
+                'zh-CN' : '从一个静态的矩形',
+                'es-ES' : 'From static rect',
+                'it-IT' : 'From static rect',
+                'ja-JP' : 'From static rect'},
+
     'FaceSwapper.model_information':{
                 'en-US' : 'Model information',
                 'ru-RU' : 'Информация о модели',