import gradio as gr
import glob
import os
import time
import random

# 生成6到12之间的随机整数
delay = random.randint(6, 12)


def generate(front_view: gr.Image, seed=-1):
    origin_file = selected_file_path.split('/')[-1].split('.')[0]
    # print(origin_file)
    # print(seed)
    # print(os.path.join("examples", f'{origin_file}.ply'))
    # 随机延时
    time.sleep(delay)
    return os.path.join("examples", f'{origin_file}.obj')
    # return "./examples/processed_xxy_v2.obj"

def get_png_paths(directory='examples'):
    """
    获取指定目录下所有jpg图片的路径列表。
    
    参数:
    directory (str): 要搜索的目录路径，默认为 'examples'。
    
    返回:
    list[str]: 包含jpg图片文件路径的列表。
    """
    # 使用glob模块匹配目录下所有的jpg文件
    jpg_paths = glob.glob(os.path.join(directory, '*.png'))
    return jpg_paths


selected_file_path = ""

def on_select(evt: gr.SelectData):
    global selected_file_path
    selected_file_path = evt.value['image']['path']
    return selected_file_path


with gr.Blocks() as demo:
    gr.Markdown(
    """
    # 3D Avatar Head Reconstruction based on 2-3DGS
    👋Welcome to use!
    
    - Easy to use
    - Run fast
    - Impressive output
    """)
    with gr.Row():
        with gr.Column(scale=1):
            frontview = gr.Image(label="FrontView")
            examples = gr.Gallery(get_png_paths(), columns=6, height=200, object_fit="contain", allow_preview=False)
            examples.select(fn=on_select, inputs=None, outputs=frontview)
        with gr.Column(scale=1):
            meshmodel = gr.Model3D(label="Mesh Model")
            seed = gr.Slider(-1, 10000, value=0, label="Seed", info="random seed")
            btn = gr.Button("Generate 3D Model")
            btn.click(fn=generate, inputs=[frontview, seed], outputs=meshmodel)
                    
if __name__ == "__main__":
    demo.launch()