Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,386 Bytes
c817722 7c7c7c9 c817722 6627736 c817722 6627736 7c7c7c9 c817722 0324d3b c817722 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
from ultralytics import YOLO
import os
import requests
from PIL import Image
from functools import lru_cache
import torch
MODEL_LINK = (
"https://github.com/akanametov/yolov8-face/releases/download/v0.0.0/yolov8l-face.pt"
)
MODEL_PATH = "yolov8l-face.pt"
def check_and_download_model():
if os.path.exists(MODEL_PATH):
return MODEL_PATH
response = requests.get(MODEL_LINK, timeout=60)
with open(MODEL_PATH, "wb") as f:
f.write(response.content)
return MODEL_PATH
model = YOLO(check_and_download_model())
device = "cuda:0" if torch.cuda.is_available() else "cpu"
model.to(device)
def adjust_box(b, image_shape, margin=10):
x1, y1, x2, y2 = b
x1 = max(0, x1 - margin)
y1 = max(0, y1 - margin)
x2 = min(image_shape[0], x2 + margin)
y2 = min(image_shape[1], y2 + margin)
return (int(x1), int(y1), int(x2), int(y2))
@lru_cache(maxsize=None)
def predict(path):
results = model.predict(path)
orimg = Image.open(path)
return results, orimg
def detect_faces(image, box_margin=0):
faces = []
results, orimg = predict(image)
for box in results[0].boxes:
coords = box.xyxy[0]
coords = adjust_box(coords, orimg.size, margin=box_margin)
crop = orimg.crop(coords)
prob = float(box.conf * 100)
prob = f"{prob:.2f}%"
faces.append((crop, prob))
return faces
|