app.py

import numpy as np
import cv2
import os
import insightface 
from insightface.app import FaceAnalysis
from insightface.data import get_image as ins_get_image

import gradio as gr

theme = gr.themes.Default(
    font=['Helvetica', 'ui-sans-serif', 'system-ui', 'sans-serif'],
    font_mono=['IBM Plex Mono', 'ui-monospace', 'Consolas', 'monospace'],
).set(
    border_color_primary='#c5c5d2',
    button_large_padding='6px 12px',
    body_text_color_subdued='#484848',
    background_fill_secondary='#eaeaea'
)

def add_bbox_padding(bbox, margin=5):
    return [
        bbox[0] - margin,
        bbox[1] - margin,
        bbox[2] + margin,
        bbox[3] + margin]


def select_handler(img, evt: gr.SelectData): 
    faces = app.get(img)
    faces = sorted(faces, key = lambda x : x.bbox[0])
    cropped_image = []
    face_index = -1
    sel_face_index = 0
    print("Coords: ", evt.index[0],evt.index[1])
    for face in faces:
        box = face.bbox.astype(np.int32)
        face_index = face_index + 1 
        if point_in_box((box[0], box[1]),(box[2],box[3]),(evt.index[0],evt.index[1])) == True:
            # print("True ", face_index)
            # print("Bbox org: ", box)
            # Add ~25% margin to the box so the face is recognized correctly
            margin = int((box[2]-box[0]) * 0.35)
            box = add_bbox_padding(box,margin)
            box = np.clip(box,0,None)
            print("Bbox exp: ", box)
            sel_face_index = face_index            
            cropped_image = img[box[1]:box[3],box[0]:box[2]]
    return cropped_image, sel_face_index

def point_in_box(bl, tr, p) :
   if (p[0] > bl[0] and p[0] < tr[0] and p[1] > bl[1] and p[1] < tr[1]) :
      return True
   else:
      return False
   
def get_faces(img):
    faces = app.get(img)
    faces = sorted(faces, key = lambda x : x.bbox[0])
    #boxed_faces = app.draw_on(img, faces)
    #for i in range(len(faces)):
    #    face = faces[i]
    #    box = face.bbox.astype(np.int32)
    #    cv2.putText(boxed_faces,'Face#:%d'%(i), (box[0]-1, box[3]+14),cv2.FONT_HERSHEY_COMPLEX,0.7,(0,0,255),2)

    return img, len(faces)

def swap_face_fct(img_source,face_index,img_swap_face):
    faces = app.get(img_source)
    faces = sorted(faces, key = lambda x : x.bbox[0])
    src_face = app.get(img_swap_face)
    src_face = sorted(src_face, key = lambda x : x.bbox[0])
    #print("index:",faces)
    res = swapper.get(img_source, faces[face_index], src_face[0], paste_back=True)
    return res

def swap_video_fct(video_path, output_path, source_face, destination_face, tolerance, preview=-1, progress=gr.Progress()):

    # Get the Destination Face parameters (the face which should be swapped)
    dest_face = app.get(destination_face)
    dest_face = sorted(dest_face, key = lambda x : x.bbox[0])

    if(len(dest_face) == 0):
        print("💡 No dest face found")
        return -1
 
    dest_face_feats = []
    dest_face_feats.append(dest_face[0].normed_embedding)
    dest_face_feats = np.array(dest_face_feats, dtype=np.float32)

    # Get the source face parameters (the face that replaces the original)
    src_face = app.get(source_face)
    src_face = sorted(src_face, key = lambda x : x.bbox[0])
    if(len(src_face) == 0):
        print("🚨 No source face found")
        return -1
    
    cap = cv2.VideoCapture(video_path)
    ret, frame = cap.read()
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fourcc = cv2.VideoWriter_fourcc(*'avc1')

    # Use the same tmp dir from gradio if no output path is set
    if(len(output_path) > 0):
        out_path = output_path
    else:
        out_path = os.path.dirname(video_path) + "/out.mp4"

    if preview == -1:
        for_range = range(frame_count)
        video_out = cv2.VideoWriter(out_path,fourcc,fps,(width,height))
    else:
        for_range = range(preview-1,preview)

    for i in for_range:
        progress(i/frame_count, desc="⏳Processing")
        cap.set(cv2.CAP_PROP_POS_FRAMES, i)
        ret, frame = cap.read()
        #frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        # Find all faces in the current frame
        faces = app.get(frame)
        faces = sorted(faces, key = lambda x : x.bbox[0])
        # No face in Scene => copy input frame

        if(len(faces) > 0):
            feats = []
            for face in faces:
                feats.append(face.normed_embedding)
            feats = np.array(feats, dtype=np.float32)
            sims = np.dot(dest_face_feats, feats.T)
            print(sims)
            # find the index of the most similar face
            max_index = np.argmax(sims)
            print("Sim:", max_index)
            if(sims[0][max_index]*100 >= (100-tolerance)):
                frame = swapper.get(frame, faces[max_index], src_face[0], paste_back=True)
        if preview == -1:
            video_out.write(frame)
    if preview == -1:
        video_out.release()
        return out_path
    else:
        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
ins_get_image

def analyze_video(video_path):
    cap = cv2.VideoCapture(video_path)
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    length = frame_count/fps
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    return f"Resolution: {width}x{height}\nLength: {length}\nFps: {fps}\nFrames: {frame_count}"
    
def update_slider(video_path):
    cap = cv2.VideoCapture(video_path)
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    length = frame_count/fps
    return gr.update(minimum=0,maximum=frame_count,value=frame_count/2)
                     
def show_preview(video_path, frame_number):
    cap = cv2.VideoCapture(video_path)
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
    ret, frame = cap.read()
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    return frame

def create_interface():
    title = '# 🧸FaceSwap UI'
    with gr.Blocks(theme='WeixuanYuan/Base_dark') as face_swap_ui:
        with gr.Tab("📀 Swap Face Image"):
            with gr.Row():
                with gr.Column():
                    image_input = gr.Image(label='🎞️Input Image (📺 Click to select a face)', scale=0.5)
                    with gr.Row():
                        analyze_button = gr.Button("⌛ Analyze")
                    with gr.Row():
                        with gr.Column():
                            face_num = gr.Number(label='📹 Recognized Faces')
                            face_index_num = gr.Number(label='📀 Face Index', precision=0)
                        selected_face = gr.Image(label='💊Face to swap', interactive=False)
                        swap_face = gr.Image(label='🧪Swap Face')
                    swap_button = gr.Button("🧩 Swap")
                with gr.Column():
                    image_output = gr.Image(label='📤Output Image',interactive=False)
                    #text_output = gr.Textbox(placeholder="What is your name?")
            swap_button.click(fn=swap_face_fct, inputs=[image_input, face_index_num, swap_face], outputs=[image_output])
            image_input.select(select_handler, image_input, [selected_face, face_index_num])
            analyze_button.click(fn=get_faces, inputs=image_input, outputs=[image_input,face_num])
        with gr.Tab("📺Swap Face Video"):
            with gr.Row():
                with gr.Column():
                    source_video = gr.Video()
                    video_info = gr.Textbox(label="📡Video Information")
                    gr.Markdown("🛠️Select a frame for preview with the slider. Then select the face which should be swapped by clicking on it with the cursor")
                    video_position = gr.Slider(label="📐Frame preview",interactive=True)
                    frame_preview = gr.Image(label="📏Frame preview")
                    face_index = gr.Textbox(label="📉Face-Index",interactive=False)
                    with gr.Row():
                        dest_face_vid = gr.Image(label="👑Face to swap",interactive=True)
                        source_face_vid = gr.Image(label="🔮New Face")
                    gr.Markdown("🔑The higher the tolerance the more likely a wrong face will be swapped. 30-40 is a good starting point.")
                    face_tolerance = gr.Slider(label="⏳Tolerance",value=40,interactive=True)
                    preview_video = gr.Button("🪞Preview")
                    video_file_path = gr.Text(label="🗳️Output Video path incl. file.mp4 (when left empty it will be put in the gradio temp dir)")
                    process_video = gr.Button("⌛Process")
                with gr.Column():
                    with gr.Column(scale=1):
                        image_output = gr.Image()
                        output_video = gr.Video(interactive=False)
                    with gr.Column(scale=1):
                        pass
            # Component Events
            source_video.upload(fn=analyze_video,inputs=source_video,outputs=video_info)
            video_info.change(fn=update_slider,inputs=source_video,outputs=video_position)
            #preview_button.click(fn=show_preview,inputs=[source_video, video_position],outputs=frame_preview)
            frame_preview.select(select_handler, frame_preview, [dest_face_vid, face_index ])
            video_position.change(show_preview,inputs=[source_video, video_position],outputs=frame_preview)
            process_video.click(fn=swap_video_fct,inputs=[source_video,video_file_path,source_face_vid,dest_face_vid, face_tolerance], outputs=output_video)
            preview_video.click(fn=swap_video_fct,inputs=[source_video,video_file_path,source_face_vid,dest_face_vid, face_tolerance, video_position], outputs=image_output)
        
        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__":

    app = FaceAnalysis(name='buffalo_l')
    app.prepare(ctx_id=0, det_size=(640, 640))
    swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=True, download_zip=True)

    create_interface()