predictor

.gitignore

@@ -0,0 +1,2 @@
+gradio_cached_examples/
+flagged/

app.py

@@ -1,48 +1,69 @@
 #!/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 cv2
 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
 
-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',
+def load_models():
+    detector_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
+    age_gender_path = huggingface_hub.hf_hub_download('iitolstykh/demo_xnet_volo_cross',
+                                           'checkpoint-377.pth.tar',
+                                           use_auth_token=HF_TOKEN)
 
-def detect(image: np.ndarray, detector: YOLO) -> np.ndarray:
-    # input is rgb image, output must be rgb too
+    predictor_cfg = Cfg(detector_path, age_gender_path)
+    predictor = Predictor(predictor_cfg)
 
-    detector_kwargs = {'conf': 0.5, 'iou': 0.5, 'half': False, 'verbose': False}
+    return predictor
 
-    image = image[:, :, ::-1]  # RGB -> BGR
-    results: Results = detector.predict(image, **detector_kwargs)[0]
-    out_im = results.plot()
+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
 
 
-detector = load_model()
+predictor = load_models()
 
 image_dir = pathlib.Path('images')
 examples = [[path.as_posix()] for path in sorted(image_dir.glob('*.jpg'))]
 
-func = functools.partial(detect, detector=detector)
+func = functools.partial(detect, predictor=predictor)
 
 gr.Interface(
     fn=func,