import gradio as gr
from NTED.NTED_module import NTED

NTED_Module = NTED()

def pose_transfer(上传人体姿态图):
    img = 上传人体姿态图
    fake_img = NTED_Module.inference(img)
    
    return fake_img

with gr.Column():
    result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')

gr.Interface(fn=pose_transfer, 
             inputs=["image"], 
             outputs=[result_gallery],
             title="谷小雨姿态驱动图像",
             examples=[["example/exp1.png"], ["example/exp2.png"], ["example/exp3.png"],\
                 ["example/exp4.png"], ["example/exp5.png"], ["example/exp6.png"]],
             ).launch(share=True)  # (server_name='0.0.0.0')  

'''
TODO
测试视频展示功能
'''