#!/usr/bin/env python

import pathlib
import os
import cv2
import gradio as gr
import huggingface_hub
import numpy as np
import functools

from ultralytics import YOLO
from ultralytics.yolo.engine.results import Results

TITLE = 'Age and Gender Estimation with Transformers from Face and Body Images in the Wild'
DESCRIPTION = 'This is an official demo for https://github.com/...'

HF_TOKEN = os.getenv('HF_TOKEN')

def load_model():
    path = huggingface_hub.hf_hub_download('iitolstykh/demo_yolov8_detector',
                                           'yolov8x_person_face.pt',
                                           use_auth_token=HF_TOKEN)
    yolo = YOLO(path)
    yolo.fuse()

    return yolo

def detect(image: np.ndarray, detector: YOLO) -> np.ndarray:
    detector_kwargs = {'conf': 0.5, 'iou': 0.5, 'half': False, 'verbose': False}

    results: Results = detector.predict(image, **detector_kwargs)[0]
    out_im = results.plot()

    return out_im


detector = load_model()

image_dir = pathlib.Path('images')
examples = [[path.as_posix()] for path in sorted(image_dir.glob('*.jpg'))]

func = functools.partial(detect, detector=detector)

gr.Interface(
    fn=func,
    inputs=gr.Image(label='Input', type='numpy'),
    outputs=gr.Image(label='Output', type='numpy'),
    examples=examples,
    examples_per_page=30,
    title=TITLE,
    description=DESCRIPTION,
).launch(show_api=False)