from ultralyticsplus import YOLO
from PIL import Image
import numpy as np
from swin import load_model, get_embeddings
import matplotlib.pyplot as plt

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:
              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")