chromadb-api / app.py
Saad0KH's picture
Update app.py
190a01a verified
raw
history blame
1.04 kB
#!/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
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", "CUDAExecutionProvider"]
)
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, kpss = detector.detect(img)
bboxes = np.round(bboxes[:, :4]).astype(int)
return bboxes
def extract