app.py

from PIL import Image, ImageOps, ImageEnhance
import mediapipe as mp
import numpy as np
import cv2
import gradio as gr
import imageio


class FaceSwapper:
    def __init__(self):
        self.face_mesh = mp.solutions.face_mesh.FaceMesh(
            static_image_mode=True,
            max_num_faces=1,
            refine_landmarks=True,
            min_detection_confidence=0.5
        )
        self.landmarks_a = None
        self.landmarks_b = None
        self.image_a = None
        self.image_b = None
        self.image_a_origin = None
        self.image_b_origin = None

    def get_face_landmarks(self, image):
        results = self.face_mesh.process(cv2.cvtColor(np.array(image), cv2.COLOR_BGR2RGB))
        if not results.multi_face_landmarks:
            return None
        landmarks = results.multi_face_landmarks[0].landmark
        w, h = image.size
        face_landmarks = [(int(p.x * w), int(p.y * h)) for p in landmarks]
        return face_landmarks

    def load_images(self, image_a, image_b):
        if image_a:
            if image_a.convert("RGBA") != self.image_a_origin:
                self.image_a = image_a.convert("RGBA")
                self.image_a_origin = image_a.convert("RGBA")
                self.landmarks_a = self.get_face_landmarks(self.image_a)
        if image_b:
            if image_b.convert("RGBA") != self.image_a_origin:
                self.image_b = image_b.convert("RGBA")
                self.image_b_origin = image_b.convert("RGBA")
                self.landmarks_b = self.get_face_landmarks(self.image_b)
        if not self.landmarks_a or not self.landmarks_b:
            return None

    def overlay_image(self, background, overlay, position):
        x, y = position

        if background.mode != 'RGBA':
            background = background.convert('RGBA')
        if overlay.mode != 'RGBA':
            overlay = overlay.convert('RGBA')

        overlay_mask = overlay.split()[3]  

        background.paste(overlay, (x, y), overlay_mask)
        return background

    def swap_faces(self, scale_factor=1.0, offset_x=0, offset_y=0, rotation_angle=0, margin_left=0.1, margin_right=0.1,
                   margin_top=0.1, margin_bottom=0.1, mirror_face_a=False):
        if not self.landmarks_a or not self.image_a or not self.image_b:
            raise ValueError("Image Load Error Or Face Not Deteceted")


        min_x_a = min([p[0] for p in self.landmarks_a])
        max_x_a = max([p[0] for p in self.landmarks_a])
        min_y_a = min([p[1] for p in self.landmarks_a])
        max_y_a = max([p[1] for p in self.landmarks_a])

        width_a = max_x_a - min_x_a
        height_a = max_y_a - min_y_a


        face_a = self.image_a.crop((
            int(min_x_a - width_a * margin_left),
            int(min_y_a - height_a * 0.85 * margin_top),
            int(max_x_a + width_a * margin_right),
            int(max_y_a + height_a * margin_bottom)
        ))

        if mirror_face_a:
            face_a = ImageOps.mirror(face_a)


        if not self.landmarks_b:

            width_b, height_b = self.image_b.size
            width_a, height_a = self.image_a.size
            center_x_b = width_b // 2
            center_y_b = height_b // 5

            scale_w = width_b / width_a * scale_factor
            scale_h = height_b / height_a * scale_factor
        else:

            min_x_b = min([p[0] for p in self.landmarks_b])
            max_x_b = max([p[0] for p in self.landmarks_b])
            min_y_b = min([p[1] for p in self.landmarks_b])
            max_y_b = max([p[1] for p in self.landmarks_b])
            width_b = max_x_b - min_x_b
            height_b = max_y_b - min_y_b
            center_x_b = min_x_b + width_b // 2
            center_y_b = min_y_b + height_b // 2

            scale_w = width_b / face_a.width * 1.4 * scale_factor
            scale_h = height_b / face_a.height * 1.4 * scale_factor

        scale = min(scale_w, scale_h)
        new_width = int(face_a.width * scale)
        new_height = int(face_a.height * scale)
        face_a_resized = face_a.resize((new_width, new_height), Image.LANCZOS)

        face_a_rotated = face_a_resized.rotate(rotation_angle, expand=True)

        final_offset_x = center_x_b - face_a_rotated.width // 2 + offset_x
        final_offset_y = center_y_b - face_a_rotated.height // 2 + offset_y

        result_image = self.overlay_image(self.image_b.copy(), face_a_rotated, (final_offset_x, final_offset_y))

        return result_image

    def adjust_colors(self, is_select_image_a, saturation, temperature, contrast, brightness):
        image = self.image_a_origin if is_select_image_a else self.image_b_origin
        if not image:
            raise ValueError("Selected image is not loaded.")

        # Convert to RGB if necessary
        if image.mode != 'RGB':
            image = image.convert('RGB')

        # Adjust brightness
        enhancer = ImageEnhance.Brightness(image)
        image = enhancer.enhance(brightness)

        # Adjust contrast
        enhancer = ImageEnhance.Contrast(image)
        image = enhancer.enhance(contrast)

        # Adjust saturation
        enhancer = ImageEnhance.Color(image)
        image = enhancer.enhance(saturation)

        # Adjust temperature by modifying color balance
        r, g, b = image.split()  # Handling three channels only
        r = r.point(lambda i: i + temperature * 10)
        b = b.point(lambda i: i - temperature * 10)
        image = Image.merge("RGB", (r, g, b))  # Re-merge as RGB

        if is_select_image_a:
            self.image_a = image
        else:
            self.image_b = image

        return image



# Function to swap faces
def process(image_a, image_b, scale_factor, offset_x, offset_y, rotation_angle, margin_left, margin_right, margin_top,
            margin_bottom, mirror_face_a):
    try:
        face_swapper = FaceSwapper()
        face_swapper.load_images(image_a, image_b)
        if face_swapper.image_a and face_swapper.image_b:
            result = face_swapper.swap_faces(
                scale_factor=scale_factor,
                offset_x=offset_x,
                offset_y=offset_y,
                rotation_angle=rotation_angle,
                margin_left=margin_left,
                margin_right=margin_right,
                margin_top=margin_top,
                margin_bottom=margin_bottom,
                mirror_face_a=mirror_face_a
            )
        return result
    except Exception as e:
        return str(e)


def adujst_color_process(image_a, image_b, scale_factor, offset_x, offset_y, rotation_angle, margin_left, margin_right, margin_top,
                         margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b):
    try:
        face_swapper = FaceSwapper()
        face_swapper.load_images(image_a, image_b)
        face_swapper.adjust_colors(True, saturation, temperature, contrast, brightness)
        face_swapper.adjust_colors(False, saturation_b, temperature_b, contrast_b, brightness_b)
        if face_swapper.image_a and face_swapper.image_b:
            result = face_swapper.swap_faces(
                scale_factor=scale_factor,
                offset_x=offset_x,
                offset_y=offset_y,
                rotation_angle=rotation_angle,
                margin_left=margin_left,
                margin_right=margin_right,
                margin_top=margin_top,
                margin_bottom=margin_bottom,
                mirror_face_a=mirror_face_a
            )
        return result
    except Exception as e:
        return str(e)


# Function for image masking (demonstrative)
def dummy(img):
    # Assuming img is a dictionary with keys 'composite', 'background', and 'layers'
    composite_image = img["composite"]
    imageio.imwrite("output_image.png", composite_image)

    alpha_channel = img["layers"][0][:, :, 3]
    mask = np.where(alpha_channel == 0, 0, 255).astype(np.uint8)

    return mask


# Function to update the image in Tab 1
def update_image(image):
    return image


# Announcement text for tab 1
announcement = """
## Backup Inpaint Mask Maker:
备用蒙版工具: [Visit Here](https://huggingface.co/spaces/BraUndress/inpaint-mask-maker2) 
## How to Use:
使用方法: [Learn More](https://telegra.ph/HowToUploadMaskOnBraundress-05-01) 
Source: BraUndress
"""

# Announcement text for tab 2
announcement_2 = """
## Backup Inpaint Mask Maker:
备用蒙版工具: [Visit Here](https://huggingface.co/spaces/BraUndress/inpaint-mask-maker2) 
## How to Use:
使用方法: [Learn More](https://telegra.ph/100-Similarity-Face-Swapped-with-Braundress-Mask-Upload-Mode-07-23) 
Source: BraUndress
"""

# Create Gradio interface with Tabs
with gr.Blocks() as app:
    tab1 = gr.Tab("Inpaint Mask Maker")
    tab2 = gr.Tab("Face Swap")

    with tab1:
        with gr.Row():
            with gr.Column():
                announcement = gr.Markdown(announcement)


        with gr.Column():
            with gr.Row():
                img = gr.ImageMask(
                    sources=["upload", "clipboard"],
                    transforms=[],
                    layers=False,
                    format="png",
                    label="base image",
                    show_label=True
                )
                img2 = gr.Image(label="mask image", show_label=True)

            btn = gr.Button("Generate Mask")
            btn.click(dummy, inputs=img, outputs=img2)

    with tab2:
        with gr.Row():
            with gr.Column():
                announcement_2 = gr.Markdown(announcement_2)
        with gr.Row():
            img_input_a = gr.Image(type="pil", label="Input Face Image A(输入人脸图片)", height=300)
            img_input_b = gr.Image(type="pil", label="Input Face Swapped Image B(输入换脸图片)", height=300)
        with gr.Row():
            scale_factor = gr.Slider(0.1, 3.5, 1, 0.1, label="Face Scale Factor(人脸放大)")
            offset_x = gr.Slider(-200, 200, 0, 1, label="Face Offset X(人脸水平横移)")
            offset_y = gr.Slider(-200, 200, 0, 1, label="Face Offset Y(人脸垂直横移)")
            rotation_angle = gr.Slider(-180, 180, 0, 1, label="Face Rotation Angle(人脸旋转)")
            margin_left = gr.Slider(0, 1, 0.1, 0.01, label="Face Margin Left(人脸左边缘扩展)")
            margin_right = gr.Slider(0, 1, 0.1, 0.01, label="Face Margin Right(人脸右边缘扩展)")
            margin_top = gr.Slider(0, 1, 0.1, 0.01, label="Face Margin Top(人脸上边缘扩展)")
            margin_bottom = gr.Slider(0, 1, 0, 0.01, label="Face Margin Bottom(人脸下边缘扩展)")
            mirror_face_a = gr.Checkbox(label="Mirror Face A(人脸镜像)")
            # btn_process = gr.Button("Swap Faces")
        with gr.Row():
            image_select = gr.Radio(['Image A'], value='Image A', label="Adjust Color For Image A(图片A调整颜色)")
            saturation = gr.Slider(0, 2.0, 1.0, 0.1, label="Saturation(饱和度)")
            temperature = gr.Slider(-10, 10, 0, 0.1, label="Color Temperature(色温)")
            contrast = gr.Slider(0, 2.0, 1.0, 0.1, label="Contrast(对比度)")
            brightness = gr.Slider(0, 2.0, 1.0, 0.1, label="Brightness(亮度)")
        with gr.Row():
            image_select_b = gr.Radio(['Image B'], value='Image B', label="Adjust Color For Image B(图片B调整颜色)")
            saturation_b = gr.Slider(0, 2.0, 1.0, 0.1, label="Saturation(饱和度)")
            temperature_b = gr.Slider(-10, 10, 0, 0.1, label="Color Temperature(色温)")
            contrast_b = gr.Slider(0, 2.0, 1.0, 0.1, label="Contrast(对比度)")
            brightness_b = gr.Slider(0, 2.0, 1.0, 0.1, label="Brightness(亮度)")
        result_image = gr.Image(type="pil", label="Result Image", height=600)
        # btn_process.click(process, inputs=[img_input_a, img_input_b, scale_factor, offset_x, offset_y, rotation_angle, margin_left, margin_right, margin_top, margin_bottom, mirror_face_a], outputs=result_image)
        scale_factor.release(adujst_color_process,
                             inputs=[img_input_a, img_input_b, scale_factor, offset_x, offset_y, rotation_angle,
                                     margin_left, margin_right, margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b],
                             outputs=result_image)
        offset_x.release(adujst_color_process, inputs=[img_input_a, img_input_b, scale_factor, offset_x, offset_y, rotation_angle,
                                          margin_left, margin_right, margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b],
                         outputs=result_image)
        offset_y.release(adujst_color_process, inputs=[img_input_a, img_input_b, scale_factor, offset_x, offset_y, rotation_angle,
                                          margin_left, margin_right, margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b],
                         outputs=result_image)
        rotation_angle.release(adujst_color_process,
                               inputs=[img_input_a, img_input_b, scale_factor, offset_x, offset_y, rotation_angle,
                                       margin_left, margin_right, margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b],
                               outputs=result_image)
        margin_left.release(adujst_color_process, inputs=[img_input_a, img_input_b, scale_factor, offset_x, offset_y, rotation_angle,
                                             margin_left, margin_right, margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b],
                            outputs=result_image)
        margin_right.release(adujst_color_process,
                             inputs=[img_input_a, img_input_b, scale_factor, offset_x, offset_y, rotation_angle,
                                     margin_left, margin_right, margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b],
                             outputs=result_image)
        margin_top.release(adujst_color_process, inputs=[img_input_a, img_input_b, scale_factor, offset_x, offset_y, rotation_angle,
                                            margin_left, margin_right, margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b],
                           outputs=result_image)
        margin_bottom.release(adujst_color_process,
                              inputs=[img_input_a, img_input_b, scale_factor, offset_x, offset_y, rotation_angle,
                                      margin_left, margin_right, margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b],
                              outputs=result_image)
        img_input_a.change(adujst_color_process, inputs=[img_input_a, img_input_b, scale_factor, offset_x, offset_y, rotation_angle,
                                            margin_left, margin_right, margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b],
                           outputs=result_image)
        img_input_b.change(adujst_color_process, inputs=[img_input_a, img_input_b, scale_factor, offset_x, offset_y, rotation_angle,
                                            margin_left, margin_right, margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b],
                           outputs=result_image)
        mirror_face_a.change(adujst_color_process,
                             inputs=[img_input_a, img_input_b, scale_factor, offset_x, offset_y, rotation_angle,
                                     margin_left, margin_right, margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b],
                             outputs=result_image)
        saturation.release(adujst_color_process,
                           inputs=[img_input_a, img_input_b,scale_factor, offset_x, offset_y, rotation_angle, margin_left, margin_right,
                                   margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b], outputs=result_image)
        temperature.release(adujst_color_process,
                            inputs=[img_input_a, img_input_b,scale_factor, offset_x, offset_y, rotation_angle, margin_left, margin_right,
                                    margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b], outputs=result_image)
        contrast.release(adujst_color_process,
                         inputs=[img_input_a, img_input_b,scale_factor, offset_x, offset_y, rotation_angle, margin_left, margin_right,
                                 margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b], outputs=result_image)
        brightness.release(adujst_color_process,
                           inputs=[img_input_a, img_input_b,scale_factor, offset_x, offset_y, rotation_angle, margin_left, margin_right,
                                   margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b], outputs=result_image)
        saturation_b.release(adujst_color_process,
                           inputs=[img_input_a, img_input_b,scale_factor, offset_x, offset_y, rotation_angle, margin_left, margin_right,
                                   margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b], outputs=result_image)
        temperature_b.release(adujst_color_process,
                            inputs=[img_input_a, img_input_b,scale_factor, offset_x, offset_y, rotation_angle, margin_left, margin_right,
                                    margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b], outputs=result_image)
        contrast_b.release(adujst_color_process,
                         inputs=[img_input_a, img_input_b,scale_factor, offset_x, offset_y, rotation_angle, margin_left, margin_right,
                                 margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b], outputs=result_image)
        brightness_b.release(adujst_color_process,
                           inputs=[img_input_a, img_input_b,scale_factor, offset_x, offset_y, rotation_angle, margin_left, margin_right,
                                   margin_top, margin_bottom, mirror_face_a, saturation, temperature, contrast, brightness, saturation_b, temperature_b, contrast_b, brightness_b], outputs=result_image)
        
        btn_send_to_tab1 = gr.Button("Send To Mask Maker(发送给蒙版制作页面)")
        btn_send_to_tab1.click(update_image, inputs=result_image, outputs=img)

app.launch()