app.py

from mtcnn.mtcnn import MTCNN
from utils import *


cloth_examples = get_cloth_examples()
pose_examples = get_pose_examples()
face_detector = MTCNN()

# Description
title = r"""
<h1 align="center">IDM-VTON + Outfit Anyone in the Wild </h1>
"""

description = r"""
<a href='https://heybeauty.ai/extension' target='_blank'><b> Go to HeyBeauty for Free Try-On! 🤗 </b></a>.<br>

This demo combines <b>IDM-VTON </b> and <b>Outfit Anyone in the Wild </b>
1. Human pose detection and reconstruction using large human model from Outfit Anyone in the Wild.
2. Use IDM-VTON for training-free try-on.
3. Applying the refine network from Outfit Anyone in the Wild.<br>
This demo is for learning purposes only.<br>
<b>IDM-VTON + Outfit Anyone in the Wild test results on </b> <a href='https://heybeautify.online/ClothData/Publics/Shows/shows/IDM_OutfitAnyone/IDM_OutfitAnyone.html' target='_blank'><b>mix01</b></a>.<br>
<b>Outfit Anyone in the Wild test results on </b> <a href='https://heybeautify.online/ClothData/Publics/Shows/shows/man_v2/man_v2.html' target='_blank'><b>man01</b></a>.<br>
<b>Outfit Anyone in the Wild test results on </b> <a href='https://heybeautify.online/ClothData/Publics/Shows/shows/cider/cider_0403.html' target='_blank'><b>woman01</b></a>.<br>
"""

css = """
.gradio-container {width: 85% !important}
"""

mk_guide = "If image does not display successfully after button clicked in your browser(mostly Mac+Chrome), try [this demo](https://openxlab.org.cn/apps/detail/jiangxiaoguo/OutfitAnyone-in-the-Wild) please"


def onUpload():
    return ""


def onClick(cloth_image, cloth_id, pose_image, pose_id, category, 
        denoise_steps, caption, request: gr.Request):
    if pose_image is None:
        return None, "no pose image found !", ""
    if isinstance(cloth_id, dict):
        cloth_id = cloth_id['label']
    if isinstance(pose_id, dict):
        pose_id = pose_id['label']
    if len(pose_id)>0 and len(cloth_id)>0:
        res = get_result_example(cloth_id, pose_id)
        assert os.path.exists(res), res
        return res, "Done! Use the pre-run results directly, the cloth size does not take effect ", ""
    else:
        try:
            client_ip = request.client.host
            x_forwarded_for = dict(request.headers).get('x-forwarded-for')
            if x_forwarded_for:
                client_ip = x_forwarded_for
                
            faces = face_detector.detect_faces(pose_image[:,:,::-1])
            if len(faces)==0:
                print(client_ip, 'faces num is 0! ', flush=True)
                return None, "Fatal Error !!! No face detected in pose image !!! ", ""
            else:
                x, y, w, h = faces[0]["box"]
                H, W = pose_image.shape[:2]
                max_face_ratio = 1/3.3
                if w/W>max_face_ratio or h/H>max_face_ratio:
                    return None, "Fatal Error !!! Headshot is not allowed in pose image!!!", ""
            if not check_warp(client_ip):
                return None, "Failed !!! Our server is under maintenance, please try again tomorrow", ""
            
            infId = upload_imgs(ApiUrl, OpenId, ApiKey, client_ip, cloth_image, pose_image)
            if infId==0:
                return None, "fail to upload", ""
            elif infId==2:
                return None, "There is a running task already, please wait and check the history tab. Please remember to give us a star on github, thx~", ""
            elif infId==3:
                return None, "can not creat task, you have exhausted free trial quota. Go to heybeauty for more free trial ! https://heybeauty.ai/extension", ""

            isPub = publicFastSwap(ApiUrl, OpenId, ApiKey, infId, category, 
                caption, denoise_steps)
            if not isPub:
                return None, "fail to public you task", ""
            info =  "task has been created successfully, you can refresh the page 1~3 mins latter, and check the following history tab"
            info = info+"任务创建成功，请1-3分钟后刷新这个页面，历史结果会显示在下面的标签页"
            return None, info, ""
        except Exception as e:
            print(e)
            return None, "fail to create task", ""

def onLoad(request: gr.Request):
    client_ip = request.client.host
    x_forwarded_for = dict(request.headers).get('x-forwarded-for')
    if x_forwarded_for:
        client_ip = x_forwarded_for
    his_datas = [None for _ in range(10)]
    info = ''
    try:
        infs = getAllFastInfs(ApiUrl, OpenId, ApiKey, client_ip)
        print(client_ip, 'history infs: ', len(infs))
        cnt = 0
        finish_n, fail_n, queue_n = 0, 0, 0
        for i, inf in enumerate(infs):
            if inf['state']==2:
                if cnt>4: continue
                pose, res = inf['pose'], inf['res']
                his_datas[cnt*2] = f"<img src=\"{pose}\" >"
                his_datas[cnt*2+1] = f"<img src=\"{res}\" >"
                finish_n += 1
                cnt += 1
            elif inf['state'] in [-1, -2, 0]:
                fail_n += 1
            elif inf['state'] in [1]:
                queue_n += 1
        info = f"{client_ip}, you have {finish_n} successed tasks, {queue_n} running tasks, {fail_n} failed tasks."
        if fail_n>0:
            info = info+" Please upload a half/full-body human image, not just a clothing image!!!"
        if queue_n>0:
            position = inf['position']
            info = info+" Wait for 3~10 mins and refresh this page, successed results will display in the history tab at the bottom. "
            info = info+f" your task position in queue is {position}. "
            info = info+f" 任务正在排队，队列位置 {position}. "
        info = info + "Go to heybeauty for better virtual try-on ! https://heybeauty.ai/extension"
        time.sleep(3)
    except Exception as e:
        print(e)
    his_datas = his_datas + [info]
    return his_datas

with gr.Blocks(css=css) as demo:
    # description
    gr.Markdown(title)
    gr.Markdown(description)
                    
    with gr.Row():
        with gr.Column():
            with gr.Column():
                cloth_image = gr.Image(value=None, type="numpy", label="")
                cloth_id = gr.Label(value=cloth_examples[0][0], label="Clothing Image", visible=False)
                example = gr.Examples(inputs=[cloth_id, cloth_image],
                                      examples_per_page=3,
                                      examples = cloth_examples)
        with gr.Column():
            with gr.Column():
                pose_image = gr.Image(value=None, type="numpy", label="")
                pose_id = gr.Label(value=pose_examples[0][0], label="Pose Image", visible=False)
                example_pose = gr.Examples(inputs=[pose_id, pose_image],
                                          examples_per_page=3,
                                          examples=pose_examples)
        with gr.Column():
            with gr.Column():
                category = gr.Dropdown(value="upper_cloth", choices=["upper_cloth", 
                    "lower_cloth", "full_body", "dresses"], interactive=True)
                denoise_steps = gr.Slider(20, 30, value=20, interactive=True, label="denoise_steps")
                caption = gr.Textbox(value="", interactive=True, label='cloth caption')
                
                info_text = gr.Textbox(value="", interactive=False, label='runtime information')
                run_button = gr.Button(value="Run")
                init_res = get_result_example(cloth_examples[0][0], pose_examples[0][0])
                res_image = gr.Image(label="result image", value=None, type="filepath")
                MK01 = gr.Markdown()

    with gr.Tab('history'):

        with gr.Row():
            MK02 = gr.Markdown()

        with gr.Row():
            his_pose_image1 = gr.HTML()
            his_res_image1 = gr.HTML()

        with gr.Row():
            his_pose_image2 = gr.HTML()
            his_res_image2 = gr.HTML()

        with gr.Row():
            his_pose_image3 = gr.HTML()
            his_res_image3 = gr.HTML()            

        with gr.Row():
            his_pose_image4 = gr.HTML()
            his_res_image4 = gr.HTML()            

        with gr.Row():
            his_pose_image5 = gr.HTML()
            his_res_image5 = gr.HTML()            

    run_button.click(fn=onClick, inputs=[cloth_image, cloth_id, pose_image, 
        pose_id, category, denoise_steps, caption, ], 
        outputs=[res_image, info_text, MK01])

    pose_image.upload(fn=onUpload, inputs=[], outputs=[pose_id],)
    cloth_image.upload(fn=onUpload, inputs=[], outputs=[cloth_id],)
    demo.load(onLoad, inputs=[], outputs=[his_pose_image1, his_res_image1,
        his_pose_image2, his_res_image2, his_pose_image3, his_res_image3,
        his_pose_image4, his_res_image4, his_pose_image5, his_res_image5,
        MK02])

if __name__ == "__main__":

    demo.queue(max_size=50)
    # demo.launch(server_name='0.0.0.0', server_port=225)
    demo.launch(server_name='0.0.0.0')