Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import numpy as np | |
| from ultralytics import YOLO | |
| import insightface | |
| from fastapi import FastAPI, Request | |
| from fastapi.responses import JSONResponse | |
| import cv2 | |
| # ---------------------------- | |
| # Load models (CPU) | |
| # ---------------------------- | |
| yolo = YOLO("yolov8n.pt") | |
| face_model = insightface.app.FaceAnalysis(name="buffalo_l") | |
| face_model.prepare(ctx_id=-1) | |
| def normalize(vec): | |
| vec = np.array(vec, dtype=np.float32) | |
| norm = np.linalg.norm(vec) | |
| if norm == 0: | |
| return vec.tolist() | |
| return (vec / norm).tolist() | |
| def process_image_np(image_np): | |
| results = yolo(image_np) | |
| faces_output = [] | |
| for r in results: | |
| boxes = r.boxes | |
| for box, cls, conf in zip(boxes.xyxy, boxes.cls, boxes.conf): | |
| if int(cls) != 0: | |
| continue | |
| if float(conf) < 0.4: | |
| continue | |
| xmin, ymin, xmax, ymax = box.cpu().numpy() | |
| xmin, ymin, xmax, ymax = map(int, [xmin, ymin, xmax, ymax]) | |
| h, w, _ = image_np.shape | |
| xmin = max(0, xmin) | |
| ymin = max(0, ymin) | |
| xmax = min(w, xmax) | |
| ymax = min(h, ymax) | |
| person_crop = image_np[ymin:ymax, xmin:xmax] | |
| if person_crop.size == 0: | |
| continue | |
| detected_faces = face_model.get(person_crop) | |
| for face in detected_faces: | |
| embedding = normalize(face.embedding) | |
| fxmin, fymin, fxmax, fymax = face.bbox.astype(int) | |
| faces_output.append({ | |
| "cx": float((fxmin + fxmax) / 2 + xmin), | |
| "cy": float((fymin + fymax) / 2 + ymin), | |
| "confidence": float(conf), | |
| "box": { | |
| "xmin": int(fxmin + xmin), | |
| "ymin": int(fymin + ymin), | |
| "xmax": int(fxmax + xmin), | |
| "ymax": int(fymax + ymin) | |
| }, | |
| "embedding": embedding | |
| }) | |
| return faces_output | |
| # ---------------------------------------- | |
| # Create Gradio Blocks | |
| # ---------------------------------------- | |
| with gr.Blocks() as demo: | |
| gr.Markdown("Face Processing API Running") | |
| # Get underlying FastAPI app from Gradio | |
| app = demo.app | |
| # ---------------------------------------- | |
| # ADD REAL REST ENDPOINT TO SAME APP | |
| # ---------------------------------------- | |
| async def detect(request: Request): | |
| body = await request.body() | |
| np_arr = np.frombuffer(body, np.uint8) | |
| image_np = cv2.imdecode(np_arr, cv2.IMREAD_COLOR) | |
| if image_np is None: | |
| return JSONResponse({"error": "Invalid image"}, status_code=400) | |
| result = process_image_np(image_np) | |
| return result | |
| # Launch | |
| demo.queue(False) | |
| demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False) |