Spaces:
Running
Running
JANGALA SAKETH commited on
Create app.py
Browse files
app.py
ADDED
|
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import numpy as np
|
| 2 |
+
import cv2
|
| 3 |
+
from fastapi import FastAPI, Request
|
| 4 |
+
from fastapi.responses import JSONResponse
|
| 5 |
+
from ultralytics import YOLO
|
| 6 |
+
import insightface
|
| 7 |
+
|
| 8 |
+
app = FastAPI()
|
| 9 |
+
|
| 10 |
+
# ----------------------------
|
| 11 |
+
# Load Models (CPU mode)
|
| 12 |
+
# ----------------------------
|
| 13 |
+
|
| 14 |
+
yolo = YOLO("yolov8n.pt")
|
| 15 |
+
|
| 16 |
+
face_model = insightface.app.FaceAnalysis(name="buffalo_l")
|
| 17 |
+
face_model.prepare(ctx_id=-1)
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def normalize(vec):
|
| 21 |
+
vec = np.array(vec, dtype=np.float32)
|
| 22 |
+
norm = np.linalg.norm(vec)
|
| 23 |
+
if norm == 0:
|
| 24 |
+
return vec.tolist()
|
| 25 |
+
return (vec / norm).tolist()
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
def process_image_np(image_np):
|
| 29 |
+
results = yolo(image_np)
|
| 30 |
+
faces_output = []
|
| 31 |
+
|
| 32 |
+
for r in results:
|
| 33 |
+
boxes = r.boxes
|
| 34 |
+
for box, cls, conf in zip(boxes.xyxy, boxes.cls, boxes.conf):
|
| 35 |
+
if int(cls) != 0:
|
| 36 |
+
continue
|
| 37 |
+
if float(conf) < 0.4:
|
| 38 |
+
continue
|
| 39 |
+
|
| 40 |
+
xmin, ymin, xmax, ymax = box.cpu().numpy()
|
| 41 |
+
xmin, ymin, xmax, ymax = map(int, [xmin, ymin, xmax, ymax])
|
| 42 |
+
|
| 43 |
+
h, w, _ = image_np.shape
|
| 44 |
+
xmin = max(0, xmin)
|
| 45 |
+
ymin = max(0, ymin)
|
| 46 |
+
xmax = min(w, xmax)
|
| 47 |
+
ymax = min(h, ymax)
|
| 48 |
+
|
| 49 |
+
person_crop = image_np[ymin:ymax, xmin:xmax]
|
| 50 |
+
if person_crop.size == 0:
|
| 51 |
+
continue
|
| 52 |
+
|
| 53 |
+
detected_faces = face_model.get(person_crop)
|
| 54 |
+
|
| 55 |
+
for face in detected_faces:
|
| 56 |
+
embedding = normalize(face.embedding)
|
| 57 |
+
fxmin, fymin, fxmax, fymax = face.bbox.astype(int)
|
| 58 |
+
|
| 59 |
+
faces_output.append({
|
| 60 |
+
"cx": float((fxmin + fxmax) / 2 + xmin),
|
| 61 |
+
"cy": float((fymin + fymax) / 2 + ymin),
|
| 62 |
+
"confidence": float(conf),
|
| 63 |
+
"box": {
|
| 64 |
+
"xmin": int(fxmin + xmin),
|
| 65 |
+
"ymin": int(fymin + ymin),
|
| 66 |
+
"xmax": int(fxmax + xmin),
|
| 67 |
+
"ymax": int(fymax + ymin)
|
| 68 |
+
},
|
| 69 |
+
"embedding": embedding
|
| 70 |
+
})
|
| 71 |
+
|
| 72 |
+
return faces_output
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
@app.post("/detect")
|
| 76 |
+
async def detect(request: Request):
|
| 77 |
+
body = await request.body()
|
| 78 |
+
|
| 79 |
+
np_arr = np.frombuffer(body, np.uint8)
|
| 80 |
+
image_np = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
|
| 81 |
+
|
| 82 |
+
if image_np is None:
|
| 83 |
+
return JSONResponse({"error": "Invalid image"}, status_code=400)
|
| 84 |
+
|
| 85 |
+
result = process_image_np(image_np)
|
| 86 |
+
return result
|