update

app.py

@@ -4,7 +4,6 @@ import glob
 import time
 import torch
 import shutil
-import gfpgan
 import argparse
 import platform
 import datetime
@@ -13,22 +12,22 @@ import insightface
 import onnxruntime
 import numpy as np
 import gradio as gr
-from moviepy.editor import VideoFileClip, ImageSequenceClip
+from tqdm import tqdm
+from moviepy.editor import VideoFileClip
 
-from face_analyser import detect_conditions, analyse_face
-from utils import trim_video, StreamerThread, ProcessBar, open_directory
+from nsfw_detector import get_nsfw_detector
+from face_swapper import Inswapper, paste_to_whole
+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_parsing import init_parser, swap_regions, mask_regions, mask_regions_to_list, SoftErosion
-from swapper import (
-    swap_face,
-    swap_face_with_condition,
-    swap_specific,
-    swap_options_list,
-)
+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")
 parser.add_argument("--out_dir", help="Default Output directory", default=os.getcwd())
+parser.add_argument("--batch_size", help="Gpu batch size", default=32)
 parser.add_argument("--cuda", action="store_true", help="Enable cuda", default=False)
 parser.add_argument(
     "--colab", action="store_true", help="Enable colab mode", default=False
@@ -40,11 +39,12 @@ user_args = parser.parse_args()
 USE_COLAB = user_args.colab
 USE_CUDA = user_args.cuda
 DEF_OUTPUT_PATH = user_args.out_dir
+BATCH_SIZE = user_args.batch_size
 WORKSPACE = None
 OUTPUT_FILE = None
 CURRENT_FRAME = None
 STREAMER = None
-DETECT_CONDITION = "left most"
+DETECT_CONDITION = "best detection"
 DETECT_SIZE = 640
 DETECT_THRESH = 0.6
 NUM_OF_SRC_SPECIFIC = 10
@@ -67,6 +67,7 @@ FACE_SWAPPER = None
 FACE_ANALYSER = None
 FACE_ENHANCER = None
 FACE_PARSER = None
+NSFW_DETECTOR = None
 
 ## ------------------------------ SET EXECUTION PROVIDER ------------------------------
 # Note: For AMD,MAC or non CUDA users, change settings here
@@ -99,25 +100,22 @@ def load_face_analyser_model(name="buffalo_l"):
         )
 
 
-def load_face_swapper_model(name="./assets/pretrained_models/inswapper_128.onnx"):
+def load_face_swapper_model(path="./assets/pretrained_models/inswapper_128.onnx"):
     global FACE_SWAPPER
-    path = os.path.join(os.path.abspath(os.path.dirname(__file__)), name)
     if FACE_SWAPPER is None:
-        FACE_SWAPPER = insightface.model_zoo.get_model(path, providers=PROVIDER)
-
-
-def load_face_enhancer_model(name="./assets/pretrained_models/GFPGANv1.4.pth"):
-    global FACE_ENHANCER
-    path = os.path.join(os.path.abspath(os.path.dirname(__file__)), name)
-    if FACE_ENHANCER is None:
-        FACE_ENHANCER = gfpgan.GFPGANer(model_path=path, upscale=1)
+        batch = int(BATCH_SIZE) if device == "cuda" else 1
+        FACE_SWAPPER = Inswapper(model_file=path, batch_size=batch, providers=PROVIDER)
 
 
-def load_face_parser_model(name="./assets/pretrained_models/79999_iter.pth"):
+def load_face_parser_model(path="./assets/pretrained_models/79999_iter.pth"):
     global FACE_PARSER
-    path = os.path.join(os.path.abspath(os.path.dirname(__file__)), name)
     if FACE_PARSER is None:
-        FACE_PARSER = init_parser(name, mode=device)
+        FACE_PARSER = init_parser(path, mode=device)
+
+def load_nsfw_detector_model(path="./assets/pretrained_models/nsfwmodel_281.pth"):
+    global NSFW_DETECTOR
+    if NSFW_DETECTOR is None:
+        NSFW_DETECTOR = get_nsfw_detector(model_path=path, device=device)
 
 
 load_face_analyser_model()
@@ -138,12 +136,18 @@ def process(
     condition,
     age,
     distance,
-    face_enhance,
+    face_enhancer_name,
     enable_face_parser,
     mask_includes,
     mask_soft_kernel,
     mask_soft_iterations,
     blur_amount,
+    face_scale,
+    enable_laplacian_blend,
+    crop_top,
+    crop_bott,
+    crop_left,
+    crop_right,
     *specifics,
 ):
     global WORKSPACE
@@ -177,12 +181,13 @@ def process(
             gr.update(value=OUTPUT_FILE, visible=True),
         )
 
-    ## ------------------------------ LOAD PENDING MODELS ------------------------------
     start_time = time.time()
-    specifics = list(specifics)
-    half = len(specifics) // 2
-    sources = specifics[:half]
-    specifics = specifics[half:]
+    total_exec_time = lambda start_time: divmod(time.time() - start_time, 60)
+    get_finsh_text = lambda start_time: f"✔️ Completed in {int(total_exec_time(start_time)[0])} min {int(total_exec_time(start_time)[1])} sec."
+
+    ## ------------------------------ PREPARE INPUTS & LOAD MODELS ------------------------------
+    yield "### \n ⌛ Loading NSFW detector model...", *ui_before()
+    load_nsfw_detector_model()
 
     yield "### \n ⌛ Loading face analyser model...", *ui_before()
     load_face_analyser_model()
@@ -190,87 +195,100 @@ def process(
     yield "### \n ⌛ Loading face swapper model...", *ui_before()
     load_face_swapper_model()
 
-    if face_enhance:
-        yield "### \n ⌛ Loading face enhancer model...", *ui_before()
-        load_face_enhancer_model()
+    if face_enhancer_name != "NONE":
+        yield f"### \n ⌛ Loading {face_enhancer_name} model...", *ui_before()
+        FACE_ENHANCER = load_face_enhancer_model(name=face_enhancer_name, device=device)
+    else:
+        FACE_ENHANCER = None
 
     if enable_face_parser:
         yield "### \n ⌛ Loading face parsing model...", *ui_before()
         load_face_parser_model()
 
-    yield "### \n ⌛ Analysing Face...", *ui_before()
-
     includes = mask_regions_to_list(mask_includes)
-    if mask_soft_iterations > 0:
-        smooth_mask = SoftErosion(kernel_size=17, threshold=0.9, iterations=int(mask_soft_iterations)).to(device)
-    else:
-        smooth_mask = None
-
-    models = {
-        "swap": FACE_SWAPPER,
-        "enhance": FACE_ENHANCER,
-        "enhance_sett": face_enhance,
-        "face_parser": FACE_PARSER,
-        "face_parser_sett": (enable_face_parser, includes, smooth_mask, int(blur_amount))
-    }
-
-    ## ------------------------------ ANALYSE SOURCE & SPECIFIC ------------------------------
-
-    analysed_source_specific = []
-    if condition == "Specific Face":
-        for source, specific in zip(sources, specifics):
-            if source is None or specific is None:
-                continue
-            analysed_source = analyse_face(
-                source,
-                FACE_ANALYSER,
-                return_single_face=True,
-                detect_condition=DETECT_CONDITION,
-            )
-            analysed_specific = analyse_face(
-                specific,
-                FACE_ANALYSER,
-                return_single_face=True,
-                detect_condition=DETECT_CONDITION,
-            )
-            analysed_source_specific.append([analysed_source, analysed_specific])
-    else:
-        source = cv2.imread(source_path)
-        analysed_source = analyse_face(
-            source,
+    smooth_mask = SoftErosion(kernel_size=17, threshold=0.9, iterations=int(mask_soft_iterations)).to(device) if mask_soft_iterations > 0 else None
+    specifics = list(specifics)
+    half = len(specifics) // 2
+    sources = specifics[:half]
+    specifics = specifics[half:]
+
+    ## ------------------------------ ANALYSE & SWAP FUNC ------------------------------
+
+    def swap_process(image_sequence):
+        yield "### \n ⌛ Checking contents...", *ui_before()
+        nsfw = NSFW_DETECTOR.is_nsfw(image_sequence)
+        if nsfw:
+            message = "NSFW Content detected !!!"
+            yield f"### \n 🔞 {message}", *ui_before()
+            assert not nsfw, message
+            return False
+        if device == "cuda": torch.cuda.empty_cache()
+
+        yield "### \n ⌛ Analysing face data...", *ui_before()
+        if condition != "Specific Face":
+            source_data = source_path, age
+        else:
+            source_data = ((sources, specifics), distance)
+        analysed_targets, analysed_sources, whole_frame_list, num_faces_per_frame = get_analysed_data(
             FACE_ANALYSER,
-            return_single_face=True,
+            image_sequence,
+            source_data,
+            swap_condition=condition,
             detect_condition=DETECT_CONDITION,
+            scale=face_scale
         )
 
+        yield "### \n ⌛ Swapping faces...", *ui_before()
+        preds, aimgs, matrs = FACE_SWAPPER.batch_forward(whole_frame_list, analysed_targets, analysed_sources)
+        torch.cuda.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()
+
+        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)
+
+                preds[idx] = cv2.resize(pred, (512,512))
+                aimgs[idx] = cv2.resize(aimgs[idx], (512,512))
+                matrs[idx] /= 0.25
+        torch.cuda.empty_cache()
+
+        split_preds = split_list_by_lengths(preds, num_faces_per_frame)
+        split_aimgs = split_list_by_lengths(aimgs, num_faces_per_frame)
+        split_matrs = split_list_by_lengths(matrs, num_faces_per_frame)
+
+        yield "### \n ⌛ Post-processing...", *ui_before()
+        for idx, frame_img in tqdm(enumerate(image_sequence), total=len(image_sequence), desc="Post-Processing"):
+            whole_img_path = frame_img
+            whole_img = cv2.imread(whole_img_path)
+            for p, a, m in zip(split_preds[idx], split_aimgs[idx], split_matrs[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)
+
+
     ## ------------------------------ IMAGE ------------------------------
 
     if input_type == "Image":
         target = cv2.imread(image_path)
-        analysed_target = analyse_face(target, FACE_ANALYSER, return_single_face=False)
-        if condition == "Specific Face":
-            swapped = swap_specific(
-                analysed_source_specific,
-                analysed_target,
-                target,
-                models,
-                threshold=distance,
-            )
-        else:
-            swapped = swap_face_with_condition(
-                target, analysed_target, analysed_source, condition, age, models
-            )
+        output_file = os.path.join(output_path, output_name + ".png")
+        cv2.imwrite(output_file, target)
 
-        filename = os.path.join(output_path, output_name + ".png")
-        cv2.imwrite(filename, swapped)
-        OUTPUT_FILE = filename
-        WORKSPACE = output_path
-        PREVIEW = swapped[:, :, ::-1]
+        for info_update in swap_process([output_file]):
+            yield info_update
 
-        tot_exec_time = time.time() - start_time
-        _min, _sec = divmod(tot_exec_time, 60)
+        OUTPUT_FILE = output_file
+        WORKSPACE = output_path
+        PREVIEW = cv2.imread(output_file)[:, :, ::-1]
 
-        yield f"Completed in {int(_min)} min {int(_sec)} sec.", *ui_after()
+        yield get_finsh_text(start_time), *ui_after()
 
     ## ------------------------------ VIDEO ------------------------------
 
@@ -278,72 +296,26 @@ def process(
         temp_path = os.path.join(output_path, output_name, "sequence")
         os.makedirs(temp_path, exist_ok=True)
 
-        video_clip = VideoFileClip(video_path)
-        duration = video_clip.duration
-        fps = video_clip.fps
-        total_frames = video_clip.reader.nframes
-
-        analysed_targets = []
-        process_bar = ProcessBar(30, total_frames)
-        yield "### \n ⌛ Analysing...", *ui_before()
-        for i, frame in enumerate(video_clip.iter_frames()):
-            analysed_targets.append(
-                analyse_face(frame, FACE_ANALYSER, return_single_face=False)
-            )
-            info_text = "Analysing Faces || "
-            info_text += process_bar.get(i)
-            print("\033[1A\033[K", end="", flush=True)
-            print(info_text)
-            if i % 10 == 0:
-                yield "### \n" + info_text, *ui_before()
-        video_clip.close()
-
+        yield "### \n ⌛ Extracting video frames...", *ui_before()
         image_sequence = []
-        video_clip = VideoFileClip(video_path)
-        audio_clip = video_clip.audio if video_clip.audio is not None else None
-        process_bar = ProcessBar(30, total_frames)
-        yield "### \n ⌛ Swapping...", *ui_before()
-        for i, frame in enumerate(video_clip.iter_frames()):
-            swapped = frame
-            analysed_target = analysed_targets[i]
-
-            if condition == "Specific Face":
-                swapped = swap_specific(
-                    analysed_source_specific,
-                    analysed_target,
-                    frame,
-                    models,
-                    threshold=distance,
-                )
-            else:
-                swapped = swap_face_with_condition(
-                    frame, analysed_target, analysed_source, condition, age, models
-                )
-
-            image_path = os.path.join(temp_path, f"frame_{i}.png")
-            cv2.imwrite(image_path, swapped[:, :, ::-1])
-            image_sequence.append(image_path)
-
-            info_text = "Swapping Faces || "
-            info_text += process_bar.get(i)
-            print("\033[1A\033[K", end="", flush=True)
-            print(info_text)
-            if i % 6 == 0:
-                PREVIEW = swapped
-                yield "### \n" + info_text, *ui_before()
-
-        yield "### \n ⌛ Merging...", *ui_before()
-        edited_video_clip = ImageSequenceClip(image_sequence, fps=fps)
-
-        if audio_clip is not None:
-            edited_video_clip = edited_video_clip.set_audio(audio_clip)
-
+        cap = cv2.VideoCapture(video_path)
+        curr_idx = 0
+        while True:
+            ret, frame = cap.read()
+            if not ret:break
+            frame_path = os.path.join(temp_path, f"frame_{curr_idx}.jpg")
+            cv2.imwrite(frame_path, frame)
+            image_sequence.append(frame_path)
+            curr_idx += 1
+        cap.release()
+        cv2.destroyAllWindows()
+
+        for info_update in swap_process(image_sequence):
+            yield info_update
+
+        yield "### \n ⌛ Merging sequence...", *ui_before()
         output_video_path = os.path.join(output_path, output_name + ".mp4")
-        edited_video_clip.set_duration(duration).write_videofile(
-            output_video_path, codec="libx264"
-        )
-        edited_video_clip.close()
-        video_clip.close()
+        merge_img_sequence_from_ref(video_path, image_sequence, output_video_path)
 
         if os.path.exists(temp_path) and not keep_output_sequence:
             yield "### \n ⌛ Removing temporary files...", *ui_before()
@@ -352,99 +324,38 @@ def process(
         WORKSPACE = output_path
         OUTPUT_FILE = output_video_path
 
-        tot_exec_time = time.time() - start_time
-        _min, _sec = divmod(tot_exec_time, 60)
-
-        yield f"✔️ Completed in {int(_min)} min {int(_sec)} sec.", *ui_after_vid()
+        yield get_finsh_text(start_time), *ui_after_vid()
 
     ## ------------------------------ DIRECTORY ------------------------------
 
     elif input_type == "Directory":
-        source = cv2.imread(source_path)
-        source = analyse_face(
-            source,
-            FACE_ANALYSER,
-            return_single_face=True,
-            detect_condition=DETECT_CONDITION,
-        )
         extensions = ["jpg", "jpeg", "png", "bmp", "tiff", "ico", "webp"]
         temp_path = os.path.join(output_path, output_name)
         if os.path.exists(temp_path):
             shutil.rmtree(temp_path)
         os.mkdir(temp_path)
-        swapped = None
 
-        files = []
+        file_paths =[]
         for file_path in glob.glob(os.path.join(directory_path, "*")):
             if any(file_path.lower().endswith(ext) for ext in extensions):
-                files.append(file_path)
-
-        files_length = len(files)
-        filename = None
-        for i, file_path in enumerate(files):
-            target = cv2.imread(file_path)
-            analysed_target = analyse_face(
-                target, FACE_ANALYSER, return_single_face=False
-            )
-
-            if condition == "Specific Face":
-                swapped = swap_specific(
-                    analysed_source_specific,
-                    analysed_target,
-                    target,
-                    models,
-                    threshold=distance,
-                )
-            else:
-                swapped = swap_face_with_condition(
-                    target, analysed_target, analysed_source, condition, age, models
-                )
+                img = cv2.imread(file_path)
+                new_file_path = os.path.join(temp_path, os.path.basename(file_path))
+                cv2.imwrite(new_file_path, img)
+                file_paths.append(new_file_path)
 
-            filename = os.path.join(temp_path, os.path.basename(file_path))
-            cv2.imwrite(filename, swapped)
-            info_text = f"### \n ⌛ Processing file {i+1} of {files_length}"
-            PREVIEW = swapped[:, :, ::-1]
-            yield info_text, *ui_before()
+        for info_update in swap_process(file_paths):
+            yield info_update
 
+        PREVIEW = cv2.imread(file_paths[-1])[:, :, ::-1]
         WORKSPACE = temp_path
-        OUTPUT_FILE = filename
-
-        tot_exec_time = time.time() - start_time
-        _min, _sec = divmod(tot_exec_time, 60)
+        OUTPUT_FILE = file_paths[-1]
 
-        yield f"✔️ Completed in {int(_min)} min {int(_sec)} sec.", *ui_after()
+        yield get_finsh_text(start_time), *ui_after()
 
     ## ------------------------------ STREAM ------------------------------
 
     elif input_type == "Stream":
-        yield "### \n ⌛ Starting...", *ui_before()
-        global STREAMER
-        STREAMER = StreamerThread(src=directory_path)
-        STREAMER.start()
-
-        while True:
-            try:
-                target = STREAMER.frame
-                analysed_target = analyse_face(
-                    target, FACE_ANALYSER, return_single_face=False
-                )
-                if condition == "Specific Face":
-                    swapped = swap_specific(
-                        target,
-                        analysed_target,
-                        analysed_source_specific,
-                        models,
-                        threshold=distance,
-                    )
-                else:
-                    swapped = swap_face_with_condition(
-                        target, analysed_target, analysed_source, condition, age, models
-                    )
-                PREVIEW = swapped[:, :, ::-1]
-                yield f"Streaming...", *ui_before()
-            except AttributeError:
-                yield "Streaming...", *ui_before()
-        STREAMER.stop()
+        pass
 
 
 ## ------------------------------ GRADIO FUNC ------------------------------
@@ -626,10 +537,6 @@ with gr.Blocks(css=css) as interface:
                     )
 
                 with gr.Tab("🪄 Other Settings"):
-                    with gr.Accordion("Enhance Face", open=True):
-                        enable_face_enhance = gr.Checkbox(
-                            label="Enable GFPGAN", value=False, interactive=True
-                        )
                     with gr.Accordion("Advanced Mask", open=False):
                         enable_face_parser_mask = gr.Checkbox(
                             label="Enable Face Parsing",
@@ -665,6 +572,30 @@ with gr.Blocks(css=css) as interface:
                             interactive=True,
                         )
 
+                    face_scale = gr.Slider(
+                        label="Face Scale",
+                        minimum=0,
+                        maximum=2,
+                        value=1,
+                        interactive=True,
+                    )
+
+                    with gr.Accordion("Crop Mask", open=False):
+                        crop_top = gr.Number(label="Top", value=0, minimum=0, interactive=True)
+                        crop_bott = gr.Number(label="Bottom", value=0, minimum=0, interactive=True)
+                        crop_left = gr.Number(label="Left", value=0, minimum=0, interactive=True)
+                        crop_right = gr.Number(label="Right", value=0, minimum=0, interactive=True)
+
+                    enable_laplacian_blend = gr.Checkbox(
+                        label="Laplacian Blending",
+                        value=True,
+                        interactive=True,
+                    )
+
+                    face_enhancer_name = gr.Dropdown(
+                        face_enhancer_list, label="Face Enhancer", value="NONE", multiselect=False, interactive=True
+                    )
+
                 source_image_input = gr.Image(
                     label="Source face", type="filepath", interactive=True
                 )
@@ -690,7 +621,7 @@ with gr.Blocks(css=css) as interface:
 
                 with gr.Group():
                     input_type = gr.Radio(
-                        ["Image", "Video"],#["Image", "Video", "Directory", "Stream"],
+                        ["Image", "Video"],
                         label="Target Type",
                         value="Video",
                     )
@@ -701,7 +632,7 @@ with gr.Blocks(css=css) as interface:
                         )
 
                     with gr.Box(visible=True) as input_video_group:
-                        vid_widget = gr.Video #gr.Video if USE_COLAB else gr.Text
+                        vid_widget = gr.Video if USE_COLAB else gr.Text
                         video_input = vid_widget(
                             label="Target Video Path", interactive=True
                         )
@@ -794,14 +725,14 @@ with gr.Blocks(css=css) as interface:
         fn=slider_changed,
         inputs=[show_trim_preview_btn, video_input, start_frame],
         outputs=[preview_image, preview_video],
-        show_progress=False,
+        show_progress=True,
     )
 
     end_frame_event = end_frame.release(
         fn=slider_changed,
         inputs=[show_trim_preview_btn, video_input, end_frame],
         outputs=[preview_image, preview_video],
-        show_progress=False,
+        show_progress=True,
     )
 
     input_type.change(
@@ -839,12 +770,18 @@ with gr.Blocks(css=css) as interface:
         swap_option,
         age,
         distance_slider,
-        enable_face_enhance,
+        face_enhancer_name,
         enable_face_parser_mask,
         mask_include,
         mask_soft_kernel,
         mask_soft_iterations,
         blur_amount,
+        face_scale,
+        enable_laplacian_blend,
+        crop_top,
+        crop_bott,
+        crop_left,
+        crop_right,
         *src_specific_inputs,
     ]
 
@@ -857,7 +794,7 @@ with gr.Blocks(css=css) as interface:
     ]
 
     swap_event = swap_button.click(
-        fn=process, inputs=swap_inputs, outputs=swap_outputs, show_progress=False
+        fn=process, inputs=swap_inputs, outputs=swap_outputs, show_progress=True
     )
 
     cancel_button.click(
@@ -871,7 +808,7 @@ with gr.Blocks(css=css) as interface:
             start_frame_event,
             end_frame_event,
         ],
-        show_progress=False,
+        show_progress=True,
     )
     output_directory_button.click(
         lambda: open_directory(path=WORKSPACE), inputs=None, outputs=None

assets/pretrained_models/RealESRGAN_x2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c830d067d54fc767b9543a8432f36d91bc2de313584e8bbfe4ac26a47339e899
+size 67061725

assets/pretrained_models/RealESRGAN_x4.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa00f09ad753d88576b21ed977e97d634976377031b178acc3b5b238df463400
+size 67040989

assets/pretrained_models/RealESRGAN_x8.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b72fb469d12f05a4770813d2603eb1b550f40df6fb8b37d6c7bc2db3d2bff5e
+size 67189359

assets/pretrained_models/codeformer.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1009e537e0c2a07d4cabce6355f53cb66767cd4b4297ec7a4a64ca4b8a5684b7
+size 376637898

assets/pretrained_models/nsfwmodel_281.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ac92f5326f0d83f24f51ba4ac9f2a79314d29199e900a8ea495a74816ad3eb67
+size 4925

face_analyser.py

@@ -1,3 +1,9 @@
+import os
+import cv2
+import numpy as np
+from tqdm import tqdm
+from utils import scale_bbox_from_center
+
 detect_conditions = [
     "left most",
     "right most",
@@ -5,11 +11,27 @@ detect_conditions = [
     "bottom most",
     "most width",
     "most height",
+    "best detection",
 ]
 
+swap_options_list = [
+    "All face",
+    "Age less than",
+    "Age greater than",
+    "All Male",
+    "All Female",
+    "Specific Face",
+]
 
-def analyse_face(image, model, return_single_face=True, detect_condition="left most"):
+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
+        for i, face in enumerate(faces):
+            landmark = face['kps']
+            center = np.mean(landmark, axis=0)
+            landmark = center + (landmark - center) * scale
+            faces[i]['kps'] = landmark
+
     if not return_single_face:
         return faces
 
@@ -30,3 +52,79 @@ def analyse_face(image, model, return_single_face=True, detect_condition="left m
         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]
+
+
+def cosine_distance(a, b):
+    a /= np.linalg.norm(a)
+    b /= np.linalg.norm(b)
+    return 1 - np.dot(a, b)
+
+
+def get_analysed_data(face_analyser, image_sequence, source_data, swap_condition="All face", detect_condition="left most", scale=1.0):
+    if swap_condition != "Specific Face":
+        source_path, age = source_data
+        source_image = cv2.imread(source_path)
+        analysed_source = analyse_face(source_image, face_analyser, return_single_face=True, detect_condition=detect_condition, scale=scale)
+    else:
+        analysed_source_specifics = []
+        source_specifics, threshold = source_data
+        for source, specific in zip(*source_specifics):
+            if source is None or specific is None:
+                continue
+            analysed_source = analyse_face(source, face_analyser, return_single_face=True, detect_condition=detect_condition, scale=scale)
+            analysed_specific = analyse_face(specific, face_analyser, return_single_face=True, detect_condition=detect_condition, scale=scale)
+            analysed_source_specifics.append([analysed_source, analysed_specific])
+
+    analysed_target_list = []
+    analysed_source_list = []
+    whole_frame_eql_list = []
+    num_faces_per_frame = []
+
+    total_frames = len(image_sequence)
+    curr_idx = 0
+    for curr_idx, frame_path in tqdm(enumerate(image_sequence), total=total_frames, desc="Analysing face data"):
+        frame = cv2.imread(frame_path)
+        analysed_faces = analyse_face(frame, face_analyser, return_single_face=False, detect_condition=detect_condition, scale=scale)
+
+        n_faces = 0
+        for analysed_face in analysed_faces:
+            if swap_condition == "All face":
+                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 == "Age less than" and analysed_face["age"] < age:
+                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 == "Age greater than" and analysed_face["age"] > age:
+                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 == "All Male" and analysed_face["gender"] == 1:
+                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 == "All Female" and analysed_face["gender"] == 0:
+                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 == "Specific Face":
+                for analysed_source, analysed_specific in analysed_source_specifics:
+                    distance = cosine_distance(analysed_specific["embedding"], analysed_face["embedding"])
+                    if distance < threshold:
+                        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

@@ -0,0 +1,39 @@
+import os
+import torch
+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 load_face_enhancer_model(name='GFPGAN', device="cpu"):
+    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)
+    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

face_parsing/__init__.py

@@ -1 +1,3 @@
-from .swap import init_parser, swap_regions, mask_regions, mask_regions_to_list, SoftErosion
+from .swap import init_parser, swap_regions, mask_regions, mask_regions_to_list, SoftErosion
+from .model import BiSeNet
+from .parse_mask import init_parsing_model, get_parsed_mask

face_parsing/parse_mask.py

@@ -0,0 +1,50 @@
+import cv2
+import torch
+import torchvision
+import numpy as np
+import torch.nn as nn
+from PIL import Image
+from tqdm import tqdm
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+
+from . model import BiSeNet
+
+transform = transforms.Compose([
+    transforms.Resize((512, 512)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+])
+
+def init_parsing_model(model_path, device="cpu"):
+    net = BiSeNet(19)
+    net.to(device)
+    net.load_state_dict(torch.load(model_path))
+    net.eval()
+    return net
+
+def transform_images(imgs):
+    tensor_images = torch.stack([transform(Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))) for img in imgs], dim=0)
+    return tensor_images
+
+def get_parsed_mask(net, imgs, classes=[1, 2, 3, 4, 5, 10, 11, 12, 13], device="cpu", batch_size=8):
+    masks = []
+    for i in tqdm(range(0, len(imgs), batch_size), total=len(imgs) // batch_size, desc="Face-parsing"):
+        batch_imgs = imgs[i:i + batch_size]
+
+        tensor_images = transform_images(batch_imgs).to(device)
+        with torch.no_grad():
+            out = net(tensor_images)[0]
+        parsing = out.argmax(dim=1).cpu().numpy()
+        batch_masks = np.isin(parsing, classes)
+
+        masks.append(batch_masks)
+
+    masks = np.concatenate(masks, axis=0)
+    # masks = np.repeat(np.expand_dims(masks, axis=1), 3, axis=1)
+
+    for i, mask in enumerate(masks):
+        cv2.imwrite(f"mask/{i}.jpg", (mask * 255).astype("uint8"))
+
+    return masks
+

face_parsing/swap.py

@@ -98,6 +98,7 @@ def get_mask(parsing, classes):
         res += parsing == val
     return res
 
+
 def swap_regions(source, target, net, smooth_mask, includes=[1,2,3,4,5,10,11,12,13], blur=10):
     parsing = image_to_parsing(source, net)
 
@@ -117,12 +118,10 @@ def swap_regions(source, target, net, smooth_mask, includes=[1,2,3,4,5,10,11,12,
     if blur > 0:
         mask = cv2.GaussianBlur(mask, (0, 0), blur)
 
-    resized_source = cv2.resize((source/255).astype("float32"), (512, 512))
-    resized_target = cv2.resize((target/255).astype("float32"), (512, 512))
-
+    resized_source = cv2.resize((source).astype("float32"), (512, 512))
+    resized_target = cv2.resize((target).astype("float32"), (512, 512))
     result = mask * resized_source + (1 - mask) * resized_target
-    normalized_result = (result - np.min(result)) / (np.max(result) - np.min(result))
-    result = cv2.resize((result*255).astype("uint8"), (source.shape[1], source.shape[0]))
+    result = cv2.resize(result.astype("uint8"), (source.shape[1], source.shape[0]))
 
     return result
 

face_swapper.py

@@ -0,0 +1,203 @@
+import time
+import torch
+import onnx
+import cv2
+import onnxruntime
+import numpy as np
+from tqdm import tqdm
+from onnx import numpy_helper
+from skimage import transform as trans
+
+arcface_dst = np.array(
+    [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366],
+     [41.5493, 92.3655], [70.7299, 92.2041]],
+    dtype=np.float32)
+
+def estimate_norm(lmk, image_size=112, mode='arcface'):
+    assert lmk.shape == (5, 2)
+    assert image_size % 112 == 0 or image_size % 128 == 0
+    if image_size % 112 == 0:
+        ratio = float(image_size) / 112.0
+        diff_x = 0
+    else:
+        ratio = float(image_size) / 128.0
+        diff_x = 8.0 * ratio
+    dst = arcface_dst * ratio
+    dst[:, 0] += diff_x
+    tform = trans.SimilarityTransform()
+    tform.estimate(lmk, dst)
+    M = tform.params[0:2, :]
+    return M
+
+
+def norm_crop2(img, landmark, image_size=112, mode='arcface'):
+    M = estimate_norm(landmark, image_size, mode)
+    warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0)
+    return warped, M
+
+
+class Inswapper():
+    def __init__(self, model_file=None, batch_size=32, providers=['CPUExecutionProvider']):
+        self.model_file = model_file
+        self.batch_size = batch_size
+
+        model = onnx.load(self.model_file)
+        graph = model.graph
+        self.emap = numpy_helper.to_array(graph.initializer[-1])
+        self.input_mean = 0.0
+        self.input_std = 255.0
+
+        self.session_options = onnxruntime.SessionOptions()
+        self.session = onnxruntime.InferenceSession(self.model_file, sess_options=self.session_options, providers=providers)
+
+        inputs = self.session.get_inputs()
+        self.input_names = [inp.name for inp in inputs]
+        outputs = self.session.get_outputs()
+        self.output_names = [out.name for out in outputs]
+        assert len(self.output_names) == 1
+        self.output_shape = outputs[0].shape
+        input_cfg = inputs[0]
+        input_shape = input_cfg.shape
+        self.input_shape = input_shape
+        self.input_size = tuple(input_shape[2:4][::-1])
+
+    def forward(self, imgs, latents):
+        batch_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
+
+    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)
+            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
+
+    def batch_forward(self, img_list, target_f_list, source_f_list):
+        num_samples = len(img_list)
+        num_batches = (num_samples + self.batch_size - 1) // self.batch_size
+
+        preds = []
+        aimgs = []
+        ms = []
+        for i in tqdm(range(num_batches), desc="Swapping face by batch"):
+            start_idx = i * self.batch_size
+            end_idx = min((i + 1) * self.batch_size, num_samples)
+
+            batch_img = img_list[start_idx:end_idx]
+            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)
+            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_
+
+
+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[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.tile(np.expand_dims(img_mask, axis=-1), (1, 1, 3))
+
+    if laplacian_blend:
+        bgr_fake = laplacian_blending(bgr_fake.astype("float32").clip(0,255), whole_img.astype("float32").clip(0,255), img_mask.clip(0,1))
+        bgr_fake = bgr_fake.astype("float32")
+
+    fake_merged = img_mask * bgr_fake + (1 - img_mask) * whole_img.astype(np.float32)
+    return fake_merged.astype("uint8")

nsfw_detector.py

@@ -0,0 +1,61 @@
+from torchvision.transforms import Normalize
+import torchvision.transforms as T
+import torch.nn as nn
+from PIL import Image
+import numpy as np
+import torch
+import timm
+from tqdm import tqdm
+
+normalize_t = Normalize((0.4814, 0.4578, 0.4082), (0.2686, 0.2613, 0.2757))
+
+#nsfw classifier
+class NSFWClassifier(nn.Module):
+    def __init__(self):
+        super().__init__()
+        nsfw_model=self
+        nsfw_model.root_model = timm.create_model('convnext_base_in22ft1k', pretrained=True)
+        nsfw_model.linear_probe = nn.Linear(1024, 1, bias=False)
+
+    def forward(self, x):
+        nsfw_model = self
+        x = normalize_t(x)
+        x = nsfw_model.root_model.stem(x)
+        x = nsfw_model.root_model.stages(x)
+        x = nsfw_model.root_model.head.global_pool(x)
+        x = nsfw_model.root_model.head.norm(x)
+        x = nsfw_model.root_model.head.flatten(x)
+        x = nsfw_model.linear_probe(x)
+        return x
+
+    def is_nsfw(self, img_paths, threshold = 0.93):
+        skip_step = 1
+        total_len = len(img_paths)
+        if total_len < 100: skip_step = 1
+        if total_len > 100 and total_len < 500: skip_step = 10
+        if total_len > 500 and total_len < 1000: skip_step = 20
+        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))
+            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
+        return False
+
+def get_nsfw_detector(model_path='nsfwmodel_281.pth', device="cpu"):
+    #load base model
+    nsfw_model = NSFWClassifier()
+    nsfw_model = nsfw_model.eval()
+    #load linear weights
+    linear_pth = model_path
+    linear_state_dict = torch.load(linear_pth, map_location='cpu')
+    nsfw_model.linear_probe.load_state_dict(linear_state_dict)
+    nsfw_model = nsfw_model.to(device)
+    return nsfw_model

requirements.txt

@@ -4,8 +4,10 @@ gradio>=3.33.1
 insightface==0.7.3
 moviepy>=1.0.3
 numpy
-opencv-python>=4.7.0.72
-opencv-python-headless>=4.7.0.72
 onnx==1.14.0
 onnxruntime==1.15.0
+opencv-python>=4.7.0.72
+opencv-python-headless>=4.7.0.72
 gfpgan==1.3.8
+timm==0.9.2
+

upscaler/RealESRGAN/__init__.py

@@ -0,0 +1 @@
+from .model import RealESRGAN

upscaler/RealESRGAN/arch_utils.py

@@ -0,0 +1,197 @@
+import math
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+from torch.nn import init as init
+from torch.nn.modules.batchnorm import _BatchNorm
+
+@torch.no_grad()
+def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs):
+    """Initialize network weights.
+
+    Args:
+        module_list (list[nn.Module] | nn.Module): Modules to be initialized.
+        scale (float): Scale initialized weights, especially for residual
+            blocks. Default: 1.
+        bias_fill (float): The value to fill bias. Default: 0
+        kwargs (dict): Other arguments for initialization function.
+    """
+    if not isinstance(module_list, list):
+        module_list = [module_list]
+    for module in module_list:
+        for m in module.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, nn.Linear):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, _BatchNorm):
+                init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+
+
+def make_layer(basic_block, num_basic_block, **kwarg):
+    """Make layers by stacking the same blocks.
+
+    Args:
+        basic_block (nn.module): nn.module class for basic block.
+        num_basic_block (int): number of blocks.
+
+    Returns:
+        nn.Sequential: Stacked blocks in nn.Sequential.
+    """
+    layers = []
+    for _ in range(num_basic_block):
+        layers.append(basic_block(**kwarg))
+    return nn.Sequential(*layers)
+
+
+class ResidualBlockNoBN(nn.Module):
+    """Residual block without BN.
+
+    It has a style of:
+        ---Conv-ReLU-Conv-+-
+         |________________|
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+            Default: 64.
+        res_scale (float): Residual scale. Default: 1.
+        pytorch_init (bool): If set to True, use pytorch default init,
+            otherwise, use default_init_weights. Default: False.
+    """
+
+    def __init__(self, num_feat=64, res_scale=1, pytorch_init=False):
+        super(ResidualBlockNoBN, self).__init__()
+        self.res_scale = res_scale
+        self.conv1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
+        self.conv2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
+        self.relu = nn.ReLU(inplace=True)
+
+        if not pytorch_init:
+            default_init_weights([self.conv1, self.conv2], 0.1)
+
+    def forward(self, x):
+        identity = x
+        out = self.conv2(self.relu(self.conv1(x)))
+        return identity + out * self.res_scale
+
+
+class Upsample(nn.Sequential):
+    """Upsample module.
+
+    Args:
+        scale (int): Scale factor. Supported scales: 2^n and 3.
+        num_feat (int): Channel number of intermediate features.
+    """
+
+    def __init__(self, scale, num_feat):
+        m = []
+        if (scale & (scale - 1)) == 0:  # scale = 2^n
+            for _ in range(int(math.log(scale, 2))):
+                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
+                m.append(nn.PixelShuffle(2))
+        elif scale == 3:
+            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
+            m.append(nn.PixelShuffle(3))
+        else:
+            raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')
+        super(Upsample, self).__init__(*m)
+
+
+def flow_warp(x, flow, interp_mode='bilinear', padding_mode='zeros', align_corners=True):
+    """Warp an image or feature map with optical flow.
+
+    Args:
+        x (Tensor): Tensor with size (n, c, h, w).
+        flow (Tensor): Tensor with size (n, h, w, 2), normal value.
+        interp_mode (str): 'nearest' or 'bilinear'. Default: 'bilinear'.
+        padding_mode (str): 'zeros' or 'border' or 'reflection'.
+            Default: 'zeros'.
+        align_corners (bool): Before pytorch 1.3, the default value is
+            align_corners=True. After pytorch 1.3, the default value is
+            align_corners=False. Here, we use the True as default.
+
+    Returns:
+        Tensor: Warped image or feature map.
+    """
+    assert x.size()[-2:] == flow.size()[1:3]
+    _, _, h, w = x.size()
+    # create mesh grid
+    grid_y, grid_x = torch.meshgrid(torch.arange(0, h).type_as(x), torch.arange(0, w).type_as(x))
+    grid = torch.stack((grid_x, grid_y), 2).float()  # W(x), H(y), 2
+    grid.requires_grad = False
+
+    vgrid = grid + flow
+    # scale grid to [-1,1]
+    vgrid_x = 2.0 * vgrid[:, :, :, 0] / max(w - 1, 1) - 1.0
+    vgrid_y = 2.0 * vgrid[:, :, :, 1] / max(h - 1, 1) - 1.0
+    vgrid_scaled = torch.stack((vgrid_x, vgrid_y), dim=3)
+    output = F.grid_sample(x, vgrid_scaled, mode=interp_mode, padding_mode=padding_mode, align_corners=align_corners)
+
+    # TODO, what if align_corners=False
+    return output
+
+
+def resize_flow(flow, size_type, sizes, interp_mode='bilinear', align_corners=False):
+    """Resize a flow according to ratio or shape.
+
+    Args:
+        flow (Tensor): Precomputed flow. shape [N, 2, H, W].
+        size_type (str): 'ratio' or 'shape'.
+        sizes (list[int | float]): the ratio for resizing or the final output
+            shape.
+            1) The order of ratio should be [ratio_h, ratio_w]. For
+            downsampling, the ratio should be smaller than 1.0 (i.e., ratio
+            < 1.0). For upsampling, the ratio should be larger than 1.0 (i.e.,
+            ratio > 1.0).
+            2) The order of output_size should be [out_h, out_w].
+        interp_mode (str): The mode of interpolation for resizing.
+            Default: 'bilinear'.
+        align_corners (bool): Whether align corners. Default: False.
+
+    Returns:
+        Tensor: Resized flow.
+    """
+    _, _, flow_h, flow_w = flow.size()
+    if size_type == 'ratio':
+        output_h, output_w = int(flow_h * sizes[0]), int(flow_w * sizes[1])
+    elif size_type == 'shape':
+        output_h, output_w = sizes[0], sizes[1]
+    else:
+        raise ValueError(f'Size type should be ratio or shape, but got type {size_type}.')
+
+    input_flow = flow.clone()
+    ratio_h = output_h / flow_h
+    ratio_w = output_w / flow_w
+    input_flow[:, 0, :, :] *= ratio_w
+    input_flow[:, 1, :, :] *= ratio_h
+    resized_flow = F.interpolate(
+        input=input_flow, size=(output_h, output_w), mode=interp_mode, align_corners=align_corners)
+    return resized_flow
+
+
+# TODO: may write a cpp file
+def pixel_unshuffle(x, scale):
+    """ Pixel unshuffle.
+
+    Args:
+        x (Tensor): Input feature with shape (b, c, hh, hw).
+        scale (int): Downsample ratio.
+
+    Returns:
+        Tensor: the pixel unshuffled feature.
+    """
+    b, c, hh, hw = x.size()
+    out_channel = c * (scale**2)
+    assert hh % scale == 0 and hw % scale == 0
+    h = hh // scale
+    w = hw // scale
+    x_view = x.view(b, c, h, scale, w, scale)
+    return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w)

upscaler/RealESRGAN/model.py

@@ -0,0 +1,90 @@
+import os
+import torch
+from torch.nn import functional as F
+from PIL import Image
+import numpy as np
+import cv2
+
+from .rrdbnet_arch import RRDBNet
+from .utils import pad_reflect, split_image_into_overlapping_patches, stich_together, \
+                   unpad_image
+
+
+HF_MODELS = {
+    2: dict(
+        repo_id='sberbank-ai/Real-ESRGAN',
+        filename='RealESRGAN_x2.pth',
+    ),
+    4: dict(
+        repo_id='sberbank-ai/Real-ESRGAN',
+        filename='RealESRGAN_x4.pth',
+    ),
+    8: dict(
+        repo_id='sberbank-ai/Real-ESRGAN',
+        filename='RealESRGAN_x8.pth',
+    ),
+}
+
+
+class RealESRGAN:
+    def __init__(self, device, scale=4):
+        self.device = device
+        self.scale = scale
+        self.model = RRDBNet(
+            num_in_ch=3, num_out_ch=3, num_feat=64,
+            num_block=23, num_grow_ch=32, scale=scale
+        )
+
+    def load_weights(self, model_path, download=True):
+        if not os.path.exists(model_path) and download:
+            from huggingface_hub import hf_hub_url, cached_download
+            assert self.scale in [2,4,8], 'You can download models only with scales: 2, 4, 8'
+            config = HF_MODELS[self.scale]
+            cache_dir = os.path.dirname(model_path)
+            local_filename = os.path.basename(model_path)
+            config_file_url = hf_hub_url(repo_id=config['repo_id'], filename=config['filename'])
+            cached_download(config_file_url, cache_dir=cache_dir, force_filename=local_filename)
+            print('Weights downloaded to:', os.path.join(cache_dir, local_filename))
+
+        loadnet = torch.load(model_path)
+        if 'params' in loadnet:
+            self.model.load_state_dict(loadnet['params'], strict=True)
+        elif 'params_ema' in loadnet:
+            self.model.load_state_dict(loadnet['params_ema'], strict=True)
+        else:
+            self.model.load_state_dict(loadnet, strict=True)
+        self.model.eval()
+        self.model.to(self.device)
+
+    @torch.cuda.amp.autocast()
+    def predict(self, lr_image, batch_size=4, patches_size=192,
+                padding=24, pad_size=15):
+        scale = self.scale
+        device = self.device
+        lr_image = np.array(lr_image)
+        lr_image = pad_reflect(lr_image, pad_size)
+
+        patches, p_shape = split_image_into_overlapping_patches(
+            lr_image, patch_size=patches_size, padding_size=padding
+        )
+        img = torch.FloatTensor(patches/255).permute((0,3,1,2)).to(device).detach()
+
+        with torch.no_grad():
+            res = self.model(img[0:batch_size])
+            for i in range(batch_size, img.shape[0], batch_size):
+                res = torch.cat((res, self.model(img[i:i+batch_size])), 0)
+
+        sr_image = res.permute((0,2,3,1)).clamp_(0, 1).cpu()
+        np_sr_image = sr_image.numpy()
+
+        padded_size_scaled = tuple(np.multiply(p_shape[0:2], scale)) + (3,)
+        scaled_image_shape = tuple(np.multiply(lr_image.shape[0:2], scale)) + (3,)
+        np_sr_image = stich_together(
+            np_sr_image, padded_image_shape=padded_size_scaled,
+            target_shape=scaled_image_shape, padding_size=padding * scale
+        )
+        sr_img = (np_sr_image*255).astype(np.uint8)
+        sr_img = unpad_image(sr_img, pad_size*scale)
+        #sr_img = Image.fromarray(sr_img)
+
+        return sr_img

upscaler/RealESRGAN/rrdbnet_arch.py

@@ -0,0 +1,121 @@
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+
+from .arch_utils import default_init_weights, make_layer, pixel_unshuffle
+
+
+class ResidualDenseBlock(nn.Module):
+    """Residual Dense Block.
+
+    Used in RRDB block in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat=64, num_grow_ch=32):
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
+        self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+        # initialization
+        default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+
+    def forward(self, x):
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        # Emperically, we use 0.2 to scale the residual for better performance
+        return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+    """Residual in Residual Dense Block.
+
+    Used in RRDB-Net in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat, num_grow_ch=32):
+        super(RRDB, self).__init__()
+        self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)
+
+    def forward(self, x):
+        out = self.rdb1(x)
+        out = self.rdb2(out)
+        out = self.rdb3(out)
+        # Emperically, we use 0.2 to scale the residual for better performance
+        return out * 0.2 + x
+
+
+class RRDBNet(nn.Module):
+    """Networks consisting of Residual in Residual Dense Block, which is used
+    in ESRGAN.
+
+    ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks.
+
+    We extend ESRGAN for scale x2 and scale x1.
+    Note: This is one option for scale 1, scale 2 in RRDBNet.
+    We first employ the pixel-unshuffle (an inverse operation of pixelshuffle to reduce the spatial size
+    and enlarge the channel size before feeding inputs into the main ESRGAN architecture.
+
+    Args:
+        num_in_ch (int): Channel number of inputs.
+        num_out_ch (int): Channel number of outputs.
+        num_feat (int): Channel number of intermediate features.
+            Default: 64
+        num_block (int): Block number in the trunk network. Defaults: 23
+        num_grow_ch (int): Channels for each growth. Default: 32.
+    """
+
+    def __init__(self, num_in_ch, num_out_ch, scale=4, num_feat=64, num_block=23, num_grow_ch=32):
+        super(RRDBNet, self).__init__()
+        self.scale = scale
+        if scale == 2:
+            num_in_ch = num_in_ch * 4
+        elif scale == 1:
+            num_in_ch = num_in_ch * 16
+        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
+        self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch)
+        self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        # upsample
+        self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        if scale == 8:
+            self.conv_up3 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x):
+        if self.scale == 2:
+            feat = pixel_unshuffle(x, scale=2)
+        elif self.scale == 1:
+            feat = pixel_unshuffle(x, scale=4)
+        else:
+            feat = x
+        feat = self.conv_first(feat)
+        body_feat = self.conv_body(self.body(feat))
+        feat = feat + body_feat
+        # upsample
+        feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))
+        feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest')))
+        if self.scale == 8:
+            feat = self.lrelu(self.conv_up3(F.interpolate(feat, scale_factor=2, mode='nearest')))
+        out = self.conv_last(self.lrelu(self.conv_hr(feat)))
+        return out

upscaler/RealESRGAN/utils.py

@@ -0,0 +1,133 @@
+import numpy as np
+import torch
+from PIL import Image
+import os
+import io
+
+def pad_reflect(image, pad_size):
+    imsize = image.shape
+    height, width = imsize[:2]
+    new_img = np.zeros([height+pad_size*2, width+pad_size*2, imsize[2]]).astype(np.uint8)
+    new_img[pad_size:-pad_size, pad_size:-pad_size, :] = image
+    
+    new_img[0:pad_size, pad_size:-pad_size, :] = np.flip(image[0:pad_size, :, :], axis=0) #top
+    new_img[-pad_size:, pad_size:-pad_size, :] = np.flip(image[-pad_size:, :, :], axis=0) #bottom
+    new_img[:, 0:pad_size, :] = np.flip(new_img[:, pad_size:pad_size*2, :], axis=1) #left
+    new_img[:, -pad_size:, :] = np.flip(new_img[:, -pad_size*2:-pad_size, :], axis=1) #right
+    
+    return new_img
+
+def unpad_image(image, pad_size):
+    return image[pad_size:-pad_size, pad_size:-pad_size, :]
+
+
+def process_array(image_array, expand=True):
+    """ Process a 3-dimensional array into a scaled, 4 dimensional batch of size 1. """
+    
+    image_batch = image_array / 255.0
+    if expand:
+        image_batch = np.expand_dims(image_batch, axis=0)
+    return image_batch
+
+
+def process_output(output_tensor):
+    """ Transforms the 4-dimensional output tensor into a suitable image format. """
+    
+    sr_img = output_tensor.clip(0, 1) * 255
+    sr_img = np.uint8(sr_img)
+    return sr_img
+
+
+def pad_patch(image_patch, padding_size, channel_last=True):
+    """ Pads image_patch with with padding_size edge values. """
+    
+    if channel_last:
+        return np.pad(
+            image_patch,
+            ((padding_size, padding_size), (padding_size, padding_size), (0, 0)),
+            'edge',
+        )
+    else:
+        return np.pad(
+            image_patch,
+            ((0, 0), (padding_size, padding_size), (padding_size, padding_size)),
+            'edge',
+        )
+
+
+def unpad_patches(image_patches, padding_size):
+    return image_patches[:, padding_size:-padding_size, padding_size:-padding_size, :]
+
+
+def split_image_into_overlapping_patches(image_array, patch_size, padding_size=2):
+    """ Splits the image into partially overlapping patches.
+    The patches overlap by padding_size pixels.
+    Pads the image twice:
+        - first to have a size multiple of the patch size,
+        - then to have equal padding at the borders.
+    Args:
+        image_array: numpy array of the input image.
+        patch_size: size of the patches from the original image (without padding).
+        padding_size: size of the overlapping area.
+    """
+    
+    xmax, ymax, _ = image_array.shape
+    x_remainder = xmax % patch_size
+    y_remainder = ymax % patch_size
+    
+    # modulo here is to avoid extending of patch_size instead of 0
+    x_extend = (patch_size - x_remainder) % patch_size
+    y_extend = (patch_size - y_remainder) % patch_size
+    
+    # make sure the image is divisible into regular patches
+    extended_image = np.pad(image_array, ((0, x_extend), (0, y_extend), (0, 0)), 'edge')
+    
+    # add padding around the image to simplify computations
+    padded_image = pad_patch(extended_image, padding_size, channel_last=True)
+    
+    xmax, ymax, _ = padded_image.shape
+    patches = []
+    
+    x_lefts = range(padding_size, xmax - padding_size, patch_size)
+    y_tops = range(padding_size, ymax - padding_size, patch_size)
+    
+    for x in x_lefts:
+        for y in y_tops:
+            x_left = x - padding_size
+            y_top = y - padding_size
+            x_right = x + patch_size + padding_size
+            y_bottom = y + patch_size + padding_size
+            patch = padded_image[x_left:x_right, y_top:y_bottom, :]
+            patches.append(patch)
+    
+    return np.array(patches), padded_image.shape
+
+
+def stich_together(patches, padded_image_shape, target_shape, padding_size=4):
+    """ Reconstruct the image from overlapping patches.
+    After scaling, shapes and padding should be scaled too.
+    Args:
+        patches: patches obtained with split_image_into_overlapping_patches
+        padded_image_shape: shape of the padded image contructed in split_image_into_overlapping_patches
+        target_shape: shape of the final image
+        padding_size: size of the overlapping area.
+    """
+    
+    xmax, ymax, _ = padded_image_shape
+    patches = unpad_patches(patches, padding_size)
+    patch_size = patches.shape[1]
+    n_patches_per_row = ymax // patch_size
+    
+    complete_image = np.zeros((xmax, ymax, 3))
+    
+    row = -1
+    col = 0
+    for i in range(len(patches)):
+        if i % n_patches_per_row == 0:
+            row += 1
+            col = 0
+        complete_image[
+        row * patch_size: (row + 1) * patch_size, col * patch_size: (col + 1) * patch_size,:
+        ] = patches[i]
+        col += 1
+    return complete_image[0: target_shape[0], 0: target_shape[1], :]

utils.py

@@ -110,3 +110,60 @@ def add_logo_to_image(img, logo=logo_image):
             roi[0] : roi[0] + logo_size, roi[1] : roi[1] + logo_size, c
         ]
     return img
+
+def split_list_by_lengths(data, length_list):
+    split_data = []
+    start_idx = 0
+    for length in length_list:
+        end_idx = start_idx + length
+        sublist = data[start_idx:end_idx]
+        split_data.append(sublist)
+        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
+    duration = video_clip.duration
+    total_frames = video_clip.reader.nframes
+    audio_clip = video_clip.audio if video_clip.audio is not None else None
+    edited_video_clip = ImageSequenceClip(image_sequence, fps=fps)
+
+    if audio_clip is not None:
+        edited_video_clip = edited_video_clip.set_audio(audio_clip)
+
+    edited_video_clip.set_duration(duration).write_videofile(
+        output_file_name, codec="libx264"
+    )
+    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
+
+    # Calculate the center point of the bbox
+    center_x = (x1 + x2) / 2
+    center_y = (y1 + y2) / 2
+
+    # Calculate the new width and height of the bbox based on the scaling factors
+    width = x2 - x1
+    height = y2 - y1
+    new_width = width * scale_width
+    new_height = height * scale_height
+
+    # Calculate the new coordinates of the bbox, considering the image boundaries
+    new_x1 = center_x - new_width / 2
+    new_y1 = center_y - new_height / 2
+    new_x2 = center_x + new_width / 2
+    new_y2 = center_y + new_height / 2
+
+    # Adjust the coordinates to ensure the bbox remains within the image boundaries
+    new_x1 = max(0, new_x1)
+    new_y1 = max(0, new_y1)
+    new_x2 = min(image_width - 1, new_x2)
+    new_y2 = min(image_height - 1, new_y2)
+
+    # Return the scaled bbox coordinates
+    scaled_bbox = [new_x1, new_y1, new_x2, new_y2]
+    return scaled_bbox