wyzely-detect/main.py

178 lines
6.2 KiB
Python

import datetime
import face_recognition
import cv2
import numpy as np
from dotenv import load_dotenv
import os
import json
import pathlib
load_dotenv()
URL = os.getenv("URL")
RUN_SCALE = os.getenv("RUN_SCALE")
VIEW_SCALE = os.getenv("VIEW_SCALE")
DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S"
# RUN_SCALE = 0.25
# VIEW_SCALE = 0.75
def find_face_from_name(name):
for face in config["faces"]:
if config["faces"][face]["name"] == name:
return face
return None
def write_config():
with open(config_path, "w") as config_file:
json.dump(config, config_file, indent=4)
# Initialize some variables
face_locations = []
face_encodings = []
face_names = []
known_face_encodings = []
known_face_names = []
process_this_frame = True
# Load the config file, if it does not exist or is blank, create it
config = {
"URL": "rtsp://localhost:8554/example",
"RUN_SCALE": "0.25",
"VIEW_SCALE": "0.75",
"faces": {
"example1": {"image": "example1.jpg", "last_seen": ""},
"example2": {"image": "example2.jpg", "last_seen": ""},
},
}
config_path = pathlib.Path("config.json")
if config_path.exists():
with open(config_path, "r") as config_file:
config = json.load(config_file)
else:
with open(config_path, "w") as config_file:
json.dump(config, config_file, indent=4)
print("Config file created, please edit it and restart the program")
exit()
if URL:
config["URL"] = URL
else:
URL = config["URL"]
if RUN_SCALE:
config["RUN_SCALE"] = RUN_SCALE
else:
RUN_SCALE = float(config["RUN_SCALE"])
if VIEW_SCALE:
config["VIEW_SCALE"] = VIEW_SCALE
else:
VIEW_SCALE = float(config["VIEW_SCALE"])
for face in config["faces"]:
# Load a sample picture and learn how to recognize it.
image = face_recognition.load_image_file(config["faces"][face]["image"])
face_encoding = face_recognition.face_encodings(image)[0]
known_face_encodings.append(face_encoding)
# Append the key to the list of known face names
known_face_names.append(face)
video_capture = cv2.VideoCapture(URL)
# Eliminate lag by setting the buffer size to 1
# This makes it so that the video capture will only grab the most recent frame
# However, this means that the video may be choppy
video_capture.set(cv2.CAP_PROP_BUFFERSIZE, 1)
while True:
# Grab a single frame of video
ret, frame = video_capture.read()
# Only process every other frame of video to save time
# if process_this_frame:
if True:
# Resize frame of video to a smaller size for faster face recognition processing
run_frame = cv2.resize(frame, (0, 0), fx=RUN_SCALE, fy=RUN_SCALE)
view_frame = cv2.resize(frame, (0, 0), fx=VIEW_SCALE, fy=VIEW_SCALE)
# Convert the image from BGR color (which OpenCV uses) to RGB color (which face_recognition uses)
rgb_run_frame = run_frame[:, :, ::-1]
# Find all the faces and face encodings in the current frame of video
# model cnn is gpu accelerated, but hog is cpu only
face_locations = face_recognition.face_locations(rgb_run_frame, model="hog")
face_encodings = face_recognition.face_encodings(rgb_run_frame, face_locations)
face_names = []
for face_encoding in face_encodings:
# See if the face is a match for the known face(s)
matches = face_recognition.compare_faces(
known_face_encodings, face_encoding
)
name = "Unknown"
# Or instead, use the known face with the smallest distance to the new face
face_distances = face_recognition.face_distance(
known_face_encodings, face_encoding
)
best_match_index = np.argmin(face_distances)
if matches[best_match_index]:
# print("For debugging, I found a face!!!! :D this should not be included in the final product lol :P")
name = known_face_names[best_match_index]
last_seen = config["faces"][name]["last_seen"]
# If it's never been seen, set the last seen time to six seconds ago so it will be seen
# Kind of a hacky way to do it, but it works... hopefully
if last_seen == "":
print(f"{name} has been seen")
config["faces"][name]["last_seen"] = (
datetime.datetime.now() - datetime.timedelta(seconds=6)
).strftime(DATETIME_FORMAT)
write_config()
# Check if the face has been seen in the last 5 seconds
if datetime.datetime.now() - datetime.datetime.strptime(
last_seen, DATETIME_FORMAT
) > datetime.timedelta(seconds=5):
print(f"{name} has been seen")
# Update the last seen time
config["faces"][name]["last_seen"] = datetime.datetime.now().strftime(
DATETIME_FORMAT
)
write_config()
face_names.append(name)
process_this_frame = not process_this_frame
# Display the results
# Iterate over each face found in the frame to draw a box around it
# Zip is used to iterate over two lists at the same time
for (top, right, bottom, left), name in zip(face_locations, face_names):
# Scale back up face locations since the frame we detected in was scaled to 1/4 size
top = int(top * (VIEW_SCALE / RUN_SCALE))
right = int(right * (VIEW_SCALE / RUN_SCALE))
bottom = int(bottom * (VIEW_SCALE / RUN_SCALE))
left = int(left * (VIEW_SCALE / RUN_SCALE))
# Draw a box around the face
cv2.rectangle(view_frame, (left, top), (right, bottom), (0, 0, 255), 2)
# Draw a label with a name below the face
cv2.rectangle(
view_frame, (left, bottom - 35), (right, bottom), (0, 0, 255), cv2.FILLED
)
font = cv2.FONT_HERSHEY_DUPLEX
cv2.putText(
view_frame, name, (left + 6, bottom - 6), font, 1.0, (255, 255, 255), 1
)
# Display the resulting image
cv2.imshow("Scaled View", view_frame)
# Hit 'q' on the keyboard to quit!
if cv2.waitKey(1) & 0xFF == ord("q"):
break
# Release handle to the webcam
video_capture.release()
cv2.destroyAllWindows()