"""
Gradio demo for TypoRef YOLOv8 Historical Document Detector.

This script defines a simple Gradio interface that allows a user to upload
an image of a historical document page.  The interface loads a YOLOv8
object detection model and applies it to the input image, overlaying
bounding boxes around detected ornaments, typography and other decorative
elements.  The resulting annotated image is returned for display in the
browser.

By default the demo uses a small pretrained YOLOv8 model from the
``ultralytics`` repository (``yolov8n``) so that it can run without any
custom weights.  If you have uploaded your own fine‑tuned weights to
Hugging Face (e.g. ``martinbadrous/TypoRef-YOLOv8-Historical-Document-Detection``)
you can replace the ``model_path`` value below with the repository ID
for your model.  The ``ultralytics`` package will automatically load
the model from Hugging Face when given a repo ID.

To launch the demo locally run ``python app.py``.  When running as a
Hugging Face Space this file will be executed automatically.
"""

import gradio as gr
from PIL import Image
from ultralytics import YOLO


def load_model(model_path: str = "ultralytics/yolov8n.pt") -> YOLO:
    """Load a YOLOv8 model from a given path or Hugging Face repo.

    Args:
        model_path: Either a local path to a ``.pt`` file or a Hugging Face
            model repo ID.  Defaults to the ultralytics YOLOv8 nano model.

    Returns:
        An instance of ``ultralytics.YOLO`` ready for inference.
    """
    model = YOLO(model_path)
    return model


def detect_objects(img: Image.Image, model: YOLO) -> Image.Image:
    """Run object detection on a single image and return a plotted result.

    The YOLOv8 model returns a list of ``Ultralytics.Results`` objects.  The
    first element contains the detections for the provided image.  The
    ``plot`` method draws the bounding boxes and class labels onto a
    numpy array.  The array is then converted back into a PIL image for
    display in the Gradio interface.

    Args:
        img: Input PIL image of a document page.
        model: A loaded YOLOv8 model.

    Returns:
        PIL image with detection boxes overlaid.
    """
    results = model(img)
    res = results[0]
    plotted = res.plot()
    return Image.fromarray(plotted)


def build_interface() -> gr.Interface:
    """Construct and return the Gradio interface for the detection demo."""
    model = load_model("ultralytics/yolov8n.pt")

    def _predict(image: Image.Image) -> Image.Image:
        return detect_objects(image, model)

    title = "TypoRef YOLOv8: Historical Document Ornament Detection"
    description = (
        "Upload a scanned page from a historical book to see how a YOLOv8 model "
        "detects graphical ornaments, typography and decorations.  This demo "
        "illustrates cultural heritage AI applied to the TypoRef dataset (16th–18th "
        "century prints).  Replace the underlying model with your own fine‑tuned "
        "weights by modifying the `model_path` in `app.py`."
    )

    iface = gr.Interface(
        fn=_predict,
        inputs=gr.Image(type="pil", label="Upload a document page"),
        outputs=gr.Image(type="pil", label="Detected ornaments & typography"),
        title=title,
        description=description,
        allow_flagging="never",
    )
    return iface


if __name__ == "__main__":
    interface = build_interface()
    interface.launch()