import gradio as gr
from gradio import Image, Textbox
from infer import CaptionModel
import numpy

# 定义 CaptionModel 对象
model_rl_grid = CaptionModel(decoder_file="rl/grid-based/model_49_5.1001_0418-1216.pth")
model_rl_region = CaptionModel(decoder_file="rl/region-based/model_49_5.0719_0420-1839.pth")
model_xe_grid = CaptionModel(decoder_file="xe/grid-based/model_49_2.8938_0417-1245.pth")
model_xe_region = CaptionModel(decoder_file="xe/region-based/model_49_2.8568_0420-1657.pth")

# 定义生成字幕的函数
def generate_caption(img):
    # read PLI.Image to numpy array
    img = numpy.array(img)
    captions_rl_grid = model_rl_grid.predict(img)[0]
    captions_rl_region = model_rl_region.predict(img)[0]
    captions_xe_grid = model_xe_grid.predict(img)[0]
    captions_xe_region = model_xe_region.predict(img)[0]
    return captions_rl_grid, captions_rl_region, captions_xe_grid, captions_xe_region
    


# 创建 Gradio 界面
interface = gr.Interface(
    fn=generate_caption,
    inputs=Image(type="pil", label="选择一张图片"),
    outputs=[
        Textbox(label="生成的字幕 1"),
        Textbox(label="生成的字幕 2"),
        Textbox(label="生成的字幕 3"),
        Textbox(label="生成的字幕 4"),
    ],
    title="自动字幕生成",
)
interface.launch()