Spaces:
Running
Running
| from ultralyticsplus import YOLO | |
| from PIL import Image | |
| import numpy as np | |
| from swin import load_model, get_embeddings | |
| import matplotlib.pyplot as plt | |
| import sqlite3 | |
| import pathlib | |
| def load_detector(): | |
| # load model | |
| model = YOLO('https://github.com/akanametov/yolov8-face/releases/download/v0.0.0/yolov8n-face.pt') | |
| # set model parameters | |
| model.overrides['conf'] = 0.25 # NMS confidence threshold | |
| model.overrides['iou'] = 0.45 # NMS IoU threshold | |
| model.overrides['agnostic_nms'] = False # NMS class-agnostic | |
| model.overrides['max_det'] = 50 # maximum number of detections per image | |
| return model | |
| def extract_faces(model, image): | |
| # perform inference | |
| results = model.predict(image) | |
| ids = np.array(results[0].boxes.xyxy).astype(np.int32) | |
| img = Image.open(image) | |
| crops = [] | |
| for id in ids: | |
| crops.append(Image.fromarray(np.array(img)[id[1] : id[3], id[0]: id[2]])) | |
| return crops, np.argmax(np.array(results[0].boxes.conf)) | |
| def findCosineDistance(source_representation, test_representation): | |
| a = np.matmul(np.transpose(source_representation), test_representation) | |
| b = np.sum(np.multiply(source_representation, source_representation)) | |
| c = np.sum(np.multiply(test_representation, test_representation)) | |
| return 1 - (a / (np.sqrt(b) * np.sqrt(c))) | |
| def recognition(imgs, ids, names, source_faces, d, source_imgs): | |
| cols = 4 | |
| rows = int(np.ceil(len(imgs)/2)) | |
| img_count = 0 | |
| fig, axes = plt.subplots(nrows=rows, ncols=cols, figsize=(15,rows*3)) | |
| for i in range(rows): | |
| for j in range(cols): | |
| if img_count < len(imgs): | |
| if(j%2): | |
| axes[i, j].set_title(f"Confidence: {1 - d[ids[img_count]][img_count]: .2f}") | |
| axes[i, j].imshow(imgs[img_count]) | |
| axes[i, j].set_axis_off() | |
| img_count+=1 | |
| else: | |
| if names[img_count] == "Unknown": | |
| axes[i, j].set_title(f"Unknown") | |
| else: | |
| axes[i, j].set_title(f"Roll No.:{source_imgs[ids[img_count]].split('/')[-1].split('.')[0]}") | |
| axes[i, j].imshow(source_faces[ids[img_count]]) | |
| axes[i, j].set_axis_off() | |
| plt.savefig("Recognition.jpg") | |
| def create(c): | |
| pathlib.Path("images").mkdir(parents=True, exist_ok=True) | |
| pathlib.Path("embeds").mkdir(parents=True, exist_ok=True) | |
| c.execute('''CREATE TABLE IF NOT EXISTS students | |
| (name text, email text, roll_no text, image text, embeddings text)''') | |
| c.execute('''CREATE TABLE IF NOT EXISTS attendance | |
| (date text, roll_no text, present text)''') | |
| def insert(c, name, email, roll_no, image, embeddings): | |
| c.execute("INSERT INTO students VALUES (?, ?, ?, ?, ?)", (name, email, roll_no, image, embeddings)) |