"""花押検出デモ (YOLOv5x) — nakamura196/yolov5-kao

重みは Space に同梱せず、モデルリポジトリ nakamura196/yolov5-kao から取得する。
推論は onnxruntime のみで完結させ、torch と vendored yolov5 への依存を外している
(コールドスタートが速く、GitHub からの実行時 clone にも依存しない)。
"""
from __future__ import annotations

import gradio as gr
import numpy as np
import onnxruntime as ort
from PIL import Image, ImageDraw
from huggingface_hub import hf_hub_download

REPO_ID = "nakamura196/yolov5-kao"
IMGSZ = 1024  # ONNX の入力は 1024x1024 固定 (batch のみ dynamic)

_session: ort.InferenceSession | None = None


def session() -> ort.InferenceSession:
    """初回呼び出し時にだけ ONNX を取得してセッションを作る (起動を待たせない)。"""
    global _session
    if _session is None:
        path = hf_hub_download(REPO_ID, "best.onnx")
        _session = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
    return _session


def letterbox(im: Image.Image) -> tuple[np.ndarray, float, int, int]:
    """長辺を IMGSZ に合わせ、余白を灰色(114)で中央パディングする。"""
    w, h = im.size
    scale = IMGSZ / max(w, h)
    nw, nh = round(w * scale), round(h * scale)
    canvas = Image.new("RGB", (IMGSZ, IMGSZ), (114, 114, 114))
    pad_l, pad_t = (IMGSZ - nw) // 2, (IMGSZ - nh) // 2
    canvas.paste(im.resize((nw, nh), Image.BILINEAR), (pad_l, pad_t))
    x = np.asarray(canvas, dtype=np.float32) / 255.0
    return np.ascontiguousarray(x.transpose(2, 0, 1))[None], scale, pad_l, pad_t


def nms(boxes: np.ndarray, scores: np.ndarray, thr: float) -> list[int]:
    """クラス非依存 greedy NMS。boxes は xyxy。"""
    order = scores.argsort()[::-1]
    area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    keep: list[int] = []
    while len(order):
        i = order[0]
        keep.append(int(i))
        if len(order) == 1:
            break
        rest = order[1:]
        x1 = np.maximum(boxes[i, 0], boxes[rest, 0])
        y1 = np.maximum(boxes[i, 1], boxes[rest, 1])
        x2 = np.minimum(boxes[i, 2], boxes[rest, 2])
        y2 = np.minimum(boxes[i, 3], boxes[rest, 3])
        inter = np.clip(x2 - x1, 0, None) * np.clip(y2 - y1, 0, None)
        order = rest[inter / (area[i] + area[rest] - inter) <= thr]
    return keep


def detect(image: Image.Image, conf: float, iou: float):
    if image is None:
        raise gr.Error("画像をアップロードしてください / Please upload an image.")

    image = image.convert("RGB")
    x, scale, pad_l, pad_t = letterbox(image)

    sess = session()
    pred = sess.run(None, {sess.get_inputs()[0].name: x})[0][0]  # (N, 6)

    # YOLOv5 の出力は (cx, cy, w, h, obj_conf, cls_conf)。score = obj * cls
    score = pred[:, 4] * pred[:, 5]
    m = score > conf
    score = score[m]
    cx, cy, bw, bh = pred[m, 0], pred[m, 1], pred[m, 2], pred[m, 3]

    # letterbox 座標 -> 元画像座標
    boxes = np.stack(
        [
            (cx - bw / 2 - pad_l) / scale,
            (cy - bh / 2 - pad_t) / scale,
            (cx + bw / 2 - pad_l) / scale,
            (cy + bh / 2 - pad_t) / scale,
        ],
        axis=1,
    )

    out = image.copy()
    rows, records = [], []
    if len(score):
        keep = nms(boxes, score, iou)
        boxes, score = boxes[keep], score[keep]
        W, H = image.size
        boxes[:, 0::2] = boxes[:, 0::2].clip(0, W)
        boxes[:, 1::2] = boxes[:, 1::2].clip(0, H)

        draw = ImageDraw.Draw(out)
        width = max(2, round(max(W, H) / 400))
        for i, ((x1, y1, x2, y2), s) in enumerate(zip(boxes, score), start=1):
            draw.rectangle([x1, y1, x2, y2], outline="#e2483d", width=width)
            draw.text((x1, max(0, y1 - 14 * width / 2)), f"{i} kao {s:.2f}", fill="#e2483d")
            rows.append([i, round(float(s), 4), round(float(x1)), round(float(y1)),
                         round(float(x2)), round(float(y2))])
            records.append({"id": i, "label": "kao", "confidence": round(float(s), 4),
                            "xmin": round(float(x1)), "ymin": round(float(y1)),
                            "xmax": round(float(x2)), "ymax": round(float(y2))})

    return out, rows, {"count": len(records), "detections": records}


TITLE = "花押検出 / Kaō Detection (YOLOv5x)"

DESCRIPTION = """
古文書の画像から**花押**（かおう・図案化された自署）の位置を矩形で検出します。
Detects *kaō* — stylized personal monogram signatures — in pre-modern Japanese documents.

モデル: [`nakamura196/yolov5-kao`](https://huggingface.co/nakamura196/yolov5-kao)
（東京大学史料編纂所「花押データベース」で学習）

⚠️ 検出のみで、**誰の花押かの同定は行いません**。
草名（自筆の草書署名）や影字（紙背からの裏写り）を誤検出する系統的な傾向が確認されています。
`conf` を 0.9 以上に上げると誤検出率は約 8% まで下がります。

⏳ CPU 推論のため 1 枚あたり十数秒かかります。初回はモデル取得（346MB）で追加の待ち時間があります。
"""

with gr.Blocks(title=TITLE) as demo:
    gr.Markdown(f"# {TITLE}\n{DESCRIPTION}")

    with gr.Row():
        with gr.Column():
            inp = gr.Image(type="pil", label="入力画像 / Input image", sources=["upload", "clipboard"])
            conf = gr.Slider(0.05, 0.95, value=0.25, step=0.05,
                             label="信頼度しきい値 / Confidence threshold")
            iou = gr.Slider(0.1, 0.9, value=0.45, step=0.05,
                            label="NMS IoU しきい値 / NMS IoU threshold")
            run = gr.Button("検出する / Detect", variant="primary")
        with gr.Column():
            out_img = gr.Image(type="pil", label="検出結果 / Detections")
            out_tbl = gr.Dataframe(
                headers=["#", "conf", "xmin", "ymin", "xmax", "ymax"],
                label="検出一覧 / Detected boxes", wrap=True,
            )
            out_json = gr.JSON(label="JSON")

    gr.Examples(examples=[["iriki.jpg", 0.25, 0.45]], inputs=[inp, conf, iou])

    gr.Markdown(
        "**例示画像の出典 / Source of the example image**: "
        "「入来院家文書」後醍醐天皇綸旨（元弘3年11月9日）／請求記号 0671-18-1 — "
        "**東京大学史料編纂所所蔵**（画像: [clioimg.hi.u-tokyo.ac.jp](https://clioimg.hi.u-tokyo.ac.jp/)、"
        "利用条件: [CC BY 相当](https://www.hi.u-tokyo.ac.jp/faq/reuse_cc-by/)）。\n\n"
        "検出結果として表示される矩形は本アプリが描画したもの（改変）であり、"
        "東京大学史料編纂所が作成・公開したものではありません。\n"
        "*The red boxes are drawn by this application (a modification); the annotated image is not "
        "produced or published by the Historiographical Institute.*"
    )

    run.click(detect, inputs=[inp, conf, iou], outputs=[out_img, out_tbl, out_json])

if __name__ == "__main__":
    # Gradio 6 では theme は launch() 側で指定する
    demo.launch(theme=gr.themes.Soft())