From b74c2b1f9a5555898c834bfb8d8d621d28db94b8 Mon Sep 17 00:00:00 2001 From: Colombo Date: Thu, 14 Nov 2019 12:08:27 +0400 Subject: [PATCH] Converter: added new color transfer mode: mix-m --- converters/ConvertMasked.py | 6 +++++- converters/ConverterConfig.py | 4 ++-- imagelib/__init__.py | 2 +- imagelib/color_transfer.py | 36 ++++++++++++++++++++++++++++++++--- 4 files changed, 41 insertions(+), 7 deletions(-) diff --git a/converters/ConvertMasked.py b/converters/ConvertMasked.py index 2520c50..12d799f 100644 --- a/converters/ConvertMasked.py +++ b/converters/ConvertMasked.py @@ -208,6 +208,8 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i elif cfg.color_transfer_mode == 7: #sot-m prd_face_bgr = imagelib.color_transfer_sot (prd_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a) prd_face_bgr = np.clip (prd_face_bgr, 0.0, 1.0) + elif cfg.color_transfer_mode == 8: #mix-m + prd_face_bgr = imagelib.color_transfer_mix (prd_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a) if cfg.mode == 'hist-match-bw': prd_face_bgr = cv2.cvtColor(prd_face_bgr, cv2.COLOR_BGR2GRAY) @@ -293,7 +295,9 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i elif cfg.color_transfer_mode == 7: #sot-m out_face_bgr = imagelib.color_transfer_sot (out_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a) out_face_bgr = np.clip (out_face_bgr, 0.0, 1.0) - + elif cfg.color_transfer_mode == 8: #mix-m + out_face_bgr = imagelib.color_transfer_mix (out_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a) + if cfg.mode == 'seamless-hist-match': out_face_bgr = imagelib.color_hist_match(out_face_bgr, dst_face_bgr, cfg.hist_match_threshold) diff --git a/converters/ConverterConfig.py b/converters/ConverterConfig.py index 3dbe64f..b7fb095 100644 --- a/converters/ConverterConfig.py +++ b/converters/ConverterConfig.py @@ -127,8 +127,8 @@ half_face_mask_mode_dict = {1:'learned', 4:'FAN-dst', 7:'learned*FAN-dst'} -ctm_dict = { 0: "None", 1:"rct", 2:"lct", 3:"mkl", 4:"mkl-m", 5:"idt", 6:"idt-m", 7:"sot-m" } -ctm_str_dict = {None:0, "rct":1, "lct":2, "mkl":3, "mkl-m":4, "idt":5, "idt-m":6, "sot-m":7 } +ctm_dict = { 0: "None", 1:"rct", 2:"lct", 3:"mkl", 4:"mkl-m", 5:"idt", 6:"idt-m", 7:"sot-m", 8:"mix-m" } +ctm_str_dict = {None:0, "rct":1, "lct":2, "mkl":3, "mkl-m":4, "idt":5, "idt-m":6, "sot-m":7, "mix-m":8 } class ConverterConfigMasked(ConverterConfig): diff --git a/imagelib/__init__.py b/imagelib/__init__.py index 43fad64..7a1ed9b 100644 --- a/imagelib/__init__.py +++ b/imagelib/__init__.py @@ -11,7 +11,7 @@ from .warp import gen_warp_params, warp_by_params from .reduce_colors import reduce_colors -from .color_transfer import color_transfer_sot, color_transfer_mkl, color_transfer_idt, color_hist_match, reinhard_color_transfer, linear_color_transfer, seamless_clone +from .color_transfer import color_transfer_mix, color_transfer_sot, color_transfer_mkl, color_transfer_idt, color_hist_match, reinhard_color_transfer, linear_color_transfer, seamless_clone from .RankSRGAN import RankSRGAN diff --git a/imagelib/color_transfer.py b/imagelib/color_transfer.py index 37f0af9..dfb65cb 100644 --- a/imagelib/color_transfer.py +++ b/imagelib/color_transfer.py @@ -54,7 +54,10 @@ def color_transfer_sot(src,trg, steps=10, batch_size=5, reg_sigmaXY=16.0, reg_si if reg_sigmaXY != 0.0: src_diff = new_src-src - new_src = src + cv2.bilateralFilter (src_diff, 0, reg_sigmaV, reg_sigmaXY ) + src_diff_filt = cv2.bilateralFilter (src_diff, 0, reg_sigmaV, reg_sigmaXY ) + if len(src_diff_filt.shape) == 2: + src_diff_filt = src_diff_filt[...,None] + new_src = src + src_diff_filt return new_src def color_transfer_mkl(x0, x1): @@ -269,11 +272,11 @@ def linear_color_transfer(target_img, source_img, mode='pca', eps=1e-5): ''' mu_t = target_img.mean(0).mean(0) t = target_img - mu_t - t = t.transpose(2,0,1).reshape(3,-1) + t = t.transpose(2,0,1).reshape( t.shape[-1],-1) Ct = t.dot(t.T) / t.shape[1] + eps * np.eye(t.shape[0]) mu_s = source_img.mean(0).mean(0) s = source_img - mu_s - s = s.transpose(2,0,1).reshape(3,-1) + s = s.transpose(2,0,1).reshape( s.shape[-1],-1) Cs = s.dot(s.T) / s.shape[1] + eps * np.eye(s.shape[0]) if mode == 'chol': chol_t = np.linalg.cholesky(Ct) @@ -364,3 +367,30 @@ def color_hist_match(src_im, tar_im, hist_match_threshold=255): matched = np.stack(to_stack, axis=-1).astype(src_im.dtype) return matched + +def color_transfer_mix(img_src,img_trg): + img_src = (img_src*255.0).astype(np.uint8) + img_trg = (img_trg*255.0).astype(np.uint8) + + img_src_lab = cv2.cvtColor(img_src, cv2.COLOR_BGR2LAB) + img_trg_lab = cv2.cvtColor(img_trg, cv2.COLOR_BGR2LAB) + + rct_light = np.clip ( linear_color_transfer(img_src_lab[...,0:1].astype(np.float32)/255.0, + img_trg_lab[...,0:1].astype(np.float32)/255.0 )[...,0]*255.0, + 0, 255).astype(np.uint8) + + img_src_lab[...,0] = (np.ones_like (rct_light)*100).astype(np.uint8) + img_src_lab = cv2.cvtColor(img_src_lab, cv2.COLOR_LAB2BGR) + + img_trg_lab[...,0] = (np.ones_like (rct_light)*100).astype(np.uint8) + img_trg_lab = cv2.cvtColor(img_trg_lab, cv2.COLOR_LAB2BGR) + + img_rct = color_transfer_sot( img_src_lab.astype(np.float32), img_trg_lab.astype(np.float32) ) + img_rct = np.clip(img_rct, 0, 255).astype(np.uint8) + + img_rct = cv2.cvtColor(img_rct, cv2.COLOR_BGR2LAB) + img_rct[...,0] = rct_light + img_rct = cv2.cvtColor(img_rct, cv2.COLOR_LAB2BGR) + + + return (img_rct / 255.0).astype(np.float32) \ No newline at end of file