30 lines
930 B
Python
30 lines
930 B
Python
import os
|
|
from typing import List
|
|
|
|
import torch
|
|
|
|
from ultralytics import YOLO
|
|
|
|
WEIGHTS_PATH = "assets/weights/yolov8n-face.pt"
|
|
DIR_PATH = os.path.dirname(os.path.realpath(__file__))
|
|
WEIGHTS_PATH = os.path.join(DIR_PATH, WEIGHTS_PATH)
|
|
MODEL = YOLO(WEIGHTS_PATH)
|
|
CONF_THRESH = 0.01
|
|
IOU_THRESH = 0.5
|
|
|
|
|
|
# TODO: currently detect_faces accepts a image path, but it can be changed to accept images in memory
|
|
def detect_faces(image_path: str) -> List[tuple[int, int, int, int]]:
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
results = MODEL.predict(source=image_path, conf=CONF_THRESH, iou=IOU_THRESH, device=device)
|
|
face_boxes = []
|
|
result = results[0].cpu().numpy()
|
|
for box in result.boxes:
|
|
xyxyn = box.xyxy[0]
|
|
x1 = int(xyxyn[0])
|
|
y1 = int(xyxyn[1])
|
|
x2 = int(xyxyn[2])
|
|
y2 = int(xyxyn[3])
|
|
face_boxes.append((x1, y1, x2, y2))
|
|
return face_boxes
|