51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
import cv2
|
|
import dlib
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from predict_api import ImagePredictor
|
|
|
|
# Initialize camera and face detector
|
|
cap = cv2.VideoCapture(0)
|
|
detector = dlib.get_frontal_face_detector()
|
|
|
|
# Initialize ImagePredictor
|
|
predictor = ImagePredictor(model_path="best300_model_0.7302241690286009.pth", class_indices_path="./class_indices.json")
|
|
|
|
while True:
|
|
# Capture frame-by-frame
|
|
ret, frame = cap.read()
|
|
|
|
# Convert the image from BGR color (which OpenCV uses) to RGB color
|
|
rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
|
|
# Perform face detection
|
|
faces = detector(rgb_image)
|
|
|
|
# Loop through each face in this frame
|
|
for rect in faces:
|
|
# Get the bounding box coordinates
|
|
x1, y1, x2, y2 = rect.left(), rect.top(), rect.right(), rect.bottom()
|
|
|
|
# Crop the face from the frame
|
|
face_image = rgb_image[y1:y2, x1:x2]
|
|
|
|
# Use ImagePredictor to predict the class of this face
|
|
result = predictor.predict(face_image)
|
|
|
|
# Draw a rectangle around the face
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
|
|
# Display the class name and score
|
|
cv2.putText(frame, f"{result['result'][0]['name']}: {round(result['result'][0]['score'],4)}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)
|
|
|
|
# Display the resulting frame
|
|
cv2.imshow('Video', frame)
|
|
|
|
# Exit loop if 'q' is pressed
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
|
|
# When everything is done, release the capture
|
|
cap.release()
|
|
cv2.destroyAllWindows() |