Compare commits
2 Commits
687372a429
...
b1f2df819c
Author | SHA1 | Date | |
---|---|---|---|
b1f2df819c | |||
246fb5a389 |
BIN
ml/assets/weights/yolov8n-face.pt
Normal file
BIN
ml/assets/weights/yolov8n-face.pt
Normal file
Binary file not shown.
29
ml/face_detection.py
Normal file
29
ml/face_detection.py
Normal file
@ -0,0 +1,29 @@
|
|||||||
|
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
|
@ -1,2 +1,93 @@
|
|||||||
|
alembic==1.12.1
|
||||||
|
anyio==4.0.0
|
||||||
|
async-timeout==4.0.3
|
||||||
|
bidict==0.22.1
|
||||||
|
certifi==2023.11.17
|
||||||
|
charset-normalizer==3.3.2
|
||||||
|
click==8.1.7
|
||||||
|
cloudpickle==2.2.1
|
||||||
|
contourpy==1.2.0
|
||||||
|
cycler==0.12.1
|
||||||
|
distro==1.8.0
|
||||||
|
exceptiongroup==1.1.3
|
||||||
|
fastapi==0.96.1
|
||||||
|
filelock==3.13.1
|
||||||
|
fonttools==4.44.3
|
||||||
|
fsspec==2023.10.0
|
||||||
|
greenlet==3.0.1
|
||||||
|
gunicorn==20.1.0
|
||||||
|
h11==0.14.0
|
||||||
|
httpcore==0.17.3
|
||||||
|
httpx==0.24.1
|
||||||
|
idna==3.4
|
||||||
|
Jinja2==3.1.2
|
||||||
|
kiwisolver==1.4.5
|
||||||
|
Mako==1.3.0
|
||||||
|
markdown-it-py==3.0.0
|
||||||
|
MarkupSafe==2.1.3
|
||||||
|
matplotlib==3.8.2
|
||||||
|
mdurl==0.1.2
|
||||||
|
mpmath==1.3.0
|
||||||
|
networkx==3.2.1
|
||||||
|
numpy==1.26.2
|
||||||
|
nvidia-cublas-cu12==12.1.3.1
|
||||||
|
nvidia-cuda-cupti-cu12==12.1.105
|
||||||
|
nvidia-cuda-nvrtc-cu12==12.1.105
|
||||||
|
nvidia-cuda-runtime-cu12==12.1.105
|
||||||
|
nvidia-cudnn-cu12==8.9.2.26
|
||||||
|
nvidia-cufft-cu12==11.0.2.54
|
||||||
|
nvidia-curand-cu12==10.3.2.106
|
||||||
|
nvidia-cusolver-cu12==11.4.5.107
|
||||||
|
nvidia-cusparse-cu12==12.1.0.106
|
||||||
|
nvidia-nccl-cu12==2.18.1
|
||||||
|
nvidia-nvjitlink-cu12==12.3.101
|
||||||
|
nvidia-nvtx-cu12==12.1.105
|
||||||
|
opencv-python==4.8.1.78
|
||||||
|
packaging==23.2
|
||||||
|
pandas==2.1.3
|
||||||
|
Pillow==10.1.0
|
||||||
|
pipdeptree==2.13.1
|
||||||
|
platformdirs==3.11.0
|
||||||
|
psutil==5.9.6
|
||||||
|
py-cpuinfo==9.0.0
|
||||||
|
pydantic==1.10.13
|
||||||
|
Pygments==2.16.1
|
||||||
|
pyparsing==3.1.1
|
||||||
|
python-dateutil==2.8.2
|
||||||
|
python-engineio==4.8.0
|
||||||
|
python-multipart==0.0.5
|
||||||
|
python-socketio==5.10.0
|
||||||
|
pytz==2023.3.post1
|
||||||
|
PyYAML==6.0.1
|
||||||
|
redis==4.6.0
|
||||||
reflex==0.3.2
|
reflex==0.3.2
|
||||||
|
requests==2.31.0
|
||||||
|
rich==13.7.0
|
||||||
|
scipy==1.11.3
|
||||||
|
seaborn==0.13.0
|
||||||
|
simple-websocket==1.0.0
|
||||||
|
six==1.16.0
|
||||||
|
sniffio==1.3.0
|
||||||
|
SQLAlchemy==1.4.41
|
||||||
|
sqlalchemy2-stubs==0.0.2a37
|
||||||
|
sqlmodel==0.0.8
|
||||||
|
starlette==0.27.0
|
||||||
|
starlette-admin==0.9.0
|
||||||
|
sympy==1.12
|
||||||
|
tabulate==0.9.0
|
||||||
|
thop==0.1.1.post2209072238
|
||||||
|
torch==2.1.1
|
||||||
|
torchvision==0.16.1
|
||||||
|
tqdm==4.66.1
|
||||||
|
triton==2.1.0
|
||||||
|
typer==0.9.0
|
||||||
|
typing_extensions==4.8.0
|
||||||
|
tzdata==2023.3
|
||||||
|
ultralytics==8.0.212
|
||||||
|
urllib3==2.1.0
|
||||||
|
uvicorn==0.20.0
|
||||||
|
watchdog==2.3.1
|
||||||
|
watchfiles==0.19.0
|
||||||
|
websockets==10.4
|
||||||
|
wrapt==1.16.0
|
||||||
|
wsproto==1.2.0
|
||||||
|
Loading…
Reference in New Issue
Block a user