diff --git a/apps/DeepFaceLive/backend/FaceMerger.py b/apps/DeepFaceLive/backend/FaceMerger.py index 02e3f23..149909e 100644 --- a/apps/DeepFaceLive/backend/FaceMerger.py +++ b/apps/DeepFaceLive/backend/FaceMerger.py @@ -37,14 +37,11 @@ class FaceMergerWorker(BackendWorker): def get_control_sheet(self) -> 'Sheet.Worker': return super().get_control_sheet() def on_start(self, weak_heap : BackendWeakHeap, reemit_frame_signal : BackendSignal, bc_in : BackendConnection, bc_out : BackendConnection): - - self.weak_heap = weak_heap self.reemit_frame_signal = reemit_frame_signal self.bc_in = bc_in self.bc_out = bc_out self.pending_bcd = None - self.is_gpu = False lib_os.set_timer_resolution(1) state, cs = self.get_state(), self.get_control_sheet() @@ -57,7 +54,7 @@ class FaceMergerWorker(BackendWorker): cs.face_opacity.call_on_number(self.on_cs_face_opacity) cs.device.enable() - + cs.device.set_choices( ['CPU'] + lib_cl.get_available_devices_info(), none_choice_name='@misc.menu_select') cs.device.select(state.device if state.device is not None else 'CPU') @@ -66,6 +63,11 @@ class FaceMergerWorker(BackendWorker): def on_cs_device(self, idxs, device : lib_cl.DeviceInfo): state, cs = self.get_state(), self.get_control_sheet() if device is not None and state.device == device: + if device != 'CPU': + dev = lib_cl.get_device(device) + dev.set_target_memory_usage(mb=512) + lib_cl.set_default_device(dev) + cs.face_x_offset.enable() cs.face_x_offset.set_config(lib_csw.Number.Config(min=-0.5, max=0.5, step=0.001, decimals=3, allow_instant_update=True)) cs.face_x_offset.set_number(state.face_x_offset if state.face_x_offset is not None else 0.0) @@ -154,15 +156,15 @@ class FaceMergerWorker(BackendWorker): cs.face_opacity.set_number(face_opacity) self.save_state() self.reemit_frame_signal.send() - + 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 ): state = self.get_state() - + frame_image = ImageProcessor(frame_image).to_ufloat32().get_image('HWC') face_align_mask_img = ImageProcessor(face_align_mask_img).to_ufloat32().get_image('HW') face_swap_mask_img = ImageProcessor(face_swap_mask_img).to_ufloat32().get_image('HW') - + if state.face_mask_type == FaceMaskType.SRC: face_mask = face_align_mask_img elif state.face_mask_type == FaceMaskType.CELEB: @@ -185,44 +187,45 @@ class FaceMergerWorker(BackendWorker): frame_final = ne.evaluate('frame_image*(1.0-frame_face_mask) + frame_face_swap_img*frame_face_mask') else: frame_final = ne.evaluate('frame_image*(1.0-frame_face_mask) + frame_image*frame_face_mask*(1.0-opacity) + frame_face_swap_img*frame_face_mask*opacity') - + return frame_final - + 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 ): state = self.get_state() - + + if state.face_mask_type == FaceMaskType.SRC: - face_mask_t = lib_cl.Tensor.from_value(face_align_mask_img, device=state.device) + face_mask_t = lib_cl.Tensor.from_value(face_align_mask_img) 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, device=state.device) + face_mask_t = lib_cl.Tensor.from_value(face_swap_mask_img) 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);', - lib_cl.Tensor.from_value(face_align_mask_img, device=state.device), - lib_cl.Tensor.from_value(face_swap_mask_img, device=state.device), + face_mask_t = lib_cl.any_wise('float X = (((float)I0) / 255.0) * (((float)I1) / 255.0); O = (X <= 0.5 ? 0 : 1);', + lib_cl.Tensor.from_value(face_align_mask_img), + lib_cl.Tensor.from_value(face_swap_mask_img), 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, device=state.device) + + face_swap_img_t = lib_cl.Tensor.from_value(face_swap_img) face_swap_img_t = face_swap_img_t.transpose( (2,0,1), op_text='O = ((O_TYPE)I) / 255.0', dtype=np.float32) - + frame_face_mask_t = lib_cl.remap_np_affine(face_mask_t, aligned_to_source_uni_mat, output_size=(frame_height, frame_width) ) frame_face_swap_img_t = lib_cl.remap_np_affine(face_swap_img_t, aligned_to_source_uni_mat, output_size=(frame_height, frame_width) ) - - frame_image_t = lib_cl.Tensor.from_value(frame_image, device=state.device).transpose( (2,0,1) ) - + + frame_image_t = lib_cl.Tensor.from_value(frame_image).transpose( (2,0,1) ) + opacity = state.face_opacity if opacity == 1.0: frame_final_t = lib_cl.any_wise('float I0f = (((float)I0) / 255.0); I1 = (I1 <= (1.0/255.0) ? 0.0 : I1 > 1.0 ? 1.0 : I1); O = I0f*(1.0-I1) + I2*I1', frame_image_t, frame_face_mask_t, frame_face_swap_img_t, dtype=np.float32) else: frame_final_t = lib_cl.any_wise('float I0f = (((float)I0) / 255.0); I1 = (I1 <= (1.0/255.0) ? 0.0 : I1 > 1.0 ? 1.0 : I1); O = I0f*(1.0-I1) + I0f*I1*(1.0-I3) + I2*I1*I3', frame_image_t, frame_face_mask_t, frame_face_swap_img_t, opacity, dtype=np.float32) - + return frame_final_t.transpose( (1,2,0) ).np() - - + + def on_tick(self): state, cs = self.get_state(), self.get_control_sheet() @@ -260,13 +263,13 @@ class FaceMergerWorker(BackendWorker): aligned_to_source_uni_mat = aligned_to_source_uni_mat.source_translated(-state.face_x_offset, -state.face_y_offset) aligned_to_source_uni_mat = aligned_to_source_uni_mat.source_scaled_around_center(state.face_scale,state.face_scale) 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 ) 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 ) # keep image in float32 in order not to extra load FaceMerger - + merged_frame_name = f'{frame_name}_merged' bcd.set_merged_frame_name(merged_frame_name) bcd.set_image(merged_frame_name, merged_frame)