Spaces:
Running
Running
| #!/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 encode_pil_to_bytes(img: Image.Image, format="WEBP", fallback_format="PNG") -> bytes: | |
| output_bytes = io.BytesIO() | |
| try: | |
| img.save(output_bytes, format=format) | |
| except Exception as e: | |
| print(f"WebP save failed: {e}, falling back to {fallback_format}") | |
| img.save(output_bytes, format=fallback_format) | |
| return output_bytes.getvalue() | |
| def save_image_with_fallback(img: Image.Image, cache_dir: pathlib.Path, fallback_format="PNG") -> pathlib.Path: | |
| try: | |
| return gr.processing_utils.save_pil_to_cache(img, cache_dir, format="WEBP") | |
| except Exception as e: | |
| print(f"WebP save failed: {e}, falling back to {fallback_format}") | |
| return gr.processing_utils.save_pil_to_cache(img, cache_dir, format=fallback_format) | |
| detector = load_model() | |
| detector.prepare(-1, nms_thresh=0.5, input_size=(640, 640)) | |
| def detect(image: np.ndarray) -> list[np.ndarray]: | |
| image = image[:, :, ::-1] # RGB -> BGR | |
| bboxes = detect_person(image, detector) | |
| person_images = extract_persons(image, bboxes) # Extract each person as a separate image | |
| return [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", postprocess=save_image_with_fallback), # 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() | |