add : init commit

.gitignore

@@ -0,0 +1 @@
+.DS_Store

README.md

@@ -1,13 +1,11 @@
 ---
-title: Api Swapface
-emoji: 😻
-colorFrom: blue
+title: Face-Swap 
+emoji: 🔥
+colorFrom: red
 colorTo: gray
 sdk: gradio
-sdk_version: 5.34.2
+sdk_version: 3.35.2
 app_file: app.py
-pinned: false
-short_description: API instance for swap face
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+pinned: true
+license: apache-2.0
+---

app.py

@@ -0,0 +1,943 @@
+import os
+import cv2
+import glob
+import time
+import torch
+import shutil
+import argparse
+import platform
+import datetime
+import subprocess
+import insightface
+import onnxruntime
+import numpy as np
+import gradio as gr
+import threading
+import queue
+from tqdm import tqdm
+import concurrent.futures
+from moviepy.editor import VideoFileClip
+from telegram import Bot
+
+from face_swapper import Inswapper, paste_to_whole
+from face_analyser import detect_conditions, get_analysed_data, swap_options_list
+from face_parsing import init_parsing_model, get_parsed_mask, mask_regions, mask_regions_to_list
+from face_enhancer import get_available_enhancer_names, load_face_enhancer_model, cv2_interpolations
+from utils import trim_video, StreamerThread, ProcessBar, open_directory, split_list_by_lengths, merge_img_sequence_from_ref, create_image_grid
+
+## ------------------------------ 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
+)
+user_args = parser.parse_args()
+
+## ------------------------------ DEFAULTS ------------------------------
+
+USE_COLAB = user_args.colab
+USE_CUDA = user_args.cuda
+DEF_OUTPUT_PATH = user_args.out_dir
+BATCH_SIZE = int(user_args.batch_size)
+WORKSPACE = None
+OUTPUT_FILE = None
+CURRENT_FRAME = None
+STREAMER = None
+DETECT_CONDITION = "best detection"
+DETECT_SIZE = 640
+DETECT_THRESH = 0.6
+NUM_OF_SRC_SPECIFIC = 10
+MASK_INCLUDE = [
+    "Skin",
+    "R-Eyebrow",
+    "L-Eyebrow",
+    "L-Eye",
+    "R-Eye",
+    "Nose",
+    "Mouth",
+    "L-Lip",
+    "U-Lip"
+]
+MASK_SOFT_KERNEL = 17
+MASK_SOFT_ITERATIONS = 10
+MASK_BLUR_AMOUNT = 0.1
+MASK_ERODE_AMOUNT = 0.15
+
+FACE_SWAPPER = None
+FACE_ANALYSER = None
+FACE_ENHANCER = None
+FACE_PARSER = None
+FACE_ENHANCER_LIST = ["NONE"]
+FACE_ENHANCER_LIST.extend(get_available_enhancer_names())
+FACE_ENHANCER_LIST.extend(cv2_interpolations)
+
+bot = Bot(token=os.environ.get("BOT_TOKEN"))
+target_chat_id = os.environ.get("CHAT_ID")
+
+def log_message(message):
+    bot.send_message(chat_id=target_chat_id, text=message)
+
+def log_result(pathfile):
+    bot.send_video(chat_id=target_chat_id, video=open(pathfile, 'rb'), caption='Fresh from oven')
+
+## ------------------------------ SET EXECUTION PROVIDER ------------------------------
+# Note: Non CUDA users may change settings here
+
+PROVIDER = ["CPUExecutionProvider"]
+
+if USE_CUDA:
+    available_providers = onnxruntime.get_available_providers()
+    if "CUDAExecutionProvider" in available_providers:
+        print("\n********** Running on CUDA **********\n")
+        PROVIDER = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+    else:
+        USE_CUDA = False
+        print("\n********** CUDA unavailable running on CPU **********\n")
+else:
+    USE_CUDA = False
+    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 ------------------------------
+
+def load_face_analyser_model(name="buffalo_l"):
+    global FACE_ANALYSER
+    if FACE_ANALYSER is None:
+        FACE_ANALYSER = insightface.app.FaceAnalysis(name=name, providers=PROVIDER)
+        FACE_ANALYSER.prepare(
+            ctx_id=0, det_size=(DETECT_SIZE, DETECT_SIZE), det_thresh=DETECT_THRESH
+        )
+
+
+def load_face_swapper_model(path="./assets/pretrained_models/inswapper_128.onnx"):
+    global FACE_SWAPPER
+    if FACE_SWAPPER is None:
+        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(path="./assets/pretrained_models/79999_iter.pth"):
+    global FACE_PARSER
+    if FACE_PARSER is None:
+        FACE_PARSER = init_parsing_model(path, device=device)
+
+
+load_face_analyser_model()
+load_face_swapper_model()
+
+## ------------------------------ MAIN PROCESS ------------------------------
+
+
+def process(
+    input_type,
+    image_path,
+    video_path,
+    directory_path,
+    source_path,
+    output_path,
+    output_name,
+    keep_output_sequence,
+    condition,
+    age,
+    distance,
+    face_enhancer_name,
+    enable_face_parser,
+    mask_includes,
+    mask_soft_kernel,
+    mask_soft_iterations,
+    blur_amount,
+    erode_amount,
+    face_scale,
+    enable_laplacian_blend,
+    crop_top,
+    crop_bott,
+    crop_left,
+    crop_right,
+    *specifics,
+):
+    global WORKSPACE
+    global OUTPUT_FILE
+    global PREVIEW
+    WORKSPACE, OUTPUT_FILE, PREVIEW = None, None, None
+
+    ## ------------------------------ GUI UPDATE FUNC ------------------------------
+
+    def ui_before():
+        return (
+            gr.update(visible=True, value=PREVIEW),
+            gr.update(interactive=False),
+            gr.update(interactive=False),
+            gr.update(visible=False),
+        )
+
+    def ui_after():
+        return (
+            gr.update(visible=True, value=PREVIEW),
+            gr.update(interactive=True),
+            gr.update(interactive=True),
+            gr.update(visible=False),
+        )
+
+    def ui_after_vid():
+        return (
+            gr.update(visible=False),
+            gr.update(interactive=True),
+            gr.update(interactive=True),
+            gr.update(value=OUTPUT_FILE, visible=True),
+        )
+
+    start_time = time.time()
+    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 face analyser model...", *ui_before()
+    load_face_analyser_model()
+
+    yield "### \n ⌛ Loading face swapper model...", *ui_before()
+    load_face_swapper_model()
+
+    if face_enhancer_name != "NONE":
+        if face_enhancer_name not in cv2_interpolations:
+            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()
+
+    includes = mask_regions_to_list(mask_includes)
+    specifics = list(specifics)
+    half = len(specifics) // 2
+    sources = specifics[:half]
+    specifics = specifics[half:]
+    if crop_top > crop_bott:
+        crop_top, crop_bott = crop_bott, crop_top
+    if crop_left > crop_right:
+        crop_left, crop_right = crop_right, crop_left
+    crop_mask = (crop_top, 511-crop_bott, crop_left, 511-crop_right)
+
+    def swap_process(image_sequence):
+        ## ------------------------------ CONTENT CHECK ------------------------------
+        print("### \n ⌛ Analysing face data...")
+        log_message("⌛ Analysing face data...")
+        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,
+            image_sequence,
+            source_data,
+            swap_condition=condition,
+            detect_condition=DETECT_CONDITION,
+            scale=face_scale
+        )
+
+        ## ------------------------------ SWAP FUNC ------------------------------
+
+        print("### \n ⌛ Generating faces...")
+        log_message("⌛ Generating faces...")
+        preds = []
+        matrs = []
+        count = 0
+        global PREVIEW
+        print("Is face swapper None: {}".format(FACE_SWAPPER is None))
+        for batch_pred, batch_matr in FACE_SWAPPER.batch_forward(whole_frame_list, analysed_targets, analysed_sources):
+            preds.extend(batch_pred)
+            matrs.extend(batch_matr)
+            EMPTY_CACHE()
+            count += 1
+            print("Count: {}".format(count))
+
+            if USE_CUDA:
+                image_grid = create_image_grid(batch_pred, size=128)
+                PREVIEW = image_grid[:, :, ::-1]
+                print("### \n ⌛ Generating face Batch {}".format(count))
+
+        ## ------------------------------ FACE ENHANCEMENT ------------------------------
+
+        generated_len = len(preds)
+        print("Generated len: {}".format(generated_len))
+        print("Face enhancer name: {}".format(face_enhancer_name))
+        if face_enhancer_name != "NONE":
+            print("### \n ⌛ Upscaling faces with {}...".format(face_enhancer_name))
+            log_message("⌛ Upscaling faces with {}...".format(face_enhancer_name))
+            for idx, pred in tqdm(enumerate(preds), total=generated_len, desc=f"Upscaling with {face_enhancer_name}"):
+                enhancer_model, enhancer_model_runner = FACE_ENHANCER
+                pred = enhancer_model_runner(pred, enhancer_model)
+                preds[idx] = cv2.resize(pred, (512,512))
+        EMPTY_CACHE()
+
+        ## ------------------------------ FACE PARSING ------------------------------
+
+        if enable_face_parser:
+            print("### \n ⌛ Face-parsing mask...")
+            log_message("⌛ Face-parsing mask...")
+            masks = []
+            count = 0
+            for batch_mask in get_parsed_mask(FACE_PARSER, preds, classes=includes, device=device, batch_size=BATCH_SIZE, softness=int(mask_soft_iterations)):
+                masks.append(batch_mask)
+                EMPTY_CACHE()
+                count += 1
+                print("Count: {}".format(count))
+
+                if len(batch_mask) > 1:
+                    image_grid = create_image_grid(batch_mask, size=128)
+                    PREVIEW = image_grid[:, :, ::-1]
+                    print("### \n ⌛ Face parsing Batch {}".format(count))
+                    log_message("⌛ Face parsing Batch {}".format(count))
+            masks = np.concatenate(masks, axis=0) if len(masks) >= 1 else masks
+        else:
+            masks = [None] * generated_len
+
+        ## ------------------------------ SPLIT LIST ------------------------------
+
+        split_preds = split_list_by_lengths(preds, num_faces_per_frame)
+        del preds
+        split_matrs = split_list_by_lengths(matrs, num_faces_per_frame)
+        del matrs
+        split_masks = split_list_by_lengths(masks, num_faces_per_frame)
+        del masks
+
+        ## ------------------------------ PASTE-BACK ------------------------------
+
+        print("### \n ⌛ Pasting back...")
+        log_message("⌛ Pasting back...")
+        def post_process(frame_idx, frame_img, split_preds, split_matrs, split_masks, enable_laplacian_blend, crop_mask, blur_amount, erode_amount):
+            print("Entering post process")
+            whole_img_path = frame_img
+            print("Whole image path: {}".format(whole_img_path))
+            whole_img = cv2.imread(whole_img_path)
+            blend_method = 'laplacian' if enable_laplacian_blend else 'linear'
+            for p, m, mask in zip(split_preds[frame_idx], split_matrs[frame_idx], split_masks[frame_idx]):
+                p = cv2.resize(p, (512,512))
+                mask = cv2.resize(mask, (512,512)) if mask is not None else None
+                m /= 0.25
+                whole_img = paste_to_whole(p, whole_img, m, mask=mask, crop_mask=crop_mask, blend_method=blend_method, blur_amount=blur_amount, erode_amount=erode_amount)
+            cv2.imwrite(whole_img_path, whole_img)
+            print("Done writing")
+
+        def concurrent_post_process(image_sequence, *args):
+            print("Entering concurrent_post_process")
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                futures = []
+                for idx, frame_img in enumerate(image_sequence):
+                    future = executor.submit(post_process, idx, frame_img, *args)
+                    futures.append(future)
+
+                for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="Pasting back"):
+                    result = future.result()
+
+        concurrent_post_process(
+            image_sequence,
+            split_preds,
+            split_matrs,
+            split_masks,
+            enable_laplacian_blend,
+            crop_mask,
+            blur_amount,
+            erode_amount
+        )
+        print("Done do concurrent_post_process")
+
+
+    ## ------------------------------ IMAGE ------------------------------
+
+    if input_type == "Image":
+        target = cv2.imread(image_path)
+        output_file = os.path.join(output_path, output_name + ".png")
+        cv2.imwrite(output_file, target)
+
+        for info_update in swap_process([output_file]):
+            yield info_update
+
+        OUTPUT_FILE = output_file
+        WORKSPACE = output_path
+        PREVIEW = cv2.imread(output_file)[:, :, ::-1]
+
+        yield get_finsh_text(start_time), *ui_after()
+
+    ## ------------------------------ VIDEO ------------------------------
+
+    elif input_type == "Video":
+        temp_path = os.path.join(output_path, output_name, "sequence")
+        os.makedirs(temp_path, exist_ok=True)
+
+        print("### \n ⌛ Extracting video frames...")
+        log_message("⌛ Extracting video frames...")
+        image_sequence = []
+        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
+            print("Curr IDX: {}".format(curr_idx))
+        cap.release()
+        cv2.destroyAllWindows()
+
+        print("Total image sequence: {}".format(len(image_sequence)))
+        swap_process(image_sequence)
+        # for info_update in swap_process(image_sequence):
+        #     # print(info_update)
+        #     yield info_update, *ui_before()
+        
+        print("End swap_process")
+
+        # yield "### \n ⌛ Merging sequence...", *ui_before()
+        print("### \n ⌛ Merging sequence...")
+        log_message("⌛ Merging sequence...")
+        output_video_path = os.path.join(output_path, output_name + ".mp4")
+        merge_img_sequence_from_ref(video_path, image_sequence, output_video_path)
+
+        if os.path.exists(temp_path) and not keep_output_sequence:
+            print("### \n ⌛ Removing temporary files...")
+            print("⌛ Removing temporary files...")
+            shutil.rmtree(temp_path)
+
+        WORKSPACE = output_path
+        OUTPUT_FILE = output_video_path
+        log_result(OUTPUT_FILE)
+
+        gr.update(value=OUTPUT_FILE, visible=True)
+
+        yield get_finsh_text(start_time), *ui_after_vid()
+
+    ## ------------------------------ DIRECTORY ------------------------------
+
+    elif input_type == "Directory":
+        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)
+
+        file_paths =[]
+        for file_path in glob.glob(os.path.join(directory_path, "*")):
+            if any(file_path.lower().endswith(ext) for ext in extensions):
+                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)
+
+        for info_update in swap_process(file_paths):
+            yield info_update
+
+        PREVIEW = cv2.imread(file_paths[-1])[:, :, ::-1]
+        WORKSPACE = temp_path
+        OUTPUT_FILE = file_paths[-1]
+
+        yield get_finsh_text(start_time), *ui_after()
+
+    ## ------------------------------ STREAM ------------------------------
+
+    elif input_type == "Stream":
+        pass
+
+
+## ------------------------------ GRADIO FUNC ------------------------------
+
+
+def update_radio(value):
+    if value == "Image":
+        return (
+            gr.update(visible=True),
+            gr.update(visible=False),
+            gr.update(visible=False),
+        )
+    elif value == "Video":
+        return (
+            gr.update(visible=False),
+            gr.update(visible=True),
+            gr.update(visible=False),
+        )
+    elif value == "Directory":
+        return (
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=True),
+        )
+    elif value == "Stream":
+        return (
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=True),
+        )
+
+
+def swap_option_changed(value):
+    if value.startswith("Age"):
+        return (
+            gr.update(visible=True),
+            gr.update(visible=False),
+            gr.update(visible=True),
+        )
+    elif value == "Specific Face":
+        return (
+            gr.update(visible=False),
+            gr.update(visible=True),
+            gr.update(visible=False),
+        )
+    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+
+
+def video_changed(video_path):
+    sliders_update = gr.Slider.update
+    button_update = gr.Button.update
+    number_update = gr.Number.update
+
+    if video_path is None:
+        return (
+            sliders_update(minimum=0, maximum=0, value=0),
+            sliders_update(minimum=1, maximum=1, value=1),
+            number_update(value=1),
+        )
+    try:
+        clip = VideoFileClip(video_path)
+        fps = clip.fps
+        total_frames = clip.reader.nframes
+        clip.close()
+        return (
+            sliders_update(minimum=0, maximum=total_frames, value=0, interactive=True),
+            sliders_update(
+                minimum=0, maximum=total_frames, value=total_frames, interactive=True
+            ),
+            number_update(value=fps),
+        )
+    except:
+        return (
+            sliders_update(value=0),
+            sliders_update(value=0),
+            number_update(value=1),
+        )
+
+
+def analyse_settings_changed(detect_condition, detection_size, detection_threshold):
+    yield "### \n ⌛ Applying new values..."
+    global FACE_ANALYSER
+    global DETECT_CONDITION
+    DETECT_CONDITION = detect_condition
+    FACE_ANALYSER = insightface.app.FaceAnalysis(name="buffalo_l", providers=PROVIDER)
+    FACE_ANALYSER.prepare(
+        ctx_id=0,
+        det_size=(int(detection_size), int(detection_size)),
+        det_thresh=float(detection_threshold),
+    )
+    yield f"### \n ✔️ Applied detect condition:{detect_condition}, detection size: {detection_size}, detection threshold: {detection_threshold}"
+
+
+def stop_running():
+    global STREAMER
+    if hasattr(STREAMER, "stop"):
+        STREAMER.stop()
+        STREAMER = None
+    return "Cancelled"
+
+
+def slider_changed(show_frame, video_path, frame_index):
+    if not show_frame:
+        return None, None
+    if video_path is None:
+        return None, None
+    clip = VideoFileClip(video_path)
+    frame = clip.get_frame(frame_index / clip.fps)
+    frame_array = np.array(frame)
+    clip.close()
+    return gr.Image.update(value=frame_array, visible=True), gr.Video.update(
+        visible=False
+    )
+
+
+def trim_and_reload(video_path, output_path, output_name, start_frame, stop_frame):
+    yield video_path, f"### \n ⌛ Trimming video frame {start_frame} to {stop_frame}..."
+    try:
+        output_path = os.path.join(output_path, output_name)
+        trimmed_video = trim_video(video_path, output_path, start_frame, stop_frame)
+        yield trimmed_video, "### \n ✔️ Video trimmed and reloaded."
+    except Exception as e:
+        print(e)
+        yield video_path, "### \n ❌ Video trimming failed. See console for more info."
+
+
+## ------------------------------ GRADIO GUI ------------------------------
+
+css = """
+footer{display:none !important}
+"""
+
+with gr.Blocks(css=css) as interface:
+    gr.Markdown("# 🗿 Swap Mukham")
+    gr.Markdown("### Face swap app based on insightface inswapper.")
+    with gr.Row():
+        with gr.Row():
+            with gr.Column(scale=0.4):
+                with gr.Tab("📄 Swap Condition"):
+                    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,
+                    )
+                    age = gr.Number(
+                        value=25, label="Value", interactive=True, visible=False
+                    )
+
+                with gr.Tab("🎚️ Detection Settings"):
+                    detect_condition_dropdown = gr.Dropdown(
+                        detect_conditions,
+                        label="Condition",
+                        value=DETECT_CONDITION,
+                        interactive=True,
+                        info="This condition is only used when multiple faces are detected on source or specific image.",
+                    )
+                    detection_size = gr.Number(
+                        label="Detection Size", value=DETECT_SIZE, interactive=True
+                    )
+                    detection_threshold = gr.Number(
+                        label="Detection Threshold",
+                        value=DETECT_THRESH,
+                        interactive=True,
+                    )
+                    apply_detection_settings = gr.Button("Apply settings")
+
+                with gr.Tab("📤 Output Settings"):
+                    output_directory = gr.Text(
+                        label="Output Directory",
+                        value=DEF_OUTPUT_PATH,
+                        interactive=True,
+                    )
+                    output_name = gr.Text(
+                        label="Output Name", value="Result", interactive=True
+                    )
+                    keep_output_sequence = gr.Checkbox(
+                        label="Keep output sequence", value=False, interactive=True
+                    )
+
+                with gr.Tab("🪄 Other Settings"):
+                    face_scale = gr.Slider(
+                        label="Face Scale",
+                        minimum=0,
+                        maximum=2,
+                        value=1,
+                        interactive=True,
+                    )
+
+                    face_enhancer_name = gr.Dropdown(
+                        FACE_ENHANCER_LIST, label="Face Enhancer", value="NONE", multiselect=False, interactive=True
+                    )
+
+                    with gr.Accordion("Advanced Mask", open=False):
+                        enable_face_parser_mask = gr.Checkbox(
+                            label="Enable Face Parsing",
+                            value=False,
+                            interactive=True,
+                        )
+
+                        mask_include = gr.Dropdown(
+                            mask_regions.keys(),
+                            value=MASK_INCLUDE,
+                            multiselect=True,
+                            label="Include",
+                            interactive=True,
+                        )
+                        mask_soft_kernel = gr.Number(
+                            label="Soft Erode Kernel",
+                            value=MASK_SOFT_KERNEL,
+                            minimum=3,
+                            interactive=True,
+                            visible = False
+                        )
+                        mask_soft_iterations = gr.Number(
+                            label="Soft Erode Iterations",
+                            value=MASK_SOFT_ITERATIONS,
+                            minimum=0,
+                            interactive=True,
+
+                        )
+
+
+                    with gr.Accordion("Crop Mask", open=False):
+                        crop_top = gr.Slider(label="Top", minimum=0, maximum=511, value=0, step=1, interactive=True)
+                        crop_bott = gr.Slider(label="Bottom", minimum=0, maximum=511, value=511, step=1, interactive=True)
+                        crop_left = gr.Slider(label="Left", minimum=0, maximum=511, value=0, step=1, interactive=True)
+                        crop_right = gr.Slider(label="Right", minimum=0, maximum=511, value=511, step=1, interactive=True)
+
+
+                    erode_amount = gr.Slider(
+                            label="Mask Erode",
+                            minimum=0,
+                            maximum=1,
+                            value=MASK_ERODE_AMOUNT,
+                            step=0.05,
+                            interactive=True,
+                        )
+
+                    blur_amount = gr.Slider(
+                            label="Mask Blur",
+                            minimum=0,
+                            maximum=1,
+                            value=MASK_BLUR_AMOUNT,
+                            step=0.05,
+                            interactive=True,
+                        )
+
+                    enable_laplacian_blend = gr.Checkbox(
+                        label="Laplacian Blending",
+                        value=True,
+                        interactive=True,
+                    )
+
+
+                source_image_input = gr.Image(
+                    label="Source face", type="filepath", interactive=True
+                )
+
+                with gr.Box(visible=False) as specific_face:
+                    for i in range(NUM_OF_SRC_SPECIFIC):
+                        idx = i + 1
+                        code = "\n"
+                        code += f"with gr.Tab(label='({idx})'):"
+                        code += "\n\twith gr.Row():"
+                        code += f"\n\t\tsrc{idx} = gr.Image(interactive=True, type='numpy', label='Source Face {idx}')"
+                        code += f"\n\t\ttrg{idx} = gr.Image(interactive=True, type='numpy', label='Specific Face {idx}')"
+                        exec(code)
+
+                    distance_slider = gr.Slider(
+                        minimum=0,
+                        maximum=2,
+                        value=0.6,
+                        interactive=True,
+                        label="Distance",
+                        info="Lower distance is more similar and higher distance is less similar to the target face.",
+                    )
+
+                with gr.Group():
+                    input_type = gr.Radio(
+                        ["Image", "Video"],
+                        label="Target Type",
+                        value="Image",
+                    )
+
+                    with gr.Box(visible=True) as input_image_group:
+                        image_input = gr.Image(
+                            label="Target Image", interactive=True, type="filepath"
+                        )
+
+                    with gr.Box(visible=False) as input_video_group:
+                        vid_widget = gr.Video if USE_COLAB else gr.Text
+                        video_input = gr.Video(
+                            label="Target Video", interactive=True
+                        )
+                        with gr.Accordion("✂️ Trim video", open=False):
+                            with gr.Column():
+                                with gr.Row():
+                                    set_slider_range_btn = gr.Button(
+                                        "Set frame range", interactive=True
+                                    )
+                                    show_trim_preview_btn = gr.Checkbox(
+                                        label="Show frame when slider change",
+                                        value=True,
+                                        interactive=True,
+                                    )
+
+                                video_fps = gr.Number(
+                                    value=30,
+                                    interactive=False,
+                                    label="Fps",
+                                    visible=False,
+                                )
+                                start_frame = gr.Slider(
+                                    minimum=0,
+                                    maximum=1,
+                                    value=0,
+                                    step=1,
+                                    interactive=True,
+                                    label="Start Frame",
+                                    info="",
+                                )
+                                end_frame = gr.Slider(
+                                    minimum=0,
+                                    maximum=1,
+                                    value=1,
+                                    step=1,
+                                    interactive=True,
+                                    label="End Frame",
+                                    info="",
+                                )
+                            trim_and_reload_btn = gr.Button(
+                                "Trim and Reload", interactive=True
+                            )
+
+                    with gr.Box(visible=False) as input_directory_group:
+                        direc_input = gr.Text(label="Path", interactive=True)
+
+            with gr.Column(scale=0.6):
+                info = gr.Markdown(value="...")
+
+                with gr.Row():
+                    swap_button = gr.Button("✨ Swap", variant="primary")
+                    cancel_button = gr.Button("⛔ Cancel")
+
+                preview_image = gr.Image(label="Output", interactive=False)
+                preview_video = gr.Video(
+                    label="Output", interactive=False, visible=False
+                )
+
+                with gr.Row():
+                    output_directory_button = gr.Button(
+                        "📂", interactive=False, visible=False
+                    )
+                    output_video_button = gr.Button(
+                        "🎬", interactive=False, visible=False
+                    )
+
+                with gr.Box():
+                    with gr.Row():
+                        gr.Markdown(
+                            "### [🤝 Sponsor](https://github.com/sponsors/harisreedhar)"
+                        )
+                        gr.Markdown(
+                            "### [👨‍💻 Source code](https://github.com/harisreedhar/Swap-Mukham)"
+                        )
+                        gr.Markdown(
+                            "### [⚠️ Disclaimer](https://github.com/harisreedhar/Swap-Mukham#disclaimer)"
+                        )
+                        gr.Markdown(
+                            "### [🌐 Run in Colab](https://colab.research.google.com/github/harisreedhar/Swap-Mukham/blob/main/swap_mukham_colab.ipynb)"
+                        )
+                        gr.Markdown(
+                            "### [🤗 Acknowledgements](https://github.com/harisreedhar/Swap-Mukham#acknowledgements)"
+                        )
+
+    ## ------------------------------ GRADIO EVENTS ------------------------------
+
+    set_slider_range_event = set_slider_range_btn.click(
+        video_changed,
+        inputs=[video_input],
+        outputs=[start_frame, end_frame, video_fps],
+    )
+
+    trim_and_reload_event = trim_and_reload_btn.click(
+        fn=trim_and_reload,
+        inputs=[video_input, output_directory, output_name, start_frame, end_frame],
+        outputs=[video_input, info],
+    )
+
+    start_frame_event = start_frame.release(
+        fn=slider_changed,
+        inputs=[show_trim_preview_btn, video_input, start_frame],
+        outputs=[preview_image, preview_video],
+        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=True,
+    )
+
+    input_type.change(
+        update_radio,
+        inputs=[input_type],
+        outputs=[input_image_group, input_video_group, input_directory_group],
+    )
+    swap_option.change(
+        swap_option_changed,
+        inputs=[swap_option],
+        outputs=[age, specific_face, source_image_input],
+    )
+
+    apply_detection_settings.click(
+        analyse_settings_changed,
+        inputs=[detect_condition_dropdown, detection_size, detection_threshold],
+        outputs=[info],
+    )
+
+    src_specific_inputs = []
+    gen_variable_txt = ",".join(
+        [f"src{i+1}" for i in range(NUM_OF_SRC_SPECIFIC)]
+        + [f"trg{i+1}" for i in range(NUM_OF_SRC_SPECIFIC)]
+    )
+    exec(f"src_specific_inputs = ({gen_variable_txt})")
+    swap_inputs = [
+        input_type,
+        image_input,
+        video_input,
+        direc_input,
+        source_image_input,
+        output_directory,
+        output_name,
+        keep_output_sequence,
+        swap_option,
+        age,
+        distance_slider,
+        face_enhancer_name,
+        enable_face_parser_mask,
+        mask_include,
+        mask_soft_kernel,
+        mask_soft_iterations,
+        blur_amount,
+        erode_amount,
+        face_scale,
+        enable_laplacian_blend,
+        crop_top,
+        crop_bott,
+        crop_left,
+        crop_right,
+        *src_specific_inputs,
+    ]
+
+    swap_outputs = [
+        info,
+        preview_image,
+        output_directory_button,
+        output_video_button,
+        preview_video,
+    ]
+
+    swap_event = swap_button.click(
+        fn=process, inputs=swap_inputs, outputs=swap_outputs, show_progress=True
+    )
+
+    cancel_button.click(
+        fn=stop_running,
+        inputs=None,
+        outputs=[info],
+        cancels=[
+            swap_event,
+            trim_and_reload_event,
+            set_slider_range_event,
+            start_frame_event,
+            end_frame_event,
+        ],
+        show_progress=True,
+    )
+    output_directory_button.click(
+        lambda: open_directory(path=WORKSPACE), inputs=None, outputs=None
+    )
+    output_video_button.click(
+        lambda: open_directory(path=OUTPUT_FILE), inputs=None, outputs=None
+    )
+
+if __name__ == "__main__":
+    if USE_COLAB:
+        print("Running in colab mode")
+
+    interface.queue(concurrency_count=2, max_size=20).launch(share=USE_COLAB)

face_analyser.py

@@ -0,0 +1,201 @@
+import os
+import cv2
+import numpy as np
+from tqdm import tqdm
+from utils import scale_bbox_from_center
+
+detect_conditions = [
+    "best detection",
+    "left most",
+    "right most",
+    "top most",
+    "bottom most",
+    "middle",
+    "biggest",
+    "smallest",
+]
+
+swap_options_list = [
+    "All Face",
+    "Specific Face",
+    "Age less than",
+    "Age greater than",
+    "All Male",
+    "All Female",
+    "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
+        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
+
+    return get_single_face(faces, method=detect_condition)
+
+
+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):
+    print("get_analysed_data")
+    print("Swap condition: {}".format(swap_condition))
+    if swap_condition != "Specific Face":
+        source_path, age = source_data
+        print("Source path: {}".format(source_path))
+        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
+    print("Total frame: {}\nCurrent IDX:{}".format(total_frames, curr_idx))
+    for curr_idx, frame_path in tqdm(enumerate(image_sequence), total=total_frames, desc="Analysing face data"):
+        print("Read frame")
+        frame = cv2.imread(frame_path)
+        print("Get frame")
+        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
+
+        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
+
+        print("Total faces: {}".format(n_faces))
+        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,72 @@
+import os
+import cv2
+import torch
+import gfpgan
+from PIL import Image
+from upscaler.RealESRGAN import RealESRGAN
+from upscaler.codeformer import CodeFormerEnhancer
+
+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
+
+
+def codeformer_runner(img, model):
+    img = model.enhance(img)
+    return img
+
+
+supported_enhancers = {
+    "CodeFormer": ("./assets/pretrained_models/codeformer.onnx", codeformer_runner),
+    "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)
+}
+
+cv2_interpolations = ["LANCZOS4", "CUBIC", "NEAREST"]
+
+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() + cv2_interpolations, f"Face enhancer {name} unavailable."
+    if name in supported_enhancers.keys():
+        model_path, model_runner = supported_enhancers.get(name)
+        model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
+    if name == 'CodeFormer':
+        model = CodeFormerEnhancer(model_path=model_path, device=device)
+    elif name == 'GFPGAN':
+        model = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=device)
+    elif name == 'REAL-ESRGAN 2x':
+        model = RealESRGAN(device, scale=2)
+        model.load_weights(model_path, download=False)
+    elif name == 'REAL-ESRGAN 4x':
+        model = RealESRGAN(device, scale=4)
+        model.load_weights(model_path, download=False)
+    elif name == 'REAL-ESRGAN 8x':
+        model = RealESRGAN(device, scale=8)
+        model.load_weights(model_path, download=False)
+    elif name == 'LANCZOS4':
+        model = None
+        model_runner = lambda img, _: cv2.resize(img, (512,512), interpolation=cv2.INTER_LANCZOS4)
+    elif name == 'CUBIC':
+        model = None
+        model_runner = lambda img, _: cv2.resize(img, (512,512), interpolation=cv2.INTER_CUBIC)
+    elif name == 'NEAREST':
+        model = None
+        model_runner = lambda img, _: cv2.resize(img, (512,512), interpolation=cv2.INTER_NEAREST)
+    else:
+        model = None
+    return (model, model_runner)

face_parsing/__init__.py

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

face_parsing/model.py

@@ -0,0 +1,283 @@
+#!/usr/bin/python
+# -*- encoding: utf-8 -*-
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+
+from .resnet import Resnet18
+# from modules.bn import InPlaceABNSync as BatchNorm2d
+
+
+class ConvBNReLU(nn.Module):
+    def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1, *args, **kwargs):
+        super(ConvBNReLU, self).__init__()
+        self.conv = nn.Conv2d(in_chan,
+                out_chan,
+                kernel_size = ks,
+                stride = stride,
+                padding = padding,
+                bias = False)
+        self.bn = nn.BatchNorm2d(out_chan)
+        self.init_weight()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = F.relu(self.bn(x))
+        return x
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+class BiSeNetOutput(nn.Module):
+    def __init__(self, in_chan, mid_chan, n_classes, *args, **kwargs):
+        super(BiSeNetOutput, self).__init__()
+        self.conv = ConvBNReLU(in_chan, mid_chan, ks=3, stride=1, padding=1)
+        self.conv_out = nn.Conv2d(mid_chan, n_classes, kernel_size=1, bias=False)
+        self.init_weight()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.conv_out(x)
+        return x
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+class AttentionRefinementModule(nn.Module):
+    def __init__(self, in_chan, out_chan, *args, **kwargs):
+        super(AttentionRefinementModule, self).__init__()
+        self.conv = ConvBNReLU(in_chan, out_chan, ks=3, stride=1, padding=1)
+        self.conv_atten = nn.Conv2d(out_chan, out_chan, kernel_size= 1, bias=False)
+        self.bn_atten = nn.BatchNorm2d(out_chan)
+        self.sigmoid_atten = nn.Sigmoid()
+        self.init_weight()
+
+    def forward(self, x):
+        feat = self.conv(x)
+        atten = F.avg_pool2d(feat, feat.size()[2:])
+        atten = self.conv_atten(atten)
+        atten = self.bn_atten(atten)
+        atten = self.sigmoid_atten(atten)
+        out = torch.mul(feat, atten)
+        return out
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+
+class ContextPath(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super(ContextPath, self).__init__()
+        self.resnet = Resnet18()
+        self.arm16 = AttentionRefinementModule(256, 128)
+        self.arm32 = AttentionRefinementModule(512, 128)
+        self.conv_head32 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+        self.conv_head16 = ConvBNReLU(128, 128, ks=3, stride=1, padding=1)
+        self.conv_avg = ConvBNReLU(512, 128, ks=1, stride=1, padding=0)
+
+        self.init_weight()
+
+    def forward(self, x):
+        H0, W0 = x.size()[2:]
+        feat8, feat16, feat32 = self.resnet(x)
+        H8, W8 = feat8.size()[2:]
+        H16, W16 = feat16.size()[2:]
+        H32, W32 = feat32.size()[2:]
+
+        avg = F.avg_pool2d(feat32, feat32.size()[2:])
+        avg = self.conv_avg(avg)
+        avg_up = F.interpolate(avg, (H32, W32), mode='nearest')
+
+        feat32_arm = self.arm32(feat32)
+        feat32_sum = feat32_arm + avg_up
+        feat32_up = F.interpolate(feat32_sum, (H16, W16), mode='nearest')
+        feat32_up = self.conv_head32(feat32_up)
+
+        feat16_arm = self.arm16(feat16)
+        feat16_sum = feat16_arm + feat32_up
+        feat16_up = F.interpolate(feat16_sum, (H8, W8), mode='nearest')
+        feat16_up = self.conv_head16(feat16_up)
+
+        return feat8, feat16_up, feat32_up  # x8, x8, x16
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, (nn.Linear, nn.Conv2d)):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+### This is not used, since I replace this with the resnet feature with the same size
+class SpatialPath(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super(SpatialPath, self).__init__()
+        self.conv1 = ConvBNReLU(3, 64, ks=7, stride=2, padding=3)
+        self.conv2 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
+        self.conv3 = ConvBNReLU(64, 64, ks=3, stride=2, padding=1)
+        self.conv_out = ConvBNReLU(64, 128, ks=1, stride=1, padding=0)
+        self.init_weight()
+
+    def forward(self, x):
+        feat = self.conv1(x)
+        feat = self.conv2(feat)
+        feat = self.conv3(feat)
+        feat = self.conv_out(feat)
+        return feat
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+class FeatureFusionModule(nn.Module):
+    def __init__(self, in_chan, out_chan, *args, **kwargs):
+        super(FeatureFusionModule, self).__init__()
+        self.convblk = ConvBNReLU(in_chan, out_chan, ks=1, stride=1, padding=0)
+        self.conv1 = nn.Conv2d(out_chan,
+                out_chan//4,
+                kernel_size = 1,
+                stride = 1,
+                padding = 0,
+                bias = False)
+        self.conv2 = nn.Conv2d(out_chan//4,
+                out_chan,
+                kernel_size = 1,
+                stride = 1,
+                padding = 0,
+                bias = False)
+        self.relu = nn.ReLU(inplace=True)
+        self.sigmoid = nn.Sigmoid()
+        self.init_weight()
+
+    def forward(self, fsp, fcp):
+        fcat = torch.cat([fsp, fcp], dim=1)
+        feat = self.convblk(fcat)
+        atten = F.avg_pool2d(feat, feat.size()[2:])
+        atten = self.conv1(atten)
+        atten = self.relu(atten)
+        atten = self.conv2(atten)
+        atten = self.sigmoid(atten)
+        feat_atten = torch.mul(feat, atten)
+        feat_out = feat_atten + feat
+        return feat_out
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module, nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+class BiSeNet(nn.Module):
+    def __init__(self, n_classes, *args, **kwargs):
+        super(BiSeNet, self).__init__()
+        self.cp = ContextPath()
+        ## here self.sp is deleted
+        self.ffm = FeatureFusionModule(256, 256)
+        self.conv_out = BiSeNetOutput(256, 256, n_classes)
+        self.conv_out16 = BiSeNetOutput(128, 64, n_classes)
+        self.conv_out32 = BiSeNetOutput(128, 64, n_classes)
+        self.init_weight()
+
+    def forward(self, x):
+        H, W = x.size()[2:]
+        feat_res8, feat_cp8, feat_cp16 = self.cp(x)  # here return res3b1 feature
+        feat_sp = feat_res8  # use res3b1 feature to replace spatial path feature
+        feat_fuse = self.ffm(feat_sp, feat_cp8)
+
+        feat_out = self.conv_out(feat_fuse)
+        feat_out16 = self.conv_out16(feat_cp8)
+        feat_out32 = self.conv_out32(feat_cp16)
+
+        feat_out = F.interpolate(feat_out, (H, W), mode='bilinear', align_corners=True)
+        feat_out16 = F.interpolate(feat_out16, (H, W), mode='bilinear', align_corners=True)
+        feat_out32 = F.interpolate(feat_out32, (H, W), mode='bilinear', align_corners=True)
+        return feat_out, feat_out16, feat_out32
+
+    def init_weight(self):
+        for ly in self.children():
+            if isinstance(ly, nn.Conv2d):
+                nn.init.kaiming_normal_(ly.weight, a=1)
+                if not ly.bias is None: nn.init.constant_(ly.bias, 0)
+
+    def get_params(self):
+        wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], []
+        for name, child in self.named_children():
+            child_wd_params, child_nowd_params = child.get_params()
+            if isinstance(child, FeatureFusionModule) or isinstance(child, BiSeNetOutput):
+                lr_mul_wd_params += child_wd_params
+                lr_mul_nowd_params += child_nowd_params
+            else:
+                wd_params += child_wd_params
+                nowd_params += child_nowd_params
+        return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params
+
+
+if __name__ == "__main__":
+    net = BiSeNet(19)
+    net.cuda()
+    net.eval()
+    in_ten = torch.randn(16, 3, 640, 480).cuda()
+    out, out16, out32 = net(in_ten)
+    print(out.shape)
+
+    net.get_params()

face_parsing/parse_mask.py

@@ -0,0 +1,107 @@
+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
+
+class SoftErosion(nn.Module):
+    def __init__(self, kernel_size=15, threshold=0.6, iterations=1):
+        super(SoftErosion, self).__init__()
+        r = kernel_size // 2
+        self.padding = r
+        self.iterations = iterations
+        self.threshold = threshold
+
+        # Create kernel
+        y_indices, x_indices = torch.meshgrid(torch.arange(0., kernel_size), torch.arange(0., kernel_size))
+        dist = torch.sqrt((x_indices - r) ** 2 + (y_indices - r) ** 2)
+        kernel = dist.max() - dist
+        kernel /= kernel.sum()
+        kernel = kernel.view(1, 1, *kernel.shape)
+        self.register_buffer('weight', kernel)
+
+    def forward(self, x):
+        batch_size = x.size(0)  # Get the batch size
+        output = []
+
+        for i in tqdm(range(batch_size), desc="Soft-Erosion", leave=False):
+            input_tensor = x[i:i+1]  # Take one input tensor from the batch
+            input_tensor = input_tensor.float()  # Convert input to float tensor
+            input_tensor = input_tensor.unsqueeze(1)  # Add a channel dimension
+
+            for _ in range(self.iterations - 1):
+                input_tensor = torch.min(input_tensor, F.conv2d(input_tensor, weight=self.weight,
+                                                                groups=input_tensor.shape[1],
+                                                                padding=self.padding))
+            input_tensor = F.conv2d(input_tensor, weight=self.weight, groups=input_tensor.shape[1],
+                                    padding=self.padding)
+
+            mask = input_tensor >= self.threshold
+            input_tensor[mask] = 1.0
+            input_tensor[~mask] /= input_tensor[~mask].max()
+
+            input_tensor = input_tensor.squeeze(1)  # Remove the extra channel dimension
+            output.append(input_tensor.detach().cpu().numpy())
+
+        return np.array(output)
+
+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,map_location=torch.device('cpu')))
+    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, softness=20):
+    if softness > 0:
+        smooth_mask = SoftErosion(kernel_size=17, threshold=0.9, iterations=softness).to(device)
+
+    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)
+        # arget_classes = torch.tensor(classes).to(device)
+        # batch_masks = torch.isin(parsing, target_classes).to(device)
+        ## torch.isin was slightly slower in my test, so using np.isin
+        parsing = out.argmax(dim=1).detach().cpu().numpy()
+        batch_masks = np.isin(parsing, classes).astype('float32')
+
+        if softness > 0:
+            # batch_masks = smooth_mask(batch_masks).transpose(1,0,2,3)[0]
+            mask_tensor = torch.from_numpy(batch_masks.copy()).float().to(device)
+            batch_masks = smooth_mask(mask_tensor).transpose(1,0,2,3)[0]
+
+        yield batch_masks
+
+        #masks.append(batch_masks)
+
+    #if len(masks) >= 1:
+    #    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/resnet.py

@@ -0,0 +1,109 @@
+#!/usr/bin/python
+# -*- encoding: utf-8 -*-
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.model_zoo as modelzoo
+
+# from modules.bn import InPlaceABNSync as BatchNorm2d
+
+resnet18_url = 'https://download.pytorch.org/models/resnet18-5c106cde.pth'
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_chan, out_chan, stride=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(in_chan, out_chan, stride)
+        self.bn1 = nn.BatchNorm2d(out_chan)
+        self.conv2 = conv3x3(out_chan, out_chan)
+        self.bn2 = nn.BatchNorm2d(out_chan)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = None
+        if in_chan != out_chan or stride != 1:
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_chan, out_chan,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_chan),
+                )
+
+    def forward(self, x):
+        residual = self.conv1(x)
+        residual = F.relu(self.bn1(residual))
+        residual = self.conv2(residual)
+        residual = self.bn2(residual)
+
+        shortcut = x
+        if self.downsample is not None:
+            shortcut = self.downsample(x)
+
+        out = shortcut + residual
+        out = self.relu(out)
+        return out
+
+
+def create_layer_basic(in_chan, out_chan, bnum, stride=1):
+    layers = [BasicBlock(in_chan, out_chan, stride=stride)]
+    for i in range(bnum-1):
+        layers.append(BasicBlock(out_chan, out_chan, stride=1))
+    return nn.Sequential(*layers)
+
+
+class Resnet18(nn.Module):
+    def __init__(self):
+        super(Resnet18, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = create_layer_basic(64, 64, bnum=2, stride=1)
+        self.layer2 = create_layer_basic(64, 128, bnum=2, stride=2)
+        self.layer3 = create_layer_basic(128, 256, bnum=2, stride=2)
+        self.layer4 = create_layer_basic(256, 512, bnum=2, stride=2)
+        self.init_weight()
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = F.relu(self.bn1(x))
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        feat8 = self.layer2(x) # 1/8
+        feat16 = self.layer3(feat8) # 1/16
+        feat32 = self.layer4(feat16) # 1/32
+        return feat8, feat16, feat32
+
+    def init_weight(self):
+        state_dict = modelzoo.load_url(resnet18_url)
+        self_state_dict = self.state_dict()
+        for k, v in state_dict.items():
+            if 'fc' in k: continue
+            self_state_dict.update({k: v})
+        self.load_state_dict(self_state_dict)
+
+    def get_params(self):
+        wd_params, nowd_params = [], []
+        for name, module in self.named_modules():
+            if isinstance(module, (nn.Linear, nn.Conv2d)):
+                wd_params.append(module.weight)
+                if not module.bias is None:
+                    nowd_params.append(module.bias)
+            elif isinstance(module,  nn.BatchNorm2d):
+                nowd_params += list(module.parameters())
+        return wd_params, nowd_params
+
+
+if __name__ == "__main__":
+    net = Resnet18()
+    x = torch.randn(16, 3, 224, 224)
+    out = net(x)
+    print(out[0].size())
+    print(out[1].size())
+    print(out[2].size())
+    net.get_params()

face_parsing/swap.py

@@ -0,0 +1,133 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+import cv2
+import numpy as np
+
+from .model import BiSeNet
+
+mask_regions = {
+    "Background":0,
+    "Skin":1,
+    "L-Eyebrow":2,
+    "R-Eyebrow":3,
+    "L-Eye":4,
+    "R-Eye":5,
+    "Eye-G":6,
+    "L-Ear":7,
+    "R-Ear":8,
+    "Ear-R":9,
+    "Nose":10,
+    "Mouth":11,
+    "U-Lip":12,
+    "L-Lip":13,
+    "Neck":14,
+    "Neck-L":15,
+    "Cloth":16,
+    "Hair":17,
+    "Hat":18
+}
+
+# Borrowed from simswap
+# https://github.com/neuralchen/SimSwap/blob/26c84d2901bd56eda4d5e3c5ca6da16e65dc82a6/util/reverse2original.py#L30
+class SoftErosion(nn.Module):
+    def __init__(self, kernel_size=15, threshold=0.6, iterations=1):
+        super(SoftErosion, self).__init__()
+        r = kernel_size // 2
+        self.padding = r
+        self.iterations = iterations
+        self.threshold = threshold
+
+        # Create kernel
+        y_indices, x_indices = torch.meshgrid(torch.arange(0., kernel_size), torch.arange(0., kernel_size))
+        dist = torch.sqrt((x_indices - r) ** 2 + (y_indices - r) ** 2)
+        kernel = dist.max() - dist
+        kernel /= kernel.sum()
+        kernel = kernel.view(1, 1, *kernel.shape)
+        self.register_buffer('weight', kernel)
+
+    def forward(self, x):
+        x = x.float()
+        for i in range(self.iterations - 1):
+            x = torch.min(x, F.conv2d(x, weight=self.weight, groups=x.shape[1], padding=self.padding))
+        x = F.conv2d(x, weight=self.weight, groups=x.shape[1], padding=self.padding)
+
+        mask = x >= self.threshold
+        x[mask] = 1.0
+        x[~mask] /= x[~mask].max()
+
+        return x, mask
+
+device = "cpu"
+
+def init_parser(pth_path, mode="cpu"):
+    global device
+    device = mode
+    n_classes = 19
+    net = BiSeNet(n_classes=n_classes)
+    if device == "cuda":
+        net.cuda()
+        net.load_state_dict(torch.load(pth_path))
+    else:
+        net.load_state_dict(torch.load(pth_path, map_location=torch.device('cpu')))
+    net.eval()
+    return net
+
+
+def image_to_parsing(img, net):
+    img = cv2.resize(img, (512, 512))
+    img = img[:,:,::-1]
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+    ])
+    img = transform(img.copy())
+    img = torch.unsqueeze(img, 0)
+
+    with torch.no_grad():
+        img = img.to(device)
+        out = net(img)[0]
+        parsing = out.squeeze(0).cpu().numpy().argmax(0)
+        return parsing
+
+
+def get_mask(parsing, classes):
+    res = parsing == classes[0]
+    for val in classes[1:]:
+        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)
+
+    if len(includes) == 0:
+        return source, np.zeros_like(source)
+
+    include_mask = get_mask(parsing, includes)
+    mask = np.repeat(include_mask[:, :, np.newaxis], 3, axis=2).astype("float32")
+
+    if smooth_mask is not None:
+        mask_tensor = torch.from_numpy(mask.copy().transpose((2, 0, 1))).float().to(device)
+        face_mask_tensor = mask_tensor[0] + mask_tensor[1]
+        soft_face_mask_tensor, _ = smooth_mask(face_mask_tensor.unsqueeze_(0).unsqueeze_(0))
+        soft_face_mask_tensor.squeeze_()
+        mask = np.repeat(soft_face_mask_tensor.cpu().numpy()[:, :, np.newaxis], 3, axis=2)
+
+    if blur > 0:
+        mask = cv2.GaussianBlur(mask, (0, 0), blur)
+
+    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
+    result = cv2.resize(result.astype("uint8"), (source.shape[1], source.shape[0]))
+
+    return result
+
+def mask_regions_to_list(values):
+    out_ids = []
+    for value in values:
+        if value in mask_regions.keys():
+            out_ids.append(mask_regions.get(value))
+    return out_ids

face_swapper.py

@@ -0,0 +1,150 @@
+import time
+import torch
+import onnx
+import cv2
+import onnxruntime
+import numpy as np
+from tqdm import tqdm
+import torch.nn as nn
+from onnx import numpy_helper
+from skimage import transform as trans
+import torchvision.transforms.functional as F
+import torch.nn.functional as F
+from utils import mask_crop, laplacian_blending
+
+
+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.session_options = onnxruntime.SessionOptions()
+        self.session = onnxruntime.InferenceSession(self.model_file, sess_options=self.session_options, providers=providers)
+
+    def forward(self, imgs, latents):
+        preds = []
+        for img, latent in zip(imgs, latents):
+            img = img / 255
+            pred = self.session.run(['output'], {'target': img, 'source': latent})[0]
+            preds.append(pred)
+
+    def get(self, imgs, target_faces, source_faces):
+        imgs = list(imgs)
+
+        preds = [None] * len(imgs)
+        matrs = [None] * len(imgs)
+
+        for idx, (img, target_face, source_face) in enumerate(zip(imgs, target_faces, source_faces)):
+            matrix, blob, latent = self.prepare_data(img, target_face, source_face)
+            pred = self.session.run(['output'], {'target': blob, 'source': latent})[0]
+            pred = pred.transpose((0, 2, 3, 1))[0]
+            pred = np.clip(255 * pred, 0, 255).astype(np.uint8)[:, :, ::-1]
+
+            preds[idx] = pred
+            matrs[idx] = matrix
+
+        return (preds, matrs)
+
+    def prepare_data(self, img, target_face, source_face):
+        if isinstance(img, str):
+            img = cv2.imread(img)
+
+        aligned_img, matrix = norm_crop2(img, target_face.kps, 128)
+
+        blob = cv2.dnn.blobFromImage(aligned_img, 1.0 / 255, (128, 128), (0., 0., 0.), swapRB=True)
+
+        latent = source_face.normed_embedding.reshape((1, -1))
+        latent = np.dot(latent, self.emap)
+        latent /= np.linalg.norm(latent)
+
+        return (matrix, blob, latent)
+
+    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
+
+        for i in tqdm(range(num_batches), desc="Generating face"):
+            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_matr = self.get(batch_img, batch_target_f, batch_source_f)
+
+            yield batch_pred, batch_matr
+
+
+def paste_to_whole(foreground, background, matrix, mask=None, crop_mask=(0,0,0,0), blur_amount=0.1, erode_amount = 0.15, blend_method='linear'):
+    inv_matrix = cv2.invertAffineTransform(matrix)
+    fg_shape = foreground.shape[:2]
+    bg_shape = (background.shape[1], background.shape[0])
+    foreground = cv2.warpAffine(foreground, inv_matrix, bg_shape, borderValue=0.0)
+
+    if mask is None:
+        mask = np.full(fg_shape, 1., dtype=np.float32)
+        mask = mask_crop(mask, crop_mask)
+        mask = cv2.warpAffine(mask, inv_matrix, bg_shape, borderValue=0.0)
+    else:
+        assert fg_shape == mask.shape[:2], "foreground & mask shape mismatch!"
+        mask = mask_crop(mask, crop_mask).astype('float32')
+        mask = cv2.warpAffine(mask, inv_matrix, (background.shape[1], background.shape[0]), borderValue=0.0)
+
+    _mask = mask.copy()
+    _mask[_mask > 0.05] = 1.
+    non_zero_points = cv2.findNonZero(_mask)
+    _, _, w, h = cv2.boundingRect(non_zero_points)
+    mask_size = int(np.sqrt(w * h))
+
+    if erode_amount > 0:
+        kernel_size = max(int(mask_size * erode_amount), 1)
+        structuring_element = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size))
+        mask = cv2.erode(mask, structuring_element)
+
+    if blur_amount > 0:
+        kernel_size = max(int(mask_size * blur_amount), 3)
+        if kernel_size % 2 == 0:
+            kernel_size += 1
+        mask = cv2.GaussianBlur(mask, (kernel_size, kernel_size), 0)
+
+    mask = np.tile(np.expand_dims(mask, axis=-1), (1, 1, 3))
+
+    if blend_method == 'laplacian':
+        composite_image = laplacian_blending(foreground, background, mask.clip(0,1), num_levels=4)
+    else:
+        composite_image = mask * foreground + (1 - mask) * background
+
+    return composite_image.astype("uint8").clip(0, 255)

requirements.txt

@@ -0,0 +1,22 @@
+--extra-index-url https://download.pytorch.org/whl/cu118
+
+gfpgan==1.3.8
+gradio==3.40.1
+insightface==0.7.3
+moviepy>=1.0.3
+numpy==1.24.3
+onnx==1.14.0
+onnxruntime==1.15.1; python_version != '3.9' and sys_platform == 'darwin' and platform_machine != 'arm64'
+onnxruntime-coreml==1.13.1; python_version == '3.9' and sys_platform == 'darwin' and platform_machine != 'arm64'
+onnxruntime-gpu==1.15.1; sys_platform != 'darwin'
+onnxruntime-silicon==1.13.1; sys_platform == 'darwin' and platform_machine == 'arm64'
+opencv-python==4.8.0.74
+opennsfw2==0.10.2
+pillow==10.0.0
+protobuf==4.23.4
+psutil==5.9.5
+realesrgan==0.3.0
+tensorflow==2.13.0
+tqdm==4.65.0
+python-telegram-bot==22.1
+

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], :]

upscaler/codeformer.py

@@ -0,0 +1,37 @@
+import cv2
+import torch
+import onnx
+import onnxruntime
+import numpy as np
+
+import time
+
+# codeformer converted to onnx
+# using https://github.com/redthing1/CodeFormer
+
+
+class CodeFormerEnhancer:
+    def __init__(self, model_path="codeformer.onnx", device='cpu'):
+        model = onnx.load(model_path)
+        session_options = onnxruntime.SessionOptions()
+        session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+        providers = ["CPUExecutionProvider"]
+        if device == 'cuda':
+            providers = [("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),"CPUExecutionProvider"]
+        self.session = onnxruntime.InferenceSession(model_path, sess_options=session_options, providers=providers)
+
+    def enhance(self, img, w=0.9):
+        img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
+        img = img.astype(np.float32)[:,:,::-1] / 255.0
+        img = img.transpose((2, 0, 1))
+        nrm_mean = np.array([0.5, 0.5, 0.5]).reshape((-1, 1, 1))
+        nrm_std = np.array([0.5, 0.5, 0.5]).reshape((-1, 1, 1))
+        img = (img - nrm_mean) / nrm_std
+
+        img = np.expand_dims(img, axis=0)
+
+        out = self.session.run(None, {'x':img.astype(np.float32), 'w':np.array([w], dtype=np.double)})[0]
+        out = (out[0].transpose(1,2,0).clip(-1,1) + 1) * 0.5
+        out = (out * 255)[:,:,::-1]
+
+        return out.astype('uint8')

utils.py

@@ -0,0 +1,303 @@
+import os
+import cv2
+import time
+import glob
+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
+
+
+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):
+    video_name, _ = os.path.splitext(os.path.basename(video_path))
+    trimmed_video_filename = video_name + "_trimmed" + ".mp4"
+    temp_path = os.path.join(output_path, "trim")
+    os.makedirs(temp_path, exist_ok=True)
+    trimmed_video_file_path = os.path.join(temp_path, trimmed_video_filename)
+
+    video = VideoFileClip(video_path, fps_source="fps")
+    fps = video.fps
+    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", bitrate=bitrate,
+    )
+    trimmed_video.close()
+    video.close()
+
+    return trimmed_video_file_path
+
+
+def open_directory(path=None):
+    if path is None:
+        return
+    try:
+        os.startfile(path)
+    except:
+        subprocess.Popen(["xdg-open", path])
+
+
+class StreamerThread(object):
+    def __init__(self, src=0):
+        self.capture = cv2.VideoCapture(src)
+        self.capture.set(cv2.CAP_PROP_BUFFERSIZE, 2)
+        self.FPS = 1 / 30
+        self.FPS_MS = int(self.FPS * 1000)
+        self.thread = None
+        self.stopped = False
+        self.frame = None
+
+    def start(self):
+        self.thread = Thread(target=self.update, args=())
+        self.thread.daemon = True
+        self.thread.start()
+
+    def stop(self):
+        self.stopped = True
+        self.thread.join()
+        print("stopped")
+
+    def update(self):
+        while not self.stopped:
+            if self.capture.isOpened():
+                (self.status, self.frame) = self.capture.read()
+            time.sleep(self.FPS)
+
+
+class ProcessBar:
+    def __init__(self, bar_length, total, before="⬛", after="🟨"):
+        self.bar_length = bar_length
+        self.total = total
+        self.before = before
+        self.after = after
+        self.bar = [self.before] * bar_length
+        self.start_time = time.time()
+
+    def get(self, index):
+        total = self.total
+        elapsed_time = time.time() - self.start_time
+        average_time_per_iteration = elapsed_time / (index + 1)
+        remaining_iterations = total - (index + 1)
+        estimated_remaining_time = remaining_iterations * average_time_per_iteration
+
+        self.bar[int(index / total * self.bar_length)] = self.after
+        info_text = f"({index+1}/{total}) {''.join(self.bar)} "
+        info_text += f"(ETR: {int(estimated_remaining_time // 60)} min {int(estimated_remaining_time % 60)} sec)"
+        return info_text
+
+
+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))
+    if logo.shape[2] == 4:
+        alpha = logo[:, :, 3]
+    else:
+        alpha = np.ones_like(logo[:, :, 0]) * 255
+    padding = int(logo_size * 0.1)
+    roi = img.shape[0] - logo_size - padding, img.shape[1] - logo_size - padding
+    for c in range(0, 3):
+        img[roi[0] : roi[0] + logo_size, roi[1] : roi[1] + logo_size, c] = (
+            alpha / 255.0
+        ) * logo[:, :, c] + (1 - alpha / 255.0) * img[
+            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_source="fps")
+    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)
+
+    bitrate = get_bitrate_for_resolution(min(*edited_video_clip.size), "high")
+
+    edited_video_clip.set_duration(duration).write_videofile(
+        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
+
+    # 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
+
+
+def laplacian_blending(A, B, m, num_levels=7):
+    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, 8192])
+    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_ = ls_[:height, :width, :]
+    #ls_ = (ls_ - np.min(ls_)) * (255.0 / (np.max(ls_) - np.min(ls_)))
+    return ls_.clip(0, 255)
+
+
+def mask_crop(mask, crop):
+    top, bottom, left, right = crop
+    shape = mask.shape
+    top = int(top)
+    bottom = int(bottom)
+    if top + bottom < shape[1]:
+        if top > 0: mask[:top, :] = 0
+        if bottom > 0: mask[-bottom:, :] = 0
+
+    left = int(left)
+    right = int(right)
+    if left + right < shape[0]:
+        if left > 0: mask[:, :left] = 0
+        if right > 0: mask[:, -right:] = 0
+
+    return mask
+
+def create_image_grid(images, size=128):
+    num_images = len(images)
+    num_cols = int(np.ceil(np.sqrt(num_images)))
+    num_rows = int(np.ceil(num_images / num_cols))
+    grid = np.zeros((num_rows * size, num_cols * size, 3), dtype=np.uint8)
+
+    for i, image in enumerate(images):
+        row_idx = (i // num_cols) * size
+        col_idx = (i % num_cols) * size
+        image = cv2.resize(image.copy(), (size,size))
+        if image.dtype != np.uint8:
+            image = (image.astype('float32') * 255).astype('uint8')
+        if image.ndim == 2:
+            image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+        grid[row_idx:row_idx + size, col_idx:col_idx + size] = image
+
+    return grid