app.py

import spaces
import gradio as gr
import subprocess
from PIL import Image,ImageEnhance,ImageFilter
import json

import mp_box
'''
Face landmark detection based Face Detection.
https://ai.google.dev/edge/mediapipe/solutions/vision/face_landmarker
from model card
https://storage.googleapis.com/mediapipe-assets/MediaPipe%20BlazeFace%20Model%20Card%20(Short%20Range).pdf
Licensed Apache License, Version 2.0
Train with google's dataset(more detail see model card)

Not Face Detector based 
https://ai.google.dev/edge/mediapipe/solutions/vision/face_detector

Bacause this is part of getting-landmark program and need control face edge.
So I don't know which one is better.never compare these.
'''

def select_box(boxes,box_type):
    if box_type == "type-3":
         box = boxes[2]
    elif box_type =="type-2":
         box = boxes[1]
    elif box_type =="type-1":
         box = boxes[0]
    else:#never happen
         box=[0,0,image.size[0],image.size[1]]
    box_width = box[2]
    box_height = box[3]
    box = mp_box.xywh_to_xyxy(box)
    return box,box_width,box_height

def process_images(image,replace_image=None,replace_image_need_crop=False,box_type="type-3",fill_color_mode=False,fill_color="black",custom_color="rgba(255,255,255,1)",image_size=1024,filter_image=False,filter_value="Sharpen",progress=gr.Progress(track_tqdm=True)):
    if image == None:
        raise gr.Error("Need Image")

    # choose box
    boxes,mp_image,face_landmarker_result = mp_box.mediapipe_to_box(image)
    box,box_width,box_height = select_box(boxes,box_type)

    # replace-mode
    if replace_image!=None:
        print("replace mode")

        if replace_image_need_crop:
            replace_boxes,mp_image,face_landmarker_result = mp_box.mediapipe_to_box(replace_image)
            replace_box,replace_box_width,replace_box_height = select_box(replace_boxes,box_type)

        if fill_color_mode:
            if replace_image_need_crop:
                cropped = replace_image.crop(replace_box)
                cropped.resize(box_width,box_height)
            else:
                cropped = replace_image.crop(box)
            image.paste(cropped,[box[0],box[1]])
            return image
        else:#scale mode
            if replace_image_need_crop:
                replace_image = replace_image.crop(replace_box)
            replace_resized = replace_image.resize((box_width,box_height),Image.Resampling.LANCZOS)
            image.paste(replace_resized,[box[0],box[1]])
            return image


    # crop-mode
    if fill_color_mode:
        # choose color
        color_map={
            "black":[0,0,0,1],
            "white":[255,255,255,1],
            "red":[255,0,0,1],
            "brown":[92,33,31,1],
            "pink":[255,192,203,1],
        }
        if fill_color == "custom":
            color_value = custom_color.strip("rgba()").split(",")
            color_value[0] = int(float(color_value[0]))
            color_value[1] = int(float(color_value[1]))
            color_value[2] = int(float(color_value[2]))
        else:
            color_value = color_map[fill_color]

        cropped = image.crop(box)

        img = Image.new('RGBA', image.size, (color_value[0], color_value[1], color_value[2]))
        img.paste(cropped,[box[0],box[1]])
        return img
    else:
        #scale up mode
        cropped = image.crop(box)
        resized = resize_image_by_max_dimension(cropped,image_size)
        filter_map={
           "Blur":ImageFilter.BLUR,"Smooth More":ImageFilter.SMOOTH_MORE,"Smooth":ImageFilter.SMOOTH,"Sharpen":ImageFilter.SHARPEN,"Edge Enhance":ImageFilter.EDGE_ENHANCE,"Edge Enhance More":ImageFilter.EDGE_ENHANCE_MORE
        }
        if filter_value not in filter_map:
            raise gr.Error(f"filter {filter_value} not found")
        if filter_image:
           
            #resized = resized.filter(ImageFilter.SHARPEN)
            #Gimp's weak 0.1-0.2?
            enhancer = ImageEnhance.Sharpness(resized)
            resized = resized.filter(filter_map[filter_value])
            #resized = enhancer.enhance(sharpen_value)


        return resized


def resize_image_by_max_dimension(image, max_size, resampling=Image.Resampling.BICUBIC):
    image_width, image_height = image.size
    
    max_dimension = max(image_width, image_height)
    
    ratio = max_size / max_dimension
    
    new_width = int(image_width * ratio)
    new_height = int(image_height * ratio)
 
    return image.resize((new_width, new_height), resampling)
    

def read_file(file_path: str) -> str:
    """read the text of target file
    """
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()

    return content

css="""
#col-left {
    margin: 0 auto;
    max-width: 640px;
}
#col-right {
    margin: 0 auto;
    max-width: 640px;
}
.grid-container {
  display: flex;
  align-items: center;
  justify-content: center;
  gap:10px
}

.image {
  width: 128px; 
  height: 128px; 
  object-fit: cover; 
}

.text {
  font-size: 16px;
}
"""

#css=css,
def update_button_label(image):
    if image == None:
        return  gr.Button(visible=bool(0)),gr.Button(visible=bool(1))
    else:
        return  gr.Button(visible=bool(1)),gr.Row(visible=bool(0))   
    
def update_visible(fill_color_mode):
     if fill_color_mode:
        return  gr.Row(visible=bool(0)),gr.Row(visible=bool(1))
     else:
        return  gr.Row(visible=bool(1)),gr.Row(visible=bool(0))    
with gr.Blocks(css=css, elem_id="demo-container") as demo:
    with gr.Column():
        gr.HTML(read_file("demo_header.html"))
        gr.HTML(read_file("demo_tools.html"))
    with gr.Row():
                with gr.Column():
                    image = gr.Image(sources=['upload','clipboard'],image_mode='RGB',elem_id="image_upload", type="pil", label="Upload")
                    box_type = gr.Dropdown(label="box-type",value="type-3",choices=["type-1","type-2","type-3"])
                    with gr.Row(elem_id="prompt-container",  equal_height=False):
                        with gr.Row():
                            btn1 = gr.Button("Face Crop", elem_id="run_button",variant="primary")
                            btn2 = gr.Button("Face Replace", elem_id="run_button2",variant="primary",visible=False)
                            
                    replace_image = gr.Image(sources=['upload','clipboard'],image_mode='RGB',elem_id="replace_upload", type="pil", label="replace image")
                    replace_image_need_crop = gr.Checkbox(label="Replace image need crop",value=False)
                    replace_image.change(update_visible,replace_image,[btn1,btn2])
                    with gr.Accordion(label="Advanced Settings", open=False):
                        fill_color_mode = gr.Checkbox(label="Fill Color Mode",value=False)
                        row1 = gr.Row(equal_height=True)
                        row2 = gr.Row(equal_height=True,visible=False)
                        fill_color_mode.change(update_visible,fill_color_mode,[row1,row2])
                        
                        with row1:
                            image_size = gr.Slider(
                            label="Image Size",info = "cropped face size",
                            minimum=8,
                            maximum=2048,
                            step=1,
                            value=1024,
                            interactive=True)

                            filter_image = gr.Checkbox(label="Filter image")
                            filter_value = gr.Dropdown(label="Filter",value="Sharpen",choices=["Blur","Smooth More","Smooth","Sharpen","Edge Enhance","Edge Enhance More"])
                        with row2:
                            
                            fill_color = gr.Dropdown(label="fill color",choices=["black","white","red","brown","pink","custom"])
                            custom_color = gr.ColorPicker(label="custom color",value="rgba(250, 218, 205, 1)")
                            
                    
                with gr.Column():
                    image_out = gr.Image(label="Output", elem_id="output-img")
                    
                    
                    
            

    
    gr.on(
        [btn1.click,btn2.click],
        fn=process_images, inputs=[image,replace_image,replace_image_need_crop,box_type,fill_color_mode,fill_color,custom_color,image_size,filter_image,filter_value], outputs =[image_out], api_name='infer'
    )
    gr.Examples(
                examples =["examples/00004200.jpg"],
                inputs=[image]
    )
    gr.HTML(read_file("demo_footer.html"))

    if __name__ == "__main__":
        demo.launch()