import cv2 as cv
import numpy as np
from PIL import Image, ImageDraw
import gradio as gr
from ultralytics import YOLO
from huggingface_hub import hf_hub_download

# ----------------- load YOLOv8n handwritten model -----------------
try:
    yolo_path = hf_hub_download(
        repo_id="armvectores/yolov8n_handwritten_text_detection",
        filename="best.pt"
    )
    model = YOLO(yolo_path)
    print("✅ Loaded YOLOv8n handwriting detector")
except Exception as e:
    print("❌ Could not load YOLO:", e)
    model = None

# ----------------- run detection -----------------
def detect_hw(image: Image.Image, conf=0.25, iou=0.45):
    if image is None:
        return None, "Brak obrazu"
    if model is None:
        return image, "Brak modelu"

    rgb = np.array(image.convert("RGB"))
    results = model.predict(source=rgb, imgsz=1280, conf=conf, iou=iou, verbose=False)

    if not results or len(results) == 0:
        return image, "Brak wyników"

    # bboxes
    boxes = []
    for b in results[0].boxes.xyxy.cpu().numpy():
        x1, y1, x2, y2 = [int(v) for v in b[:4]]
        boxes.append([x1, y1, x2, y2])

    # draw
    pil = image.convert("RGB").copy()
    dr = ImageDraw.Draw(pil)
    for (x1, y1, x2, y2) in boxes:
        dr.rectangle([x1, y1, x2, y2], outline=(255, 140, 0), width=3)
        dr.text((x1, max(0, y1-12)), "handwritten", fill=(255, 140, 0))

    return pil, f"Znaleziono {len(boxes)} obszarów odręcznych"

# ----------------- Gradio UI -----------------
with gr.Blocks(title="YOLO Handwritten Detector") as demo:
    gr.Markdown("### 🖊 YOLOv8n Handwritten Text Detection\nModel: `armvectores/yolov8n_handwritten_text_detection`")

    with gr.Row():
        inp = gr.Image(type="pil", label="Wgraj dokument/obraz")
        out_img = gr.Image(type="pil", label="Podgląd z bboxami")
    out_txt = gr.Textbox(label="Info")

    conf = gr.Slider(0, 1, 0.25, step=0.05, label="Confidence")
    iou  = gr.Slider(0, 1, 0.45, step=0.05, label="IoU threshold")

    run_btn = gr.Button("Uruchom detekcję")

    run_btn.click(detect_hw, inputs=[inp, conf, iou], outputs=[out_img, out_txt])

demo.queue()
if __name__ == "__main__":
    demo.launch(ssr_mode=False)