forked from codingforentrepreneurs/OpenCV-REST-API
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathutils.py
More file actions
30 lines (26 loc) · 1.16 KB
/
utils.py
File metadata and controls
30 lines (26 loc) · 1.16 KB
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
from api import OUTPUT_DIR
import os
import cv2
def get_cascade(cascade_xml= 'haarcascade_frontalface_default.xml'):
cascade_path = os.path.join(cv2.data.haarcascades, cascade_xml)
cascade = cv2.CascadeClassifier(cascade_path)
return cascade
def extract(img_path):
cascade = get_cascade()
filename, _ = os.path.splitext(os.path.basename(img_path))
final_output_dir = os.path.join(OUTPUT_DIR, f"{filename}")
faces_output_dir = os.path.join(final_output_dir, "faces")
os.makedirs(final_output_dir, exist_ok=True)
os.makedirs(faces_output_dir, exist_ok=True)
frame = cv2.imread(img_path)
og_frame = cv2.imread(img_path)
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
faces = cascade.detectMultiScale(gray, scaleFactor=1.5, minNeighbors=3)
for i, (x, y, w, h) in enumerate(faces):
roi_color = frame[y:y+h, x:x+w] # y: y+height, x: x+width
path = os.path.join(faces_output_dir, f"{i}.jpg")
color = (255, 0, 0) # BGR
cv2.rectangle(og_frame, (x, y), (x+w, y+h), color, 2)
cv2.imwrite(path, roi_color)
cv2.imwrite(os.path.join(final_output_dir, 'main.jpg'), og_frame)
return final_output_dir