diff --git a/.gitignore b/.gitignore index 8699180..c00d8ea 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,8 @@ data/* !data/sample_images !data/sample_test +reports/ + __pycache__ # Under work diff --git a/counter.py b/counter.py new file mode 100644 index 0000000..fa10a8b --- /dev/null +++ b/counter.py @@ -0,0 +1,42 @@ +import os + +from sympy import Id + +class IDCounter(object): + def __init__(self, threshold, ID, sessionId, path): + self.count_thresh = threshold + self.names = [f"{entity['first']} {entity['last']}" for entity in ID] + self.counters = {i:0 for i in self.names} + self.attendance = {i:False for i in self.names} + + self.sessionId = sessionId + self.path = path + with open(path, 'wb') as f: + f.write(f"Attendance report for {sessionId}\n\nRegistered:\n".encode("utf-8")) + + def write(self, val): + with open(self.path, 'ab') as f: + f.write(val.encode("utf-8")) + + def update(self, id): + if id in self.counters: + self.counters[id] += 1 + if self.counters[id] >= self.count_thresh: + if self.attendance[id] == False: + self.attendance[id] = True + self.write(f"✔\t{id}\n") + return True + else: + self.counters[id] = 1 + return False + + def showReport(self): + append = f"\nNot registered:\n" + for person in self.attendance: + if self.attendance[person] == False: + append += f"❌\t{person}\n" + append += '\n\n[TECLARS ATTENDANCE REPORT COMPLETE]\n' + + self.write(append) + + os.system(f"notepad.exe {self.path}") diff --git a/get_embed.py b/embed.py similarity index 51% rename from get_embed.py rename to embed.py index d78c59a..f26f1a2 100644 --- a/get_embed.py +++ b/embed.py @@ -1,14 +1,30 @@ from facenet_pytorch import MTCNN, InceptionResnetV1 from tqdm import tqdm from PIL import Image +import argparse import torch import os import json -with open('data\id.json', 'r') as f: +parser = argparse.ArgumentParser(description='TECLARS Embedding Generator') +parser.add_argument('-i', '--id_path', type=str, default='./data/id.json', + help='Path of ID data JSON file') +parser.add_argument('-p', '--img_dir', type=str, default='./data/images', + help='Directory of where image data is stored') +parser.add_argument('-o', '--out_path', type=str, default='./data/embeddings.pt', + help='Path of ID data JSON file') +parser.add_argument('-d', '--device', type=str, default='auto', + help='Device to compute algorithm on') + +args = parser.parse_args() + +with open(args.id_path, 'r') as f: ID = json.load(f) -device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') +if args.device == "auto": + device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') +else: + device = torch.device(args.device) print('Running on device: {}'.format(device)) print("Aligning for faces from data") @@ -17,7 +33,7 @@ aligned = [] for idx, person in tqdm(enumerate(ID)): - fpath = os.path.join('./data/images', person['image']) + fpath = os.path.join(args.img_dir, person['image']) with Image.open(fpath) as x: x_aligned, prob = mtcnn(x, return_prob=True) @@ -34,6 +50,6 @@ print("Computing face embeddings") aligned = torch.stack(aligned).to(device) embeddings = resnet(aligned).detach().cpu() -torch.save({'embedding': embeddings}, 'data/embeddings.pt') +torch.save({'embedding': embeddings}, args.out_path) print("\n[All done]") diff --git a/run.py b/run.py index 5fe1381..97b4cf9 100644 --- a/run.py +++ b/run.py @@ -1,3 +1,4 @@ +from genericpath import exists from facenet_pytorch import MTCNN, InceptionResnetV1 import torch import torch.nn as nn @@ -9,6 +10,9 @@ import time import argparse +from counter import IDCounter + + def euclidean_distance(out, refs): return (out - refs).norm(dim=1) @@ -75,6 +79,15 @@ def identify_faces(pil_image): return None def video(): + if args.session_id is None: + sess_id = "sheets-bill-" + input("Enter session ID: sheets-bill-") + else: + sess_id = "sheets-bill-" + args.session_id + + os.makedirs(args.output_dir, exist_ok=True) + report_path = os.path.join(args.output_dir, f"{sess_id}.txt") + counter = IDCounter(args.number, ID, sess_id, report_path) + print("Starting video capture\n") video_capture = cv2.VideoCapture(args.camera, cv2.CAP_DSHOW) @@ -103,13 +116,17 @@ def video(): if args.show_unrecognised and entity is None: cv2.rectangle(frame, (bounds[0], bounds[1]), (bounds[2], bounds[3]), (0, 0, 255), 2) else: - text = f"{entity['first']} {entity['last']}: {round(confidence*100, 2)}%" - cv2.putText(frame, text, (bounds[0], bounds[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) - cv2.rectangle(frame, (bounds[0], bounds[1]), (bounds[2], bounds[3]), (0, 255, 0), 2) - - if entity not in registered_students: - print(f"{entity['first']} {entity['last']} registered") - registered_students.append(entity) + name = f"{entity['first']} {entity['last']}" + if counter.update(name): + if entity not in registered_students: + print(f"{entity['first']} {entity['last']} registered") + registered_students.append(entity) + color = (0, 255, 0) + else: + color = (0, 210, 255) + text = f"{name}: {round(confidence*100, 2)}%" + cv2.putText(frame, text, (bounds[0], bounds[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) + cv2.rectangle(frame, (bounds[0], bounds[1]), (bounds[2], bounds[3]), color, 2) cv2.imshow('TECLARS Main UI', frame) @@ -128,6 +145,8 @@ def video(): for e, entity in enumerate(registered_students) ])) + counter.showReport() + def test(): print("Beginning testing") for filename in os.listdir(args.test_path): @@ -188,6 +207,8 @@ def check(): parser.add_argument('-r', '--threshold', type=float, default=0.8, help='Probability above which a face will be considered recognised') +parser.add_argument('-n', '--number', type=int, default=5, + help='Minimum number of frames above which a face will be considered recognised') parser.add_argument('-g', '--margin', type=float, default=0.1, help='Minimum probability margin above next likely face for the face to be considered recognised') parser.add_argument('-t', '--temp', type=float, default=2, @@ -198,10 +219,14 @@ def check(): help='Device to compute algorithm on') parser.add_argument('-m', '--mode', type=str, choices=['cosine', 'euclidean'], default='cosine', help='Distance function for evaluating the similarity between face embeddings') -parser.add_argument('-u', '--show_unrecognised', action="store_true", +parser.add_argument('-u', '--show_unrecognised', action="store_false", help='Remove bounding boxes around unrecognised faces') parser.add_argument('-i', '--ignore', type=int, default=100, help='Ignore faraway faces with a width smaller than this value (set 0 to include all faces)') +parser.add_argument('-s', '--session_id', type=str, default=None, + help="Session ID") +parser.add_argument('-o', '--output_dir', type=str, default="./reports", + help="Directory to output reports") parser.add_argument('-x', '--dev', action='store_true', help="Enable developer options") diff --git a/session.py b/session.py new file mode 100644 index 0000000..e69de29