Duplicate from bluefoxcreation/SwapMukham

Co-authored-by: BlueFox <[email protected]>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+
+*.pyc

README.md

@@ -0,0 +1,14 @@
+---
+title: Swap Mukham
+emoji: 💻
+colorFrom: purple
+colorTo: gray
+sdk: gradio
+sdk_version: 3.35.2
+app_file: app.py
+pinned: false
+license: unknown
+duplicated_from: bluefoxcreation/SwapMukham
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,924 @@
+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 nsfw_checker import NSFWChecker
+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
+NSFW_DETECTOR = None
+FACE_ENHANCER_LIST = ["NONE"]
+FACE_ENHANCER_LIST.extend(get_available_enhancer_names())
+FACE_ENHANCER_LIST.extend(cv2_interpolations)
+
+## ------------------------------ 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)
+
+def load_nsfw_detector_model(path="./assets/pretrained_models/open-nsfw.onnx"):
+    global NSFW_DETECTOR
+    if NSFW_DETECTOR is None:
+        NSFW_DETECTOR = NSFWChecker(model_path=path, providers=PROVIDER)
+
+
+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 NSFW detector model...", *ui_before()
+    load_nsfw_detector_model()
+
+    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 ------------------------------
+
+        yield "### \n ⌛ Checking contents...", *ui_before()
+        nsfw = NSFW_DETECTOR.is_nsfw(image_sequence)
+        if nsfw:
+            message = "NSFW Content detected !!!"
+            yield f"### \n 🔞 {message}", *ui_before()
+            assert not nsfw, message
+            return False
+        EMPTY_CACHE()
+
+        ## ------------------------------ ANALYSE FACE ------------------------------
+
+        yield "### \n ⌛ Analysing face data...", *ui_before()
+        if condition != "Specific Face":
+            source_data = source_path, age
+        else:
+            source_data = ((sources, specifics), distance)
+        analysed_targets, analysed_sources, whole_frame_list, num_faces_per_frame = get_analysed_data(
+            FACE_ANALYSER,
+            image_sequence,
+            source_data,
+            swap_condition=condition,
+            detect_condition=DETECT_CONDITION,
+            scale=face_scale
+        )
+
+        ## ------------------------------ SWAP FUNC ------------------------------
+
+        yield "### \n ⌛ Generating faces...", *ui_before()
+        preds = []
+        matrs = []
+        count = 0
+        global PREVIEW
+        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
+
+            if USE_CUDA:
+                image_grid = create_image_grid(batch_pred, size=128)
+                PREVIEW = image_grid[:, :, ::-1]
+                yield f"### \n ⌛ Generating face Batch {count}", *ui_before()
+
+        ## ------------------------------ FACE ENHANCEMENT ------------------------------
+
+        generated_len = len(preds)
+        if face_enhancer_name != "NONE":
+            yield f"### \n ⌛ Upscaling faces with {face_enhancer_name}...", *ui_before()
+            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:
+            yield "### \n ⌛ Face-parsing mask...", *ui_before()
+            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
+
+                if len(batch_mask) > 1:
+                    image_grid = create_image_grid(batch_mask, size=128)
+                    PREVIEW = image_grid[:, :, ::-1]
+                    yield f"### \n ⌛ Face parsing Batch {count}", *ui_before()
+            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 ------------------------------
+
+        yield "### \n ⌛ Pasting back...", *ui_before()
+        def post_process(frame_idx, frame_img, split_preds, split_matrs, split_masks, enable_laplacian_blend, crop_mask, blur_amount, erode_amount):
+            whole_img_path = frame_img
+            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)
+
+        def concurrent_post_process(image_sequence, *args):
+            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
+        )
+
+
+    ## ------------------------------ 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)
+
+        yield "### \n ⌛ Extracting video frames...", *ui_before()
+        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
+        cap.release()
+        cv2.destroyAllWindows()
+
+        for info_update in swap_process(image_sequence):
+            yield info_update
+
+        yield "### \n ⌛ Merging sequence...", *ui_before()
+        output_video_path = os.path.join(output_path, output_name + ".mp4")
+        merge_img_sequence_from_ref(video_path, image_sequence, output_video_path)
+
+        if os.path.exists(temp_path) and not keep_output_sequence:
+            yield "### \n ⌛ Removing temporary files...", *ui_before()
+            shutil.rmtree(temp_path)
+
+        WORKSPACE = output_path
+        OUTPUT_FILE = output_video_path
+
+        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)

assets/pretrained_models/79999_iter.pth

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

assets/pretrained_models/GFPGANv1.4.pth

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

assets/pretrained_models/RealESRGAN_x2.pth

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

assets/pretrained_models/RealESRGAN_x4.pth

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

assets/pretrained_models/RealESRGAN_x8.pth

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

assets/pretrained_models/codeformer.onnx

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

assets/pretrained_models/inswapper_128.onnx

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

assets/pretrained_models/open-nsfw.onnx

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

assets/pretrained_models/readme.md

@@ -0,0 +1,4 @@
+## Downolad these models here
+- [inswapper_128.onnx](https://huggingface.co/deepinsight/inswapper/resolve/main/inswapper_128.onnx)
+- [GFPGANv1.4.pth](https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth)
+- [79999_iter.pth](https://drive.google.com/open?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812)

face_analyser.py

@@ -0,0 +1,194 @@
+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):
+    if swap_condition != "Specific Face":
+        source_path, age = source_data
+        source_image = cv2.imread(source_path)
+        analysed_source = analyse_face(source_image, face_analyser, return_single_face=True, detect_condition=detect_condition, scale=scale)
+    else:
+        analysed_source_specifics = []
+        source_specifics, threshold = source_data
+        for source, specific in zip(*source_specifics):
+            if source is None or specific is None:
+                continue
+            analysed_source = analyse_face(source, face_analyser, return_single_face=True, detect_condition=detect_condition, scale=scale)
+            analysed_specific = analyse_face(specific, face_analyser, return_single_face=True, detect_condition=detect_condition, scale=scale)
+            analysed_source_specifics.append([analysed_source, analysed_specific])
+
+    analysed_target_list = []
+    analysed_source_list = []
+    whole_frame_eql_list = []
+    num_faces_per_frame = []
+
+    total_frames = len(image_sequence)
+    curr_idx = 0
+    for curr_idx, frame_path in tqdm(enumerate(image_sequence), total=total_frames, desc="Analysing face data"):
+        frame = cv2.imread(frame_path)
+        analysed_faces = analyse_face(frame, face_analyser, return_single_face=False, detect_condition=detect_condition, scale=scale)
+
+        n_faces = 0
+        for analysed_face in analysed_faces:
+            if swap_condition == "All Face":
+                analysed_target_list.append(analysed_face)
+                analysed_source_list.append(analysed_source)
+                whole_frame_eql_list.append(frame_path)
+                n_faces += 1
+            elif swap_condition == "Age less than" and analysed_face["age"] < age:
+                analysed_target_list.append(analysed_face)
+                analysed_source_list.append(analysed_source)
+                whole_frame_eql_list.append(frame_path)
+                n_faces += 1
+            elif swap_condition == "Age greater than" and analysed_face["age"] > age:
+                analysed_target_list.append(analysed_face)
+                analysed_source_list.append(analysed_source)
+                whole_frame_eql_list.append(frame_path)
+                n_faces += 1
+            elif swap_condition == "All Male" and analysed_face["gender"] == 1:
+                analysed_target_list.append(analysed_face)
+                analysed_source_list.append(analysed_source)
+                whole_frame_eql_list.append(frame_path)
+                n_faces += 1
+            elif swap_condition == "All Female" and analysed_face["gender"] == 0:
+                analysed_target_list.append(analysed_face)
+                analysed_source_list.append(analysed_source)
+                whole_frame_eql_list.append(frame_path)
+                n_faces += 1
+            elif swap_condition == "Specific Face":
+                for analysed_source, analysed_specific in analysed_source_specifics:
+                    distance = cosine_distance(analysed_specific["embedding"], analysed_face["embedding"])
+                    if distance < threshold:
+                        analysed_target_list.append(analysed_face)
+                        analysed_source_list.append(analysed_source)
+                        whole_frame_eql_list.append(frame_path)
+                        n_faces += 1
+
+        if swap_condition == "Left Most":
+            analysed_face = get_single_face(analysed_faces, method="left most")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
+        elif swap_condition == "Right Most":
+            analysed_face = get_single_face(analysed_faces, method="right most")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
+        elif swap_condition == "Top Most":
+            analysed_face = get_single_face(analysed_faces, method="top most")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
+        elif swap_condition == "Bottom Most":
+            analysed_face = get_single_face(analysed_faces, method="bottom most")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
+        elif swap_condition == "Middle":
+            analysed_face = get_single_face(analysed_faces, method="middle")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
+        elif swap_condition == "Biggest":
+            analysed_face = get_single_face(analysed_faces, method="biggest")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
+        elif swap_condition == "Smallest":
+            analysed_face = get_single_face(analysed_faces, method="smallest")
+            analysed_target_list.append(analysed_face)
+            analysed_source_list.append(analysed_source)
+            whole_frame_eql_list.append(frame_path)
+            n_faces += 1
+
+        num_faces_per_frame.append(n_faces)
+
+    return analysed_target_list, analysed_source_list, whole_frame_eql_list, num_faces_per_frame

face_enhancer.py

@@ -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))
+    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)

gfpgan/weights/detection_Resnet50_Final.pth

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

gfpgan/weights/parsing_parsenet.pth

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

nsfw_checker/LICENSE.md

@@ -0,0 +1,11 @@
+
+Copyright 2016, Yahoo Inc.
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+

nsfw_checker/__init__.py

@@ -0,0 +1 @@
+from . opennsfw import NSFWChecker

nsfw_checker/opennsfw.py

@@ -0,0 +1,37 @@
+import cv2
+import torch
+import onnx
+import onnxruntime
+import numpy as np
+from tqdm import tqdm
+
+# https://github.com/yahoo/open_nsfw
+
+class NSFWChecker:
+    def __init__(self, model_path=None, providers=["CPUExecutionProvider"]):
+        model = onnx.load(model_path)
+        self.input_name = model.graph.input[0].name
+        session_options = onnxruntime.SessionOptions()
+        self.session = onnxruntime.InferenceSession(model_path, sess_options=session_options, providers=providers)
+
+    def is_nsfw(self, img_paths, threshold = 0.85):
+        skip_step = 1
+        total_len = len(img_paths)
+        if total_len < 100: skip_step = 1
+        if total_len > 100 and total_len < 500: skip_step = 10
+        if total_len > 500 and total_len < 1000: skip_step = 20
+        if total_len > 1000 and total_len < 10000: skip_step = 50
+        if total_len > 10000: skip_step = 100
+
+        for idx in tqdm(range(0, total_len, skip_step), total=int(total_len // skip_step), desc="Checking for NSFW contents"):
+            img = cv2.imread(img_paths[idx])
+            img = cv2.resize(img, (224,224)).astype('float32')
+            img -= np.array([104, 117, 123], dtype=np.float32)
+            img = np.expand_dims(img, axis=0)
+
+            score = self.session.run(None, {self.input_name:img})[0][0][1]
+
+            if score > threshold:
+                print(f"Detected nsfw score:{score}")
+                return True
+        return False

nsfw_detector.py

@@ -0,0 +1,65 @@
+from torchvision.transforms import Normalize
+import torchvision.transforms as T
+import torch.nn as nn
+from PIL import Image
+import numpy as np
+import torch
+import timm
+from tqdm import tqdm
+
+# https://github.com/Whiax/NSFW-Classifier/raw/main/nsfwmodel_281.pth
+normalize_t = Normalize((0.4814, 0.4578, 0.4082), (0.2686, 0.2613, 0.2757))
+
+#nsfw classifier
+class NSFWClassifier(nn.Module):
+    def __init__(self):
+        super().__init__()
+        nsfw_model=self
+        nsfw_model.root_model = timm.create_model('convnext_base_in22ft1k', pretrained=True)
+        nsfw_model.linear_probe = nn.Linear(1024, 1, bias=False)
+
+    def forward(self, x):
+        nsfw_model = self
+        x = normalize_t(x)
+        x = nsfw_model.root_model.stem(x)
+        x = nsfw_model.root_model.stages(x)
+        x = nsfw_model.root_model.head.global_pool(x)
+        x = nsfw_model.root_model.head.norm(x)
+        x = nsfw_model.root_model.head.flatten(x)
+        x = nsfw_model.linear_probe(x)
+        return x
+
+    def is_nsfw(self, img_paths, threshold = 0.98):
+        skip_step = 1
+        total_len = len(img_paths)
+        if total_len < 100: skip_step = 1
+        if total_len > 100 and total_len < 500: skip_step = 10
+        if total_len > 500 and total_len < 1000: skip_step = 20
+        if total_len > 1000 and total_len < 10000: skip_step = 50
+        if total_len > 10000: skip_step = 100
+
+        for idx in tqdm(range(0, total_len, skip_step), total=int(total_len // skip_step), desc="Checking for NSFW contents"):
+            _img = Image.open(img_paths[idx]).convert('RGB')
+            img = _img.resize((224, 224))
+            img = np.array(img)/255
+            img = T.ToTensor()(img).unsqueeze(0).float()
+            if next(self.parameters()).is_cuda:
+                img = img.cuda()
+            with torch.no_grad():
+                score = self.forward(img).sigmoid()[0].item()
+            if score > threshold:
+                print(f"Detected nsfw score:{score}")
+                _img.save("nsfw.jpg")
+                return True
+        return False
+
+def get_nsfw_detector(model_path='nsfwmodel_281.pth', device="cpu"):
+    #load base model
+    nsfw_model = NSFWClassifier()
+    nsfw_model = nsfw_model.eval()
+    #load linear weights
+    linear_pth = model_path
+    linear_state_dict = torch.load(linear_pth, map_location='cpu')
+    nsfw_model.linear_probe.load_state_dict(linear_state_dict)
+    nsfw_model = nsfw_model.to(device)
+    return nsfw_model

requirements.txt

@@ -0,0 +1,12 @@
+torch==2.0.1
+torchvision
+gradio>=3.33.1
+insightface==0.7.3
+moviepy>=1.0.3
+numpy
+onnx==1.14.0
+onnxruntime==1.15.0
+opencv-python>=4.7.0.72
+opencv-python-headless>=4.7.0.72
+gfpgan==1.3.8
+

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