import gradio as gr
import torch
from PIL import Image, ImageDraw
import json

# Model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', source="local") 

def yolo(im, size=1024):
    g = (size / max(im.size))  # gain
    im = im.resize((int(x * g) for x in im.size), Image.BICUBIC)  # resize

    results = model(im)  # inference
    df = results.pandas().xyxy[0].to_json(orient="records")
    res = json.loads(df)

    detected_images = []

    im_draw = im.copy()

    draw = ImageDraw.Draw(im_draw)

    # JSONデータ内の座標に基づいて矩形を描画
    for item in res:

        xmin = item['xmin']#  * w
        ymin = item['ymin']#  * h
        xmax = item['xmax']#  * w
        ymax = item['ymax']#  * h
        draw.rectangle([(xmin, ymin), (xmax, ymax)], outline="red", width=2)

        # Extract each detected object into a separate image
        detected_object = im.crop((xmin, ymin, xmax, ymax))
        detected_images.append(detected_object)

    return [
        res,
        im_draw,
        detected_images
    ]


inputs = [
    gr.Image(type='pil', label="Original Image"),
]
outputs = [
    gr.JSON(label="JSON Output"),
    gr.Image(type="pil", label="Output Image"),
    gr.Gallery(label="Detected Objects", object_fit="contain")
]

title = "YOLOv5 Kunshujo"
description = "YOLOv5 Kunshujo Gradio demo for object detection. Upload an image or click an example image to use."
article = "<p style='text-align: center'>YOLOv5 Kunshujo is an object detection model trained on the <a href=\"https://github.com/utda/kunshujo-layout-dataset\">Kunshujo layout dataset</a>.<br/>Examples: 張交帖.[5] (国立国会図書館 National Diet Library, JAPAN)</p>"

examples = [
    ['2586696_R0000008.jpg'],
    ['2586696_R0000009.jpg']
]
demo = gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, allow_flagging="never")

demo.launch()