fix FaceMerger

This commit is contained in:
iperov 2021-10-01 18:10:44 +04:00
commit 933bb99fb9

View file

@ -37,14 +37,11 @@ class FaceMergerWorker(BackendWorker):
def get_control_sheet(self) -> 'Sheet.Worker': return super().get_control_sheet() 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): def on_start(self, weak_heap : BackendWeakHeap, reemit_frame_signal : BackendSignal, bc_in : BackendConnection, bc_out : BackendConnection):
self.weak_heap = weak_heap self.weak_heap = weak_heap
self.reemit_frame_signal = reemit_frame_signal self.reemit_frame_signal = reemit_frame_signal
self.bc_in = bc_in self.bc_in = bc_in
self.bc_out = bc_out self.bc_out = bc_out
self.pending_bcd = None self.pending_bcd = None
self.is_gpu = False
lib_os.set_timer_resolution(1) lib_os.set_timer_resolution(1)
state, cs = self.get_state(), self.get_control_sheet() state, cs = self.get_state(), self.get_control_sheet()
@ -66,6 +63,11 @@ class FaceMergerWorker(BackendWorker):
def on_cs_device(self, idxs, device : lib_cl.DeviceInfo): def on_cs_device(self, idxs, device : lib_cl.DeviceInfo):
state, cs = self.get_state(), self.get_control_sheet() state, cs = self.get_state(), self.get_control_sheet()
if device is not None and state.device == device: 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.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_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) cs.face_x_offset.set_number(state.face_x_offset if state.face_x_offset is not None else 0.0)
@ -191,28 +193,29 @@ class FaceMergerWorker(BackendWorker):
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_mask_img, face_swap_img, face_swap_mask_img, aligned_to_source_uni_mat, frame_width, frame_height ):
state = self.get_state() state = self.get_state()
if state.face_mask_type == FaceMaskType.SRC: 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) 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: 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) 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: 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);', 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_align_mask_img),
lib_cl.Tensor.from_value(face_swap_mask_img, device=state.device), lib_cl.Tensor.from_value(face_swap_mask_img),
dtype=np.uint8).transpose( (2,0,1) ) 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_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) 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_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_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 opacity = state.face_opacity
if opacity == 1.0: if opacity == 1.0: