update

app.py

@@ -17,13 +17,12 @@ import concurrent.futures
 from moviepy.editor import VideoFileClip
 
 from nsfw_detector import get_nsfw_detector
-from face_swapper import Inswapper, paste_to_whole
+from face_swapper import Inswapper, paste_to_whole, place_foreground_on_background
 from face_analyser import detect_conditions, get_analysed_data, swap_options_list
-from face_enhancer import load_face_enhancer_model, face_enhancer_list, gfpgan_enhance, realesrgan_enhance
+from face_enhancer import get_available_enhancer_names, load_face_enhancer_model
 from face_parsing import init_parser, swap_regions, mask_regions, mask_regions_to_list, SoftErosion
 from utils import trim_video, StreamerThread, ProcessBar, open_directory, split_list_by_lengths, merge_img_sequence_from_ref
 
-
 ## ------------------------------ USER ARGS ------------------------------
 
 parser = argparse.ArgumentParser(description="Swap-Mukham Face Swapper")
@@ -69,9 +68,12 @@ FACE_ANALYSER = None
 FACE_ENHANCER = None
 FACE_PARSER = None
 NSFW_DETECTOR = None
+FACE_ENHANCER_LIST = ["NONE"]
+FACE_ENHANCER_LIST.extend(get_available_enhancer_names())
+
 
 ## ------------------------------ SET EXECUTION PROVIDER ------------------------------
-# Note: For AMD,MAC or non CUDA users, change settings here
+# Note: Non CUDA users may change settings here
 
 PROVIDER = ["CPUExecutionProvider"]
 
@@ -88,7 +90,7 @@ else:
     print("\n********** Running on CPU **********\n")
 
 device = "cuda" if USE_CUDA else "cpu"
-
+EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == "cuda" else None
 
 ## ------------------------------ LOAD MODELS ------------------------------
 
@@ -223,7 +225,7 @@ def process(
             yield f"### \n 🔞 {message}", *ui_before()
             assert not nsfw, message
             return False
-        if device == "cuda": torch.cuda.empty_cache()
+        EMPTY_CACHE()
 
         yield "### \n ⌛ Analysing face data...", *ui_before()
         if condition != "Specific Face":
@@ -241,26 +243,24 @@ def process(
 
         yield "### \n ⌛ Swapping faces...", *ui_before()
         preds, aimgs, matrs = FACE_SWAPPER.batch_forward(whole_frame_list, analysed_targets, analysed_sources)
-        torch.cuda.empty_cache()
+        EMPTY_CACHE()
 
         if enable_face_parser:
             yield "### \n ⌛ Applying face-parsing mask...", *ui_before()
             for idx, (pred, aimg) in tqdm(enumerate(zip(preds, aimgs)), total=len(preds), desc="Face parsing"):
                 preds[idx] = swap_regions(pred, aimg, FACE_PARSER, smooth_mask, includes=includes, blur=int(blur_amount))
-            torch.cuda.empty_cache()
+        EMPTY_CACHE()
 
         if face_enhancer_name != "NONE":
             yield f"### \n ⌛ Enhancing faces with {face_enhancer_name}...", *ui_before()
             for idx, pred in tqdm(enumerate(preds), total=len(preds), desc=f"{face_enhancer_name}"):
-                if face_enhancer_name == 'GFPGAN':
-                    pred = gfpgan_enhance(pred, FACE_ENHANCER)
-                elif face_enhancer_name.startswith("REAL-ESRGAN"):
-                    pred = realesrgan_enhance(pred, FACE_ENHANCER)
-
+                enhancer_model, enhancer_model_runner = FACE_ENHANCER
+                pred = enhancer_model_runner(pred, enhancer_model)
                 preds[idx] = cv2.resize(pred, (512,512))
                 aimgs[idx] = cv2.resize(aimgs[idx], (512,512))
                 matrs[idx] /= 0.25
-        torch.cuda.empty_cache()
+
+        EMPTY_CACHE()
 
         split_preds = split_list_by_lengths(preds, num_faces_per_frame)
         del preds
@@ -270,19 +270,19 @@ def process(
         del matrs
 
         yield "### \n ⌛ Post-processing...", *ui_before()
-        def process_frame(frame_idx, frame_img, split_preds, split_aimgs, split_matrs, enable_laplacian_blend, crop_top, crop_bott, crop_left, crop_right):
+        def post_process(frame_idx, frame_img, split_preds, split_aimgs, split_matrs, enable_laplacian_blend, crop_top, crop_bott, crop_left, crop_right):
             whole_img_path = frame_img
             whole_img = cv2.imread(whole_img_path)
             for p, a, m in zip(split_preds[frame_idx], split_aimgs[frame_idx], split_matrs[frame_idx]):
                 whole_img = paste_to_whole(p, a, m, whole_img, laplacian_blend=enable_laplacian_blend, crop_mask=(crop_top, crop_bott, crop_left, crop_right))
             cv2.imwrite(whole_img_path, whole_img)
 
-        def optimize_processing(image_sequence, split_preds, split_aimgs, split_matrs, enable_laplacian_blend, crop_top, crop_bott, crop_left, crop_right):
+        def concurrent_post_process(image_sequence, split_preds, split_aimgs, split_matrs, enable_laplacian_blend, crop_top, crop_bott, crop_left, crop_right):
             with concurrent.futures.ThreadPoolExecutor() as executor:
                 futures = []
                 for idx, frame_img in enumerate(image_sequence):
                     future = executor.submit(
-                        process_frame,
+                        post_process,
                         idx,
                         frame_img,
                         split_preds,
@@ -302,8 +302,7 @@ def process(
                     except Exception as e:
                         print(f"An error occurred: {e}")
 
-        # Usage:
-        optimize_processing(
+        concurrent_post_process(
             image_sequence,
             split_preds,
             split_aimgs,
@@ -432,13 +431,13 @@ def update_radio(value):
 
 
 def swap_option_changed(value):
-    if value == swap_options_list[1] or value == swap_options_list[2]:
+    if value.startswith("Age"):
         return (
             gr.update(visible=True),
             gr.update(visible=False),
             gr.update(visible=True),
         )
-    elif value == swap_options_list[5]:
+    elif value == "Specific Face":
         return (
             gr.update(visible=False),
             gr.update(visible=True),
@@ -497,7 +496,7 @@ def stop_running():
     if hasattr(STREAMER, "stop"):
         STREAMER.stop()
         STREAMER = None
-    return "Cancelled"
+    yield "cancelled !"
 
 
 def slider_changed(show_frame, video_path, frame_index):
@@ -538,8 +537,10 @@ with gr.Blocks(css=css) as interface:
         with gr.Row():
             with gr.Column(scale=0.4):
                 with gr.Tab("📄 Swap Condition"):
-                    swap_option = gr.Radio(
+                    swap_option = gr.Dropdown(
                         swap_options_list,
+                        info="Choose which face or faces in the target image to swap.",
+                        multiselect=False,
                         show_label=False,
                         value=swap_options_list[0],
                         interactive=True,
@@ -636,7 +637,7 @@ with gr.Blocks(css=css) as interface:
                     )
 
                     face_enhancer_name = gr.Dropdown(
-                        face_enhancer_list, label="Face Enhancer", value="NONE", multiselect=False, interactive=True
+                        FACE_ENHANCER_LIST, label="Face Enhancer", value="NONE", multiselect=False, interactive=True
                     )
 
                 source_image_input = gr.Image(
@@ -675,8 +676,8 @@ with gr.Blocks(css=css) as interface:
                         )
 
                     with gr.Box(visible=True) as input_video_group:
-                        # vid_widget = gr.Video if USE_COLAB else gr.Text
-                        video_input = gr.Video(
+                        vid_widget = gr.Video if USE_COLAB else gr.Text
+                        video_input = vid_widget(
                             label="Target Video Path", interactive=True
                         )
                         with gr.Accordion("✂️ Trim video", open=False):
@@ -837,7 +838,7 @@ with gr.Blocks(css=css) as interface:
     ]
 
     swap_event = swap_button.click(
-        fn=process, inputs=swap_inputs, outputs=swap_outputs, show_progress=True
+        fn=process, inputs=swap_inputs, outputs=swap_outputs, show_progress=True,
     )
 
     cancel_button.click(

face_analyser.py

@@ -5,24 +5,58 @@ from tqdm import tqdm
 from utils import scale_bbox_from_center
 
 detect_conditions = [
+    "best detection",
     "left most",
     "right most",
     "top most",
     "bottom most",
-    "most width",
-    "most height",
-    "best detection",
+    "middle",
+    "biggest",
+    "smallest",
 ]
 
 swap_options_list = [
-    "All face",
+    "All Face",
+    "Specific Face",
     "Age less than",
     "Age greater than",
     "All Male",
     "All Female",
-    "Specific Face",
+    "Left Most",
+    "Right Most",
+    "Top Most",
+    "Bottom Most",
+    "Middle",
+    "Biggest",
+    "Smallest",
 ]
 
+def get_single_face(faces, method="best detection"):
+    total_faces = len(faces)
+    if total_faces == 1:
+        return faces[0]
+
+    print(f"{total_faces} face detected. Using {method} face.")
+    if method == "best detection":
+        return sorted(faces, key=lambda face: face["det_score"])[-1]
+    elif method == "left most":
+        return sorted(faces, key=lambda face: face["bbox"][0])[0]
+    elif method == "right most":
+        return sorted(faces, key=lambda face: face["bbox"][0])[-1]
+    elif method == "top most":
+        return sorted(faces, key=lambda face: face["bbox"][1])[0]
+    elif method == "bottom most":
+        return sorted(faces, key=lambda face: face["bbox"][1])[-1]
+    elif method == "middle":
+        return sorted(faces, key=lambda face: (
+                (face["bbox"][0] + face["bbox"][2]) / 2 - 0.5) ** 2 +
+                ((face["bbox"][1] + face["bbox"][3]) / 2 - 0.5) ** 2)[len(faces) // 2]
+    elif method == "biggest":
+        return sorted(faces, key=lambda face: (face["bbox"][2] - face["bbox"][0]) * (face["bbox"][3] - face["bbox"][1]))[-1]
+    elif method == "smallest":
+        return sorted(faces, key=lambda face: (face["bbox"][2] - face["bbox"][0]) * (face["bbox"][3] - face["bbox"][1]))[0]
+
+
 def analyse_face(image, model, return_single_face=True, detect_condition="best detection", scale=1.0):
     faces = model.get(image)
     if scale != 1: # landmark-scale
@@ -35,25 +69,7 @@ def analyse_face(image, model, return_single_face=True, detect_condition="best d
     if not return_single_face:
         return faces
 
-    total_faces = len(faces)
-    if total_faces == 1:
-        return faces[0]
-
-    print(f"{total_faces} face detected. Using {detect_condition} face.")
-    if detect_condition == "left most":
-        return sorted(faces, key=lambda face: face["bbox"][0])[0]
-    elif detect_condition == "right most":
-        return sorted(faces, key=lambda face: face["bbox"][0])[-1]
-    elif detect_condition == "top most":
-        return sorted(faces, key=lambda face: face["bbox"][1])[0]
-    elif detect_condition == "bottom most":
-        return sorted(faces, key=lambda face: face["bbox"][1])[-1]
-    elif detect_condition == "most width":
-        return sorted(faces, key=lambda face: face["bbox"][2])[-1]
-    elif detect_condition == "most height":
-        return sorted(faces, key=lambda face: face["bbox"][3])[-1]
-    elif detect_condition == "best detection":
-        return sorted(faces, key=lambda face: face["det_score"])[-1]
+    return get_single_face(faces, method=detect_condition)
 
 
 def cosine_distance(a, b):
@@ -90,7 +106,7 @@ def get_analysed_data(face_analyser, image_sequence, source_data, swap_condition
 
         n_faces = 0
         for analysed_face in analysed_faces:
-            if swap_condition == "All face":
+            if swap_condition == "All Face":
                 analysed_target_list.append(analysed_face)
                 analysed_source_list.append(analysed_source)
                 whole_frame_eql_list.append(frame_path)
@@ -124,6 +140,55 @@ def get_analysed_data(face_analyser, image_sequence, source_data, swap_condition
                         whole_frame_eql_list.append(frame_path)
                         n_faces += 1
 
+        if swap_condition == "Left Most":
+            analysed_face = get_single_face(analysed_faces, method="left most")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
+        elif swap_condition == "Right Most":
+            analysed_face = get_single_face(analysed_faces, method="right most")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
+        elif swap_condition == "Top Most":
+            analysed_face = get_single_face(analysed_faces, method="top most")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
+        elif swap_condition == "Bottom Most":
+            analysed_face = get_single_face(analysed_faces, method="bottom most")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
+        elif swap_condition == "Middle":
+            analysed_face = get_single_face(analysed_faces, method="middle")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
+        elif swap_condition == "Biggest":
+            analysed_face = get_single_face(analysed_faces, method="biggest")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
+        elif swap_condition == "Smallest":
+            analysed_face = get_single_face(analysed_faces, method="smallest")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
         num_faces_per_frame.append(n_faces)
 
     return analysed_target_list, analysed_source_list, whole_frame_eql_list, num_faces_per_frame

face_enhancer.py

@@ -4,36 +4,49 @@ import gfpgan
 from PIL import Image
 from upscaler.RealESRGAN import RealESRGAN
 
-face_enhancer_list = ['NONE', 'GFPGAN', 'REAL-ESRGAN 2x', 'REAL-ESRGAN 4x', 'REAL-ESRGAN 8x']
+
+def gfpgan_runner(img, model):
+    _, imgs, _ = model.enhance(img, paste_back=True, has_aligned=True)
+    return imgs[0]
+
+
+def realesrgan_runner(img, model):
+    img = model.predict(img)
+    return img
+
+
+supported_enhancers = {
+    "GFPGAN": ("./assets/pretrained_models/GFPGANv1.4.pth", gfpgan_runner),
+    "REAL-ESRGAN 2x": ("./assets/pretrained_models/RealESRGAN_x2.pth", realesrgan_runner),
+    "REAL-ESRGAN 4x": ("./assets/pretrained_models/RealESRGAN_x4.pth", realesrgan_runner),
+    "REAL-ESRGAN 8x": ("./assets/pretrained_models/RealESRGAN_x8.pth", realesrgan_runner)
+}
+
+
+def get_available_enhancer_names():
+    available = []
+    for name, data in supported_enhancers.items():
+        path = os.path.join(os.path.abspath(os.path.dirname(__file__)), data[0])
+        if os.path.exists(path):
+            available.append(name)
+    return available
+
 
 def load_face_enhancer_model(name='GFPGAN', device="cpu"):
+    assert name in get_available_enhancer_names(), f"Face enhancer {name} unavailable."
+    model_path, model_runner = supported_enhancers.get(name)
+    model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
     if name == 'GFPGAN':
-        model_path = "./assets/pretrained_models/GFPGANv1.4.pth"
-        model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
-        model = gfpgan.GFPGANer(model_path=model_path, upscale=1)
+        model = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=device)
     elif name == 'REAL-ESRGAN 2x':
-        model_path = "./assets/pretrained_models/RealESRGAN_x2.pth"
-        model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
         model = RealESRGAN(device, scale=2)
         model.load_weights(model_path, download=False)
     elif name == 'REAL-ESRGAN 4x':
-        model_path = "./assets/pretrained_models/RealESRGAN_x4.pth"
-        model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
         model = RealESRGAN(device, scale=4)
         model.load_weights(model_path, download=False)
     elif name == 'REAL-ESRGAN 8x':
-        model_path = "./assets/pretrained_models/RealESRGAN_x8.pth"
-        model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
         model = RealESRGAN(device, scale=8)
         model.load_weights(model_path, download=False)
     else:
         model = None
-    return model
-
-def gfpgan_enhance(img, model, has_aligned=True):
-    _, imgs, _ = model.enhance(img, paste_back=True, has_aligned=has_aligned)
-    return imgs[0]
-
-def realesrgan_enhance(img, model):
-    img = model.predict(img)
-    return img
+    return (model, model_runner)

face_swapper.py

@@ -7,6 +7,8 @@ import numpy as np
 from tqdm import tqdm
 from onnx import numpy_helper
 from skimage import transform as trans
+import torchvision.transforms.functional as F
+from utils import make_white_image, laplacian_blending
 
 arcface_dst = np.array(
     [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366],
@@ -62,33 +64,44 @@ class Inswapper():
         self.input_size = tuple(input_shape[2:4][::-1])
 
     def forward(self, imgs, latents):
-        batch_preds = []
+        preds = []
         for img, latent in zip(imgs, latents):
             img = (img - self.input_mean) / self.input_std
             pred = self.session.run(self.output_names, {self.input_names[0]: img, self.input_names[1]: latent})[0]
-            batch_preds.append(pred)
-        return batch_preds
+            preds.append(pred)
 
     def get(self, imgs, target_faces, source_faces):
-        batch_preds = []
-        batch_aimgs = []
-        batch_ms = []
-        for img, target_face, source_face in zip(imgs, target_faces, source_faces):
-            if isinstance(img, str):
-                img = cv2.imread(img)
-            aimg, M = norm_crop2(img, target_face.kps, self.input_size[0])
-            blob = cv2.dnn.blobFromImage(aimg, 1.0 / self.input_std, self.input_size,
-                                         (self.input_mean, self.input_mean, self.input_mean), swapRB=True)
-            latent = source_face.normed_embedding.reshape((1, -1))
-            latent = np.dot(latent, self.emap)
-            latent /= np.linalg.norm(latent)
+        imgs = list(imgs)
+
+        preds = [None] * len(imgs)
+        aimgs = [None] * len(imgs)
+        matrs = [None] * len(imgs)
+
+        for idx, (img, target_face, source_face) in enumerate(zip(imgs, target_faces, source_faces)):
+            aimg, M, blob, latent = self.prepare_data(img, target_face, source_face)
+            aimgs[idx] = aimg
+            matrs[idx] = M
             pred = self.session.run(self.output_names, {self.input_names[0]: blob, self.input_names[1]: latent})[0]
             pred = pred.transpose((0, 2, 3, 1))[0]
             pred = np.clip(255 * pred, 0, 255).astype(np.uint8)[:, :, ::-1]
-            batch_preds.append(pred)
-            batch_aimgs.append(aimg)
-            batch_ms.append(M)
-        return batch_preds, batch_aimgs, batch_ms
+            preds[idx] = pred
+
+        return (preds, aimgs, matrs)
+
+    def prepare_data(self, img, target_face, source_face):
+        if isinstance(img, str):
+            img = cv2.imread(img)
+
+        aimg, M = norm_crop2(img, target_face.kps, self.input_size[0])
+
+        blob = cv2.dnn.blobFromImage(aimg, 1.0 / self.input_std, self.input_size,
+                (self.input_mean, self.input_mean, self.input_mean), swapRB=True)
+
+        latent = source_face.normed_embedding.reshape((1, -1))
+        latent = np.dot(latent, self.emap)
+        latent /= np.linalg.norm(latent)
+
+        return (aimg, M, blob, latent)
 
     def batch_forward(self, img_list, target_f_list, source_f_list):
         num_samples = len(img_list)
@@ -96,8 +109,9 @@ class Inswapper():
 
         preds = []
         aimgs = []
-        ms = []
-        for i in tqdm(range(num_batches), desc="Swapping face by batch"):
+        matrs = []
+
+        for i in tqdm(range(num_batches), desc="Swapping face"):
             start_idx = i * self.batch_size
             end_idx = min((i + 1) * self.batch_size, num_samples)
 
@@ -105,86 +119,26 @@ class Inswapper():
             batch_target_f = target_f_list[start_idx:end_idx]
             batch_source_f = source_f_list[start_idx:end_idx]
 
-            batch_pred, batch_aimg, batch_m = self.get(batch_img, batch_target_f, batch_source_f)
+            batch_pred, batch_aimg, batch_matr = self.get(batch_img, batch_target_f, batch_source_f)
             preds.extend(batch_pred)
             aimgs.extend(batch_aimg)
-            ms.extend(batch_m)
-        return preds, aimgs, ms
-
-
-def laplacian_blending(A, B, m, num_levels=4):
-    assert A.shape == B.shape
-    assert B.shape == m.shape
-    height = m.shape[0]
-    width = m.shape[1]
-    size_list = np.array([4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096])
-    size = size_list[np.where(size_list > max(height, width))][0]
-    GA = np.zeros((size, size, 3), dtype=np.float32)
-    GA[:height, :width, :] = A
-    GB = np.zeros((size, size, 3), dtype=np.float32)
-    GB[:height, :width, :] = B
-    GM = np.zeros((size, size, 3), dtype=np.float32)
-    GM[:height, :width, :] = m
-    gpA = [GA]
-    gpB = [GB]
-    gpM = [GM]
-    for i in range(num_levels):
-        GA = cv2.pyrDown(GA)
-        GB = cv2.pyrDown(GB)
-        GM = cv2.pyrDown(GM)
-        gpA.append(np.float32(GA))
-        gpB.append(np.float32(GB))
-        gpM.append(np.float32(GM))
-    lpA  = [gpA[num_levels-1]]
-    lpB  = [gpB[num_levels-1]]
-    gpMr = [gpM[num_levels-1]]
-    for i in range(num_levels-1,0,-1):
-        LA = np.subtract(gpA[i-1], cv2.pyrUp(gpA[i]))
-        LB = np.subtract(gpB[i-1], cv2.pyrUp(gpB[i]))
-        lpA.append(LA)
-        lpB.append(LB)
-        gpMr.append(gpM[i-1])
-    LS = []
-    for la,lb,gm in zip(lpA,lpB,gpMr):
-        ls = la * gm + lb * (1.0 - gm)
-        LS.append(ls)
-    ls_ = LS[0]
-    for i in range(1,num_levels):
-        ls_ = cv2.pyrUp(ls_)
-        ls_ = cv2.add(ls_, LS[i])
-    ls_ = np.clip(ls_[:height, :width, :], 0, 255)
-    return ls_
+            matrs.extend(batch_matr)
+
+        return (preds, aimgs, matrs)
 
 
 def paste_to_whole(bgr_fake, aimg, M, whole_img, laplacian_blend=True, crop_mask=(0,0,0,0)):
     IM = cv2.invertAffineTransform(M)
 
-    img_white = np.full((aimg.shape[0], aimg.shape[1]), 255, dtype=np.float32)
-
-    top = int(crop_mask[0])
-    bottom = int(crop_mask[1])
-    if top + bottom < aimg.shape[1]:
-        if top > 0: img_white[:top, :] = 0
-        if bottom > 0: img_white[-bottom:, :] = 0
-
-    left = int(crop_mask[2])
-    right = int(crop_mask[3])
-    if left + right < aimg.shape[0]:
-        if left > 0: img_white[:, :left] = 0
-        if right > 0: img_white[:, -right:] = 0
-
-    bgr_fake = cv2.warpAffine(
-        bgr_fake, IM, (whole_img.shape[1], whole_img.shape[0]), borderValue=0.0
-    )
-    img_white = cv2.warpAffine(
-        img_white, IM, (whole_img.shape[1], whole_img.shape[0]), borderValue=0.0
-    )
+    img_white = make_white_image(aimg.shape[:2], crop=crop_mask, white_value=255)
+
+    bgr_fake = cv2.warpAffine(bgr_fake, IM, (whole_img.shape[1], whole_img.shape[0]), borderValue=0.0)
+    img_white = cv2.warpAffine(img_white, IM, (whole_img.shape[1], whole_img.shape[0]), borderValue=0.0)
+
     img_white[img_white > 20] = 255
     img_mask = img_white
     mask_h_inds, mask_w_inds = np.where(img_mask == 255)
-    mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
-    mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
-    mask_size = int(np.sqrt(mask_h * mask_w))
+    mask_size = int(np.sqrt(np.ptp(mask_h_inds) * np.ptp(mask_w_inds)))
 
     k = max(mask_size // 10, 10)
     img_mask = cv2.erode(img_mask, np.ones((k, k), np.uint8), iterations=1)
@@ -201,3 +155,11 @@ def paste_to_whole(bgr_fake, aimg, M, whole_img, laplacian_blend=True, crop_mask
 
     fake_merged = img_mask * bgr_fake + (1 - img_mask) * whole_img.astype(np.float32)
     return fake_merged.astype("uint8")
+
+def place_foreground_on_background(foreground, background, matrix):
+    matrix = cv2.invertAffineTransform(matrix)
+    mask = np.ones(foreground.shape, dtype="float32")
+    foreground = cv2.warpAffine(foreground, matrix, (background.shape[1], background.shape[0]), borderValue=0.0)
+    mask = cv2.warpAffine(mask, matrix, (background.shape[1], background.shape[0]), borderValue=0.0)
+    composite_image = mask * foreground + (1 - mask) * background
+    return composite_image

nsfw_detector.py

@@ -7,6 +7,7 @@ import torch
 import timm
 from tqdm import tqdm
 
+# https://github.com/Whiax/NSFW-Classifier/raw/main/nsfwmodel_281.pth
 normalize_t = Normalize((0.4814, 0.4578, 0.4082), (0.2686, 0.2613, 0.2757))
 
 #nsfw classifier
@@ -28,7 +29,7 @@ class NSFWClassifier(nn.Module):
         x = nsfw_model.linear_probe(x)
         return x
 
-    def is_nsfw(self, img_paths, threshold = 0.93):
+    def is_nsfw(self, img_paths, threshold = 0.98):
         skip_step = 1
         total_len = len(img_paths)
         if total_len < 100: skip_step = 1
@@ -37,16 +38,19 @@ class NSFWClassifier(nn.Module):
         if total_len > 1000 and total_len < 10000: skip_step = 50
         if total_len > 10000: skip_step = 100
 
-        for idx in tqdm(range(0, total_len, skip_step), total=total_len, desc="Checking for NSFW contents"):
-            img = Image.open(img_paths[idx]).convert('RGB')
-            img = img.resize((224, 224))
+        for idx in tqdm(range(0, total_len, skip_step), total=int(total_len // skip_step), desc="Checking for NSFW contents"):
+            _img = Image.open(img_paths[idx]).convert('RGB')
+            img = _img.resize((224, 224))
             img = np.array(img)/255
             img = T.ToTensor()(img).unsqueeze(0).float()
             if next(self.parameters()).is_cuda:
                 img = img.cuda()
             with torch.no_grad():
                 score = self.forward(img).sigmoid()[0].item()
-            if score > threshold:return True
+            if score > threshold:
+                print(f"Detected nsfw score:{score}")
+                _img.save("nsfw.jpg")
+                return True
         return False
 
 def get_nsfw_detector(model_path='nsfwmodel_281.pth', device="cpu"):

swapper.py

@@ -1,106 +0,0 @@
-import cv2
-import numpy as np
-from insightface.utils import face_align
-from face_parsing.swap import swap_regions
-from utils import add_logo_to_image
-
-swap_options_list = [
-    "All face",
-    "Age less than",
-    "Age greater than",
-    "All Male",
-    "All Female",
-    "Specific Face",
-]
-
-
-def swap_face(whole_img, target_face, source_face, models):
-    inswapper = models.get("swap")
-    face_enhancer = models.get("enhance", None)
-    face_parser = models.get("face_parser", None)
-    fe_enable = models.get("enhance_sett", False)
-
-    bgr_fake, M = inswapper.get(whole_img, target_face, source_face, paste_back=False)
-    image_size = 128 if not fe_enable else 512
-    aimg, _ = face_align.norm_crop2(whole_img, target_face.kps, image_size=image_size)
-
-    if face_parser is not None:
-        fp_enable, includes, smooth_mask, blur_amount = models.get("face_parser_sett")
-        if fp_enable:
-            bgr_fake = swap_regions(
-                bgr_fake, aimg, face_parser, smooth_mask, includes=includes, blur=blur_amount
-            )
-
-    if fe_enable:
-        _, bgr_fake, _ = face_enhancer.enhance(
-            bgr_fake, paste_back=True, has_aligned=True
-        )
-        bgr_fake = bgr_fake[0]
-        M /= 0.25
-
-    IM = cv2.invertAffineTransform(M)
-
-    img_white = np.full((aimg.shape[0], aimg.shape[1]), 255, dtype=np.float32)
-    bgr_fake = cv2.warpAffine(
-        bgr_fake, IM, (whole_img.shape[1], whole_img.shape[0]), borderValue=0.0
-    )
-    img_white = cv2.warpAffine(
-        img_white, IM, (whole_img.shape[1], whole_img.shape[0]), borderValue=0.0
-    )
-    img_white[img_white > 20] = 255
-    img_mask = img_white
-    mask_h_inds, mask_w_inds = np.where(img_mask == 255)
-    mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
-    mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
-    mask_size = int(np.sqrt(mask_h * mask_w))
-
-    k = max(mask_size // 10, 10)
-    img_mask = cv2.erode(img_mask, np.ones((k, k), np.uint8), iterations=1)
-
-    k = max(mask_size // 20, 5)
-    kernel_size = (k, k)
-    blur_size = tuple(2 * i + 1 for i in kernel_size)
-    img_mask = cv2.GaussianBlur(img_mask, blur_size, 0) / 255
-
-    img_mask = np.reshape(img_mask, [img_mask.shape[0], img_mask.shape[1], 1])
-    fake_merged = img_mask * bgr_fake + (1 - img_mask) * whole_img.astype(np.float32)
-    fake_merged = add_logo_to_image(fake_merged.astype("uint8"))
-    return fake_merged
-
-
-def swap_face_with_condition(
-    whole_img, target_faces, source_face, condition, age, models
-):
-    swapped = whole_img.copy()
-
-    for target_face in target_faces:
-        if condition == "All face":
-            swapped = swap_face(swapped, target_face, source_face, models)
-        elif condition == "Age less than" and target_face["age"] < age:
-            swapped = swap_face(swapped, target_face, source_face, models)
-        elif condition == "Age greater than" and target_face["age"] > age:
-            swapped = swap_face(swapped, target_face, source_face, models)
-        elif condition == "All Male" and target_face["gender"] == 1:
-            swapped = swap_face(swapped, target_face, source_face, models)
-        elif condition == "All Female" and target_face["gender"] == 0:
-            swapped = swap_face(swapped, target_face, source_face, models)
-
-    return swapped
-
-
-def swap_specific(source_specifics, target_faces, whole_img, models, threshold=0.6):
-    swapped = whole_img.copy()
-
-    for source_face, specific_face in source_specifics:
-        specific_embed = specific_face["embedding"]
-        specific_embed /= np.linalg.norm(specific_embed)
-
-        for target_face in target_faces:
-            target_embed = target_face["embedding"]
-            target_embed /= np.linalg.norm(target_embed)
-            cosine_distance = 1 - np.dot(specific_embed, target_embed)
-            if cosine_distance > threshold:
-                continue
-            swapped = swap_face(swapped, target_face, source_face, models)
-
-    return swapped

utils.py

@@ -2,13 +2,64 @@ import os
 import cv2
 import time
 import glob
+import torch
 import shutil
 import platform
 import datetime
 import subprocess
+import numpy as np
 from threading import Thread
 from moviepy.editor import VideoFileClip, ImageSequenceClip
 from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip
+from face_parsing import init_parser, swap_regions, mask_regions, mask_regions_to_list, SoftErosion
+
+
+logo_image = cv2.imread("./assets/images/logo.png", cv2.IMREAD_UNCHANGED)
+
+
+quality_types = ["poor", "low", "medium", "high", "best"]
+
+
+bitrate_quality_by_resolution = {
+    240: {"poor": "300k", "low": "500k", "medium": "800k", "high": "1000k", "best": "1200k"},
+    360: {"poor": "500k","low": "800k","medium": "1200k","high": "1500k","best": "2000k"},
+    480: {"poor": "800k","low": "1200k","medium": "2000k","high": "2500k","best": "3000k"},
+    720: {"poor": "1500k","low": "2500k","medium": "4000k","high": "5000k","best": "6000k"},
+    1080: {"poor": "2500k","low": "4000k","medium": "6000k","high": "7000k","best": "8000k"},
+    1440: {"poor": "4000k","low": "6000k","medium": "8000k","high": "10000k","best": "12000k"},
+    2160: {"poor": "8000k","low": "10000k","medium": "12000k","high": "15000k","best": "20000k"}
+}
+
+
+crf_quality_by_resolution = {
+    240: {"poor": 45, "low": 35, "medium": 28, "high": 23, "best": 20},
+    360: {"poor": 35, "low": 28, "medium": 23, "high": 20, "best": 18},
+    480: {"poor": 28, "low": 23, "medium": 20, "high": 18, "best": 16},
+    720: {"poor": 23, "low": 20, "medium": 18, "high": 16, "best": 14},
+    1080: {"poor": 20, "low": 18, "medium": 16, "high": 14, "best": 12},
+    1440: {"poor": 18, "low": 16, "medium": 14, "high": 12, "best": 10},
+    2160: {"poor": 16, "low": 14, "medium": 12, "high": 10, "best": 8}
+}
+
+
+def get_bitrate_for_resolution(resolution, quality):
+    available_resolutions = list(bitrate_quality_by_resolution.keys())
+    closest_resolution = min(available_resolutions, key=lambda x: abs(x - resolution))
+    return bitrate_quality_by_resolution[closest_resolution][quality]
+
+
+def get_crf_for_resolution(resolution, quality):
+    available_resolutions = list(crf_quality_by_resolution.keys())
+    closest_resolution = min(available_resolutions, key=lambda x: abs(x - resolution))
+    return crf_quality_by_resolution[closest_resolution][quality]
+
+
+def get_video_bitrate(video_file):
+    ffprobe_cmd = ['ffprobe', '-v', 'error', '-select_streams', 'v:0', '-show_entries',
+        'stream=bit_rate', '-of', 'default=noprint_wrappers=1:nokey=1', video_file]
+    result = subprocess.run(ffprobe_cmd, stdout=subprocess.PIPE)
+    kbps = max(int(result.stdout) // 1000, 10)
+    return str(kbps) + 'k'
 
 
 def trim_video(video_path, output_path, start_frame, stop_frame):
@@ -23,9 +74,11 @@ def trim_video(video_path, output_path, start_frame, stop_frame):
     start_time = start_frame / fps
     duration = (stop_frame - start_frame) / fps
 
+    bitrate = get_bitrate_for_resolution(min(*video.size), "high")
+
     trimmed_video = video.subclip(start_time, start_time + duration)
     trimmed_video.write_videofile(
-        trimmed_video_file_path, codec="libx264", audio_codec="aac"
+        trimmed_video_file_path, codec="libx264", audio_codec="aac", bitrate=bitrate,
     )
     trimmed_video.close()
     video.close()
@@ -91,9 +144,6 @@ class ProcessBar:
         return info_text
 
 
-logo_image = cv2.imread("./assets/images/logo.png", cv2.IMREAD_UNCHANGED)
-
-
 def add_logo_to_image(img, logo=logo_image):
     logo_size = int(img.shape[1] * 0.1)
     logo = cv2.resize(logo, (logo_size, logo_size))
@@ -111,6 +161,7 @@ def add_logo_to_image(img, logo=logo_image):
         ]
     return img
 
+
 def split_list_by_lengths(data, length_list):
     split_data = []
     start_idx = 0
@@ -121,6 +172,7 @@ def split_list_by_lengths(data, length_list):
         start_idx = end_idx
     return split_data
 
+
 def merge_img_sequence_from_ref(ref_video_path, image_sequence, output_file_name):
     video_clip = VideoFileClip(ref_video_path)
     fps = video_clip.fps
@@ -132,12 +184,15 @@ def merge_img_sequence_from_ref(ref_video_path, image_sequence, output_file_name
     if audio_clip is not None:
         edited_video_clip = edited_video_clip.set_audio(audio_clip)
 
+    bitrate = get_bitrate_for_resolution(min(*edited_video_clip.size), "high")
+
     edited_video_clip.set_duration(duration).write_videofile(
-        output_file_name, codec="libx264"
+        output_file_name, codec="libx264", bitrate=bitrate,
     )
     edited_video_clip.close()
     video_clip.close()
 
+
 def scale_bbox_from_center(bbox, scale_width, scale_height, image_width, image_height):
     # Extract the coordinates of the bbox
     x1, y1, x2, y2 = bbox
@@ -167,3 +222,72 @@ def scale_bbox_from_center(bbox, scale_width, scale_height, image_width, image_h
     # Return the scaled bbox coordinates
     scaled_bbox = [new_x1, new_y1, new_x2, new_y2]
     return scaled_bbox
+
+
+def laplacian_blending(A, B, m, num_levels=4):
+    assert A.shape == B.shape
+    assert B.shape == m.shape
+    height = m.shape[0]
+    width = m.shape[1]
+    size_list = np.array([4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096])
+    size = size_list[np.where(size_list > max(height, width))][0]
+    GA = np.zeros((size, size, 3), dtype=np.float32)
+    GA[:height, :width, :] = A
+    GB = np.zeros((size, size, 3), dtype=np.float32)
+    GB[:height, :width, :] = B
+    GM = np.zeros((size, size, 3), dtype=np.float32)
+    GM[:height, :width, :] = m
+    gpA = [GA]
+    gpB = [GB]
+    gpM = [GM]
+    for i in range(num_levels):
+        GA = cv2.pyrDown(GA)
+        GB = cv2.pyrDown(GB)
+        GM = cv2.pyrDown(GM)
+        gpA.append(np.float32(GA))
+        gpB.append(np.float32(GB))
+        gpM.append(np.float32(GM))
+    lpA  = [gpA[num_levels-1]]
+    lpB  = [gpB[num_levels-1]]
+    gpMr = [gpM[num_levels-1]]
+    for i in range(num_levels-1,0,-1):
+        LA = np.subtract(gpA[i-1], cv2.pyrUp(gpA[i]))
+        LB = np.subtract(gpB[i-1], cv2.pyrUp(gpB[i]))
+        lpA.append(LA)
+        lpB.append(LB)
+        gpMr.append(gpM[i-1])
+    LS = []
+    for la,lb,gm in zip(lpA,lpB,gpMr):
+        ls = la * gm + lb * (1.0 - gm)
+        LS.append(ls)
+    ls_ = LS[0]
+    for i in range(1,num_levels):
+        ls_ = cv2.pyrUp(ls_)
+        ls_ = cv2.add(ls_, LS[i])
+    ls_ = np.clip(ls_[:height, :width, :], 0, 255)
+    return ls_
+
+
+def make_white_image(shape, crop=None, white_value=255):
+    img_white = np.full((shape[0], shape[1]), white_value, dtype=np.float32)
+    if crop is not None:
+        top = int(crop[0])
+        bottom = int(crop[1])
+        if top + bottom < shape[1]:
+            if top > 0: img_white[:top, :] = 0
+            if bottom > 0: img_white[-bottom:, :] = 0
+
+        left = int(crop[2])
+        right = int(crop[3])
+        if left + right < shape[0]:
+            if left > 0: img_white[:, :left] = 0
+            if right > 0: img_white[:, -right:] = 0
+
+    return img_white
+
+
+def remove_hair(img, model=None):
+    if model is None:
+        path = "./assets/pretrained_models/79999_iter.pth"
+        model = init_parser(path, mode="cuda" if torch.cuda.is_available() else "cpu")
+