Update app.py

app.py

@@ -1,7 +1,209 @@
-import gradio as gr
+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"]
@@ -12,36 +214,152 @@ def dummy(img):
 
     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)
+备用蒙版工具: [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
 """
 
-with gr.Blocks() as demo:
-    # 将公告栏放在一个单独的Column中
-    with gr.Row():
+# 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():
-            announcement = gr.Markdown(announcement)
-    
-    # 创建一个新的Column来放置ImageMask和按钮
-    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():
-            img = gr.ImageMask(
-                sources=["upload"], 
-                layers=True, 
-                transforms=[], 
-                format="png", 
-                label="base image", 
-                show_label=True
-            )
-            img2 = gr.Image(label="mask image", show_label=True)
+            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 = gr.Button("Generate Mask")
-        btn.click(dummy, inputs=img, outputs=img2)
+        btn_send_to_tab1 = gr.Button("Send To Mask Maker(发送给蒙版制作页面)")
+        btn_send_to_tab1.click(update_image, inputs=result_image, outputs=img)
 
-demo.launch(debug=True)
+app.launch()