#!/usr/bin/env python

import os
import shlex
import subprocess

if os.getenv('SYSTEM') == 'spaces':
    GITHUB_TOKEN = os.getenv('GITHUB_TOKEN')
    GITHUB_USER = os.getenv('GITHUB_USER')
    git_repo = f"https://{GITHUB_TOKEN}@github.com/{GITHUB_USER}/xnet_demo.git"
    subprocess.call(shlex.split(f'pip install git+{git_repo}'))

import pathlib
import os
import gradio as gr
import huggingface_hub
import numpy as np
import functools
from dataclasses import dataclass

from xnet.predictor import Predictor


@dataclass
class Cfg:
    detector_weights: str
    checkpoint: str
    device: str = "cpu"
    with_persons: bool = True
    disable_faces: bool = False
    draw: bool = True


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_models():
    detector_path = huggingface_hub.hf_hub_download('iitolstykh/demo_yolov8_detector',
                                           'yolov8x_person_face.pt',
                                           use_auth_token=HF_TOKEN)

    age_gender_path = huggingface_hub.hf_hub_download('iitolstykh/demo_xnet_volo_cross',
                                           'checkpoint-377.pth.tar',
                                           use_auth_token=HF_TOKEN)

    predictor_cfg = Cfg(detector_path, age_gender_path)
    predictor = Predictor(predictor_cfg)

    return predictor

def detect(image: np.ndarray, predictor: Predictor) -> np.ndarray:
    # input is rgb image, output must be rgb too

    image = image[:, :, ::-1]  # RGB -> BGR
    detected_objects, out_im = predictor.recognize(image)
    return out_im[:, :, ::-1]  # BGR -> RGB


predictor = load_models()

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

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

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)