import gradio as gr
import json
from ultralytics import YOLO

# Model Heading and Description
model_heading = "YOLOv11x くずし字認識サービス（一文字）"
description = """YOLOv11x くずし字認識サービス（一文字） Gradio demo for classification. Upload an image or click an example image to use."""

article = "<p style='text-align: center'>YOLOv11x くずし字認識サービス（一文字） is a classification model trained on the <a href=\"http://codh.rois.ac.jp/char-shape/\">日本古典籍くずし字データセット</a>.</p>"

image_path= [
       
    ['U+4F4E_200004148_00022_1_X1018_Y0469.jpg'],
    ['U+5F3E_200015779_00112_1_X0978_Y2642.jpg'],
    ['U+7CBE_100249537_00088_1_X1463_Y0823.jpg']
]

# Load YOLO model
model = YOLO('yolo11x-cls.pt')

def YOLOv11x_img_inference(
    image: gr.Image = None,
):
    """
    YOLOv11x inference function
    Args:
        image: Input image
    Returns:
        top5_json: JSON format of top 5 class names and confidence
    """
    results = model.predict(image)
    result = results[0]
    class_names = result.names  # クラスIDとクラス名のマッピング

    # 上位5件のクラスIDと信頼度を取得して、nameとconfのペアでリストに変換
    top5_list = [
        {
            "name": chr(int(class_names[class_id][2:], 16)),  # Unicodeコードポイントを文字に変換
            "conf": float(conf)
        } 
        for class_id, conf in zip(result.probs.top5, result.probs.top5conf)
    ]

    # JSON形式に変換
    top5_json = json.dumps(top5_list, ensure_ascii=False, indent=2)
    
    return top5_json

    
inputs_image = [
    gr.Image(type="filepath", label="Input Image"),
]

outputs_image =[
    gr.JSON(label="Output JSON")
]
demo = gr.Interface(
    fn=YOLOv11x_img_inference,
    inputs=inputs_image,
    outputs=outputs_image,
    title=model_heading,
    description=description,
    examples=image_path,
    article=article,
    cache_examples=False
    # allow_flagging=False
)

demo.launch(share=False)