import gradio as gr
import json
from PIL import Image, ImageDraw
from ultralytics import YOLO

# Model Heading and Description
model_heading = "YOLOv11x Character"
description = """YOLOv11x Character Gradio demo for object detection. Upload an image or click an example image to use."""

article = "<p style='text-align: center'>YOLOv11x Character is an object detection model trained on the <a href=\"http://codh.rois.ac.jp/char-shape/\">日本古典籍くずし字データセット</a>.</p>"

image_path= [
       
    ['『源氏物語』(東京大学総合図書館所蔵).jpg', 0.25, 0.45],
    ['『源氏物語』(京都大学所蔵).jpg', 0.25, 0.45],
    ['『平家物語』(国文学研究資料館提供).jpg', 0.25, 0.45]
]

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


def get_color(score):
    """Returns color based on confidence score."""
    if score > 0.75:
        return "blue"       # 高スコアに濃い青
    elif score > 0.5:
        return "deepskyblue"  # 中スコアに明るい青
    elif score > 0.25:
        return "lightblue"  # 低スコアに薄い青
    else:
        return "gray"       # 非常に低いスコアにグレー

def draw_boxes(image_path, results):
    # Open image
    image = Image.open(image_path)
    draw = ImageDraw.Draw(image)

    # 画像の短辺に基づいて矩形の線の太さを調整
    min_dimension = min(image.size)  # 画像の短辺を取得
    line_width = max(1, min_dimension // 200)  # 線の太さを短辺の1%程度に設定（最小値は1）
    
    # Draw boxes
    for item in results:
        box = item['box']
        # label = item['class']
        score = item['confidence']
        
        # Define box coordinates
        x1, y1, x2, y2 = box["x1"], box["y1"], box["x2"], box["y2"]

        color = get_color(score)
        
        # Draw rectangle and label
        draw.rectangle([x1, y1, x2, y2], outline=color, width=line_width)
        # draw.text((x1, y1), f"{label} {score:.2f}", fill=color)
    
    return image

def YOLOv11x_img_inference(
    image: gr.Image = None,
    conf_threshold: gr.Slider = 0.25,
    iou_threshold: gr.Slider = 0.45,
):
    """
    YOLOv11x inference function
    Args:
        image: Input image
        conf_threshold: Confidence threshold
        iou_threshold: IOU threshold
    Returns:
        Rendered image
        JSON output
    """
    results = model.predict(image, conf=conf_threshold, iou=iou_threshold, device="cpu")          

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

    # Draw boxes on image
    result_image = draw_boxes(image, json_data)
    
    return result_image, 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")
]
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
)

demo.css = """
.json-holder {
    height: 300px;
    overflow: auto;
}
"""

demo.launch(share=False)