import gradio as gr
import json

from ultralyticsplus import YOLO, render_result

# Model Heading and Description
model_heading = "東京大学史料編纂所所蔵『日本一鑑』の行検出モデル"
description = """東京大学史料編纂所所蔵『日本一鑑』の行検出モデルは、YOLOv8をベースにしたモデルです。このモデルは、東京大学史料編纂所所蔵『日本一鑑』の行を検出するために訓練されています。モデルは、画像を入力として受け取り、画像内の東京大学史料編纂所所蔵『日本一鑑』の行を検出し、その位置を示すバウンディングボックスを描画します。モデルは、信頼度スコアとIOUスコアの2つのしきい値を受け取ります。信頼度スコアは、検出されたオブジェクトの信頼度を示し、IOUスコアは、検出されたオブジェクトの位置の正確さを示します。"""

image_path= [
        ['test/00000017.jpg', 0.1, 0.1],
    ['test/00000020.jpg', 0.1, 0.1]
]

# Load YOLO model
model = YOLO('best.pt')

def yolov8_img_inference(
    image: gr.Image = None,
    conf_threshold: gr.Slider = 0.25,
    iou_threshold: gr.Slider = 0.45,
):
    """
    YOLOv8 inference function
    Args:
        image: Input image
        conf_threshold: Confidence threshold
        iou_threshold: IOU threshold
    Returns:
        Rendered image
    """
    results = model.predict(image, conf=conf_threshold, iou=iou_threshold)
    render = render_result(model=model, image=image, result=results[0])            

    json_data = json.loads(results[0].tojson())
    
    return render, json_data

    
inputs_image = [
    gr.Image(type="filepath", label="Input Image"),
    gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold"),
    gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold"),
]

outputs_image =[
    gr.Image(type="filepath", label="Output Image"),
    gr.JSON(label="Output JSON")
]
interface_image = gr.Interface(
    fn=yolov8_img_inference,
    inputs=inputs_image,
    outputs=outputs_image,
    title=model_heading,
    description=description,
    examples=image_path,
    cache_examples=False,
    theme='huggingface'
)


'''
demo = gr.TabbedInterface(
    [interface_image],
    tab_names=['Image inference']
).queue() # .launch()
'''

demo = interface_image

# if __name__ == "__main__":
demo.launch()