import gradio as gr
from PIL import Image
import time,os
from pathlib import Path 
base_dir = '猫咪表情包'
selected_dir = 'selected'
files = [str(x) for x in 
         Path(base_dir).rglob('*.jp*g') 
         if 'checkpoint' not in str(x)]
files_idx=0
selected_count=0
    
def show_img(path):
    return Image.open(path)
def fn_before(done,todo):
    global files_idx
    if files_idx>0:
        files_idx-=1
    path=files[files_idx]
    img=show_img(path)
    return done,todo,path,img
def fn_next(done,todo):
    global files_idx
    if files_idx<len(files)-1:
        files_idx+=1
    path=files[files_idx]
    img=show_img(path)
    return done,todo,path,img
def save_selected(img_path,done,todo):
    global msg,selected_count
    if img_path not in msg:
        msg=msg+'  '+img_path
        selected_count+=1
        done=selected_count
        todo=len(files)-selected_count
    return msg,done,todo 
def get_default_msg():
    global msg
    msg=''
    return msg
    
    
with gr.Blocks() as demo:
    with gr.Row():
        total = gr.Number(len(files),label='总数量')
        with gr.Row(scale = 1):
            bn_before = gr.Button("上一张")
            bn_next = gr.Button("下一张")
        with gr.Row(scale = 2):
            done = gr.Number(value=selected_count,label='已完成')
            todo = gr.Number(value=len(files)-selected_count,label='待完成')
    with gr.Row():
        with gr.Column():
            path = gr.Text(files[0],lines=1, label='当前图片路径')
            feedback_button = gr.Button("选择图片",variant="primary")
            msg = gr.TextArea(value=get_default_msg,lines=3,max_lines = 5)
        with gr.Column():
            img = gr.Image(value = show_img(files[0]),type='pil')
    
    bn_before.click(fn_before,
                 inputs= [done,todo], 
                 outputs=[done,todo,path,img])
    bn_next.click(fn_next,
                 inputs= [done,todo], 
                 outputs=[done,todo,path,img])
    feedback_button.click(save_selected,
                         inputs = [path,done,todo],
                         outputs =[msg,done,todo]
                         )
gr.close_all()
demo.launch()