FaceMerger: added color_transfer rct

apps/DeepFaceLive/backend/FaceMerger.py

@@ -5,6 +5,7 @@ import numexpr as ne
 import numpy as np
 from xlib import avecl as lib_cl
 from xlib import os as lib_os
+from xlib.image import color_transfer as lib_ct
 from xlib.image import ImageProcessor
 from xlib.mp import csw as lib_csw
 from xlib.python import all_is_not_None
@@ -93,6 +94,11 @@ class FaceMergerWorker(BackendWorker):
             cs.face_mask_blur.set_config(lib_csw.Number.Config(min=0, max=400, step=1, decimals=0, allow_instant_update=True))
             cs.face_mask_blur.set_number(state.face_mask_blur if state.face_mask_blur is not None else 25.0)
 
+            cs.color_transfer.call_on_selected(self.on_cs_color_transfer)
+            cs.color_transfer.enable()
+            cs.color_transfer.set_choices(['none','rct'])
+            cs.color_transfer.select(state.color_transfer if state.color_transfer is not None else 'none')
+
             cs.interpolation.call_on_selected(self.on_cs_interpolation)
             cs.interpolation.enable()
             cs.interpolation.set_choices(['bilinear','bicubic','lanczos4'], none_choice_name=None)
@@ -107,7 +113,6 @@ class FaceMergerWorker(BackendWorker):
             self.save_state()
             self.restart()
 
-
     def on_cs_face_x_offset(self, face_x_offset):
         state, cs = self.get_state(), self.get_control_sheet()
         cfg = cs.face_x_offset.get_config()
@@ -154,6 +159,12 @@ class FaceMergerWorker(BackendWorker):
         self.save_state()
         self.reemit_frame_signal.send()
 
+    def on_cs_color_transfer(self, idx, color_transfer):
+        state, cs = self.get_state(), self.get_control_sheet()
+        state.color_transfer = color_transfer
+        self.save_state()
+        self.reemit_frame_signal.send()
+
     def on_cs_interpolation(self, idx, interpolation):
         state, cs = self.get_state(), self.get_control_sheet()
         state.interpolation = interpolation
@@ -171,7 +182,7 @@ class FaceMergerWorker(BackendWorker):
     _cpu_interp = {'bilinear' : ImageProcessor.Interpolation.LINEAR,
                    'bicubic'  : ImageProcessor.Interpolation.CUBIC,
                    'lanczos4' : ImageProcessor.Interpolation.LANCZOS4}
-    def _merge_on_cpu(self, frame_image, face_align_mask_img, face_swap_img, face_swap_mask_img, aligned_to_source_uni_mat, frame_width, frame_height ):
+    def _merge_on_cpu(self, frame_image, face_align_img, face_align_mask_img, face_swap_img, face_swap_mask_img, aligned_to_source_uni_mat, frame_width, frame_height ):
         state = self.get_state()
 
         interpolation = self._cpu_interp[state.interpolation]
@@ -188,13 +199,17 @@ class FaceMergerWorker(BackendWorker):
             face_mask = face_align_mask_img*face_swap_mask_img
 
         # Combine face mask
-        face_mask_ip = ImageProcessor(face_mask).erode_blur(state.face_mask_erode, state.face_mask_blur, fade_to_border=True) \
+        face_mask = ImageProcessor(face_mask).erode_blur(state.face_mask_erode, state.face_mask_blur, fade_to_border=True).get_image('HWC')
-                                                .warpAffine(aligned_to_source_uni_mat, frame_width, frame_height)
-        face_mask_ip.clip2( (1.0/255.0), 0.0, 1.0, 1.0)
-        frame_face_mask = face_mask_ip.get_image('HWC')
 
-        frame_face_swap_img = ImageProcessor(face_swap_img) \
+        frame_face_mask = ImageProcessor(face_mask).warpAffine(aligned_to_source_uni_mat, frame_width, frame_height).clip2( (1.0/255.0), 0.0, 1.0, 1.0).get_image('HWC')
-                                .to_ufloat32().warpAffine(aligned_to_source_uni_mat, frame_width, frame_height, interpolation=interpolation).get_image('HWC')
+
+        face_swap_img = ImageProcessor(face_swap_img).to_ufloat32().get_image('HWC')
+
+        if state.color_transfer == 'rct':
+            face_align_img = ImageProcessor(face_align_img).to_ufloat32().get_image('HWC')
+            face_swap_img = lib_ct.rct(face_swap_img, face_align_img, target_mask=face_mask, source_mask=face_mask)
+
+        frame_face_swap_img = ImageProcessor(face_swap_img).warpAffine(aligned_to_source_uni_mat, frame_width, frame_height, interpolation=interpolation).get_image('HWC')
 
         # Combine final frame
         opacity = state.face_opacity
@@ -209,16 +224,14 @@ class FaceMergerWorker(BackendWorker):
                    'bicubic'  : lib_cl.EInterpolation.CUBIC,
                    'lanczos4' : lib_cl.EInterpolation.LANCZOS4}
 
-    def _merge_on_gpu(self, frame_image, face_align_mask_img, face_swap_img, face_swap_mask_img, aligned_to_source_uni_mat, frame_width, frame_height ):
+    def _merge_on_gpu(self, frame_image, face_align_img, face_align_mask_img, face_swap_img, face_swap_mask_img, aligned_to_source_uni_mat, frame_width, frame_height ):
         state = self.get_state()
         interpolation = self._gpu_interp[state.interpolation]
 
         if state.face_mask_type == FaceMaskType.SRC:
-            face_mask_t = lib_cl.Tensor.from_value(face_align_mask_img)
+            face_mask_t = lib_cl.Tensor.from_value(face_align_mask_img).transpose( (2,0,1), op_text='O = (I < 128 ? 0 : 1);', dtype=np.uint8)
-            face_mask_t = face_mask_t.transpose( (2,0,1), op_text='O = (I < 128 ? 0 : 1);', dtype=np.uint8)
         elif state.face_mask_type == FaceMaskType.CELEB:
-            face_mask_t = lib_cl.Tensor.from_value(face_swap_mask_img)
+            face_mask_t = lib_cl.Tensor.from_value(face_swap_mask_img).transpose( (2,0,1), op_text='O = (I < 128 ? 0 : 1);', dtype=np.uint8)
-            face_mask_t = face_mask_t.transpose( (2,0,1), op_text='O = (I < 128 ? 0 : 1);', dtype=np.uint8)
 
         elif state.face_mask_type == FaceMaskType.SRC_M_CELEB:
             face_mask_t = lib_cl.any_wise('float X = (((float)I0) / 255.0) * (((float)I1) / 255.0); O = (X <= 0.5 ? 0 : 1);',
@@ -227,9 +240,11 @@ class FaceMergerWorker(BackendWorker):
                                           dtype=np.uint8).transpose( (2,0,1) )
 
         face_mask_t = lib_cl.binary_morph(face_mask_t, state.face_mask_erode, state.face_mask_blur, fade_to_border=True, dtype=np.float32)
+        face_swap_img_t  = lib_cl.Tensor.from_value(face_swap_img ).transpose( (2,0,1), op_text='O = ((O_TYPE)I) / 255.0', dtype=np.float32)
 
-        face_swap_img_t = lib_cl.Tensor.from_value(face_swap_img)
+        if state.color_transfer == 'rct':
-        face_swap_img_t = face_swap_img_t.transpose( (2,0,1), op_text='O = ((O_TYPE)I) / 255.0', dtype=np.float32)
+            face_align_img_t = lib_cl.Tensor.from_value(face_align_img).transpose( (2,0,1), op_text='O = ((O_TYPE)I) / 255.0', dtype=np.float32)
+            face_swap_img_t = lib_cl.rct(face_swap_img_t, face_align_img_t, target_mask_t=face_mask_t, source_mask_t=face_mask_t)
 
         frame_face_mask_t     = lib_cl.remap_np_affine(face_mask_t,     aligned_to_source_uni_mat, interpolation=lib_cl.EInterpolation.LINEAR, output_size=(frame_height, frame_width), post_op_text='O = (O <= (1.0/255.0) ? 0.0 : O > 1.0 ? 1.0 : O);' )
         frame_face_swap_img_t = lib_cl.remap_np_affine(face_swap_img_t, aligned_to_source_uni_mat, interpolation=interpolation, output_size=(frame_height, frame_width), post_op_text='O = clamp(O, 0.0, 1.0);' )
@@ -244,7 +259,6 @@ class FaceMergerWorker(BackendWorker):
 
         return frame_final_t.transpose( (1,2,0) ).np()
 
-
     def on_tick(self):
         state, cs = self.get_state(), self.get_control_sheet()
 
@@ -260,14 +274,14 @@ class FaceMergerWorker(BackendWorker):
 
                 if frame_image is not None:
                     for fsi in bcd.get_face_swap_info_list():
-                        face_align_img_shape, _ = bcd.get_image_shape_dtype(fsi.face_align_image_name)
+                        face_align_img      = bcd.get_image(fsi.face_align_image_name)
                         face_align_mask_img = bcd.get_image(fsi.face_align_mask_name)
                         face_swap_img       = bcd.get_image(fsi.face_swap_image_name)
                         face_swap_mask_img  = bcd.get_image(fsi.face_swap_mask_name)
                         image_to_align_uni_mat = fsi.image_to_align_uni_mat
 
-                        if all_is_not_None(face_align_img_shape, face_align_mask_img, face_swap_img, face_swap_mask_img, image_to_align_uni_mat):
+                        if all_is_not_None(face_align_img, face_align_mask_img, face_swap_img, face_swap_mask_img, image_to_align_uni_mat):
-                            face_height, face_width = face_align_img_shape[:2]
+                            face_height, face_width = face_align_img.shape[:2]
                             frame_height, frame_width = frame_image.shape[:2]
                             aligned_to_source_uni_mat = image_to_align_uni_mat.invert()
                             aligned_to_source_uni_mat = aligned_to_source_uni_mat.source_translated(-state.face_x_offset, -state.face_y_offset)
@@ -275,9 +289,9 @@ class FaceMergerWorker(BackendWorker):
                             aligned_to_source_uni_mat = aligned_to_source_uni_mat.to_exact_mat (face_width, face_height, frame_width, frame_height)
 
                             if state.device == 'CPU':
-                                merged_frame = self._merge_on_cpu(frame_image, face_align_mask_img, face_swap_img, face_swap_mask_img, aligned_to_source_uni_mat, frame_width, frame_height )
+                                merged_frame = self._merge_on_cpu(frame_image, face_align_img, face_align_mask_img, face_swap_img, face_swap_mask_img, aligned_to_source_uni_mat, frame_width, frame_height )
                             else:
-                                merged_frame = self._merge_on_gpu(frame_image, face_align_mask_img, face_swap_img, face_swap_mask_img, aligned_to_source_uni_mat, frame_width, frame_height )
+                                merged_frame = self._merge_on_gpu(frame_image, face_align_img, face_align_mask_img, face_swap_img, face_swap_mask_img, aligned_to_source_uni_mat, frame_width, frame_height )
                             # keep image in float32 in order not to extra load FaceMerger
 
                             merged_image_name = f'{frame_image_name}_merged'
@@ -310,6 +324,7 @@ class Sheet:
             self.face_mask_scale = lib_csw.Number.Client()
             self.face_mask_erode = lib_csw.Number.Client()
             self.face_mask_blur = lib_csw.Number.Client()
+            self.color_transfer = lib_csw.DynamicSingleSwitch.Client()
             self.interpolation = lib_csw.DynamicSingleSwitch.Client()
             self.face_opacity = lib_csw.Number.Client()
 
@@ -323,6 +338,7 @@ class Sheet:
             self.face_mask_type = lib_csw.DynamicSingleSwitch.Host()
             self.face_mask_erode = lib_csw.Number.Host()
             self.face_mask_blur = lib_csw.Number.Host()
+            self.color_transfer = lib_csw.DynamicSingleSwitch.Host()
             self.interpolation = lib_csw.DynamicSingleSwitch.Host()
             self.face_opacity = lib_csw.Number.Host()
 
@@ -334,6 +350,6 @@ class WorkerState(BackendWorkerState):
     face_mask_type = None
     face_mask_erode : int = None
     face_mask_blur : int = None
+    color_transfer = None
     interpolation = None
     face_opacity : float = None
-

apps/DeepFaceLive/ui/QFaceMerger.py

@@ -37,6 +37,9 @@ class QFaceMerger(QBackendPanel):
         q_face_mask_blur_label = QLabelPopupInfo(label=L('@QFaceMerger.face_mask_blur') )
         q_face_mask_blur       = QSpinBoxCSWNumber(cs.face_mask_blur, reflect_state_widgets=[q_face_mask_blur_label])
         
+        q_color_transfer_label = QLabelPopupInfo(label=L('@QFaceMerger.color_transfer') )
+        q_color_transfer       = QComboBoxCSWDynamicSingleSwitch(cs.color_transfer, reflect_state_widgets=[q_color_transfer_label])
+
         q_interpolation_label = QLabelPopupInfo(label=L('@QFaceMerger.interpolation') )
         q_interpolation       = QComboBoxCSWDynamicSingleSwitch(cs.interpolation, reflect_state_widgets=[q_interpolation_label])
 
@@ -60,6 +63,9 @@ class QFaceMerger(QBackendPanel):
         grid_l.addLayout(lib_qt.QXVBoxLayout([q_face_mask_erode_label,q_face_mask_blur_label]), row, 0, alignment=Qt.AlignmentFlag.AlignRight | Qt.AlignmentFlag.AlignVCenter  )
         grid_l.addLayout(lib_qt.QXHBoxLayout([q_face_mask_erode,q_face_mask_blur], spacing=3), row, 1, alignment=Qt.AlignmentFlag.AlignLeft )
         row += 1
+        grid_l.addWidget(q_color_transfer_label, row, 0, alignment=Qt.AlignmentFlag.AlignRight | Qt.AlignmentFlag.AlignVCenter  )
+        grid_l.addWidget(q_color_transfer, row, 1, alignment=Qt.AlignmentFlag.AlignLeft )
+        row += 1
         grid_l.addWidget(q_interpolation_label, row, 0, alignment=Qt.AlignmentFlag.AlignRight | Qt.AlignmentFlag.AlignVCenter  )
         grid_l.addWidget(q_interpolation, row, 1, alignment=Qt.AlignmentFlag.AlignLeft )
         row += 1

localization/localization.py

@@ -547,6 +547,11 @@ class Localization:
                 'ru-RU' : 'Размытие маски',
                 'zh-CN' : '遮罩边缘羽化'},
                 
+    'QFaceMerger.color_transfer':{
+                'en-US' : 'Color transfer',
+                'ru-RU' : 'Перенос цвета',
+                'zh-CN' : '彩色转印'},
+                
     'QFaceMerger.interpolation':{
                 'en-US' : 'Interpolation',
                 'ru-RU' : 'Интерполяция',