187 lines
6.4 KiB
Python
187 lines
6.4 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
|
|
DISPLAY = False
|
|
|
|
|
|
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": ""},
|
|
},
|
|
"display": True
|
|
}
|
|
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"])
|
|
if DISPLAY:
|
|
config["DISPLAY"] = DISPLAY
|
|
else:
|
|
DISPLAY = config["display"]
|
|
|
|
print(f"Current config: {config}")
|
|
|
|
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 if DISPLAY is set to true
|
|
if DISPLAY:
|
|
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()
|