"""
Gradio APP for Yolov5 medieval registers layout analysis.

Date: 20/12/2022
Author: lterriel
"""
import glob
import sys, types

import gradio as gr
from PIL import Image
from huggingface_hub import hf_hub_download

# --- HF errors shim (compat with old yolov5 import path)
try:
    import huggingface_hub.errors as hf_errors
    if 'huggingface_hub.utils' not in sys.modules:
        sys.modules['huggingface_hub.utils'] = types.ModuleType('huggingface_hub.utils')
    sys.modules['huggingface_hub.utils._errors'] = hf_errors
except Exception:
    pass

import yolov5  # must come after the shim

# Add new models here:
model_names = [
    "lterriel/endp-yolov5x-35e-bs4",
]

def load_model(model_name: str) -> str:
    # download local checkpoint from HF
    return hf_hub_download(repo_id=model_name, filename="best.pt")

def yolo_inference(im, model_name, threshold=0.50):
    model_path = load_model(model_name)
    model = yolov5.load(model_path)
    model.conf = float(threshold)
    results = model(im)  # inference
    numpy_image = results.render()[0]
    return Image.fromarray(numpy_image)

title = "YOLOv5 - Medieval Register Segmentation"
description = """<h5 style='text-align: center;'>YOLOv5 Gradio demo for medieval register layout analysis.</h3><br>
<ul>
<li>YOLOv5 source code : <a href='https://github.com/ultralytics/yolov5'>Source code</a></li>
<li>The examples provided are digitized images extracted from the <i>registers of the chapter of Notre-Dame (1326-1504)</i> available on the <a href='https://www.siv.archives-nationales.culture.gouv.fr/siv/rechercheconsultation/consultation/ir/consultationIR.action?irId=FRAN_IR_059635' target='_blank'>SIV of the National Archives of France</a></li>
</ul>
"""

inputs = [
    gr.Image(type="pil", label="document image"),
    gr.Dropdown(choices=model_names, label="Model", value=model_names[0]),
    gr.Slider(minimum=0, maximum=1, step=0.01, value=0.50, label="Confidence threshold"),
]

examples = [[str(file), model_names[0], 0.50] for file in glob.glob("./images_examples/*.jpg")]

demo = gr.Interface(
    fn=yolo_inference,
    inputs=inputs,
    outputs=gr.Image(type="pil", label="annotated document", height=800),
    title=title,
    description=description,
    theme="default",
    examples=examples,
    cache_examples=False,   # avoid startup execution
)

if __name__ == "__main__":
    demo.launch(debug=True)