#!/usr/bin/env python

from __future__ import annotations
import pathlib
import cv2
import gradio as gr
import huggingface_hub
import insightface
import numpy as np
import onnxruntime as ort
from PIL import Image
import io

TITLE = "insightface Person Detection"
DESCRIPTION = "https://github.com/deepinsight/insightface/tree/master/examples/person_detection"


def load_model():
    path = huggingface_hub.hf_hub_download("public-data/insightface", "models/scrfd_person_2.5g.onnx")
    options = ort.SessionOptions()
    options.intra_op_num_threads = 8
    options.inter_op_num_threads = 8
    session = ort.InferenceSession(
        path, sess_options=options, providers=["CPUExecutionProvider"]
    )
    model = insightface.model_zoo.retinaface.RetinaFace(model_file=path, session=session)
    return model


def detect_person(
    img: np.ndarray, detector: insightface.model_zoo.retinaface.RetinaFace
) -> np.ndarray:
    bboxes, _ = detector.detect(img)
    bboxes = np.round(bboxes[:, :4]).astype(int)
    return bboxes


def extract_persons(image: np.ndarray, bboxes: np.ndarray) -> list[np.ndarray]:
    person_images = []
    for bbox in bboxes:
        x1, y1, x2, y2 = bbox
        person_image = image[y1:y2, x1:x2]  # Crop the detected person
        person_images.append(person_image)
    return person_images


def convert_to_png(image: np.ndarray) -> bytes:
    """Convert a NumPy image array to a PNG byte stream."""
    pil_image = Image.fromarray(image)
    buffer = io.BytesIO()
    pil_image.save(buffer, format="PNG")
    buffer.seek(0)
    return buffer.read()


detector = load_model()
detector.prepare(-1, nms_thresh=0.5, input_size=(640, 640))


def detect(image: np.ndarray) -> list[bytes]:
    image = image[:, :, ::-1]  # RGB -> BGR
    bboxes = detect_person(image, detector)
    person_images = extract_persons(image, bboxes)  # Extract each person as a separate image
    return [convert_to_png(person_img[:, :, ::-1]) for person_img in person_images]  # BGR -> RGB


examples = sorted(pathlib.Path("images").glob("*.jpg"))

demo = gr.Interface(
    fn=detect,
    inputs=gr.Image(label="Input", type="numpy"),
    outputs=gr.Gallery(label="Detected Persons"),  # Display multiple images in a gallery
    examples=examples,
    cache_examples=False,  # Disable caching of examples
    examples_per_page=30,
    title=TITLE,
    description=DESCRIPTION,
)

if __name__ == "__main__":
    demo.queue(max_size=10).launch()