utils/face_rec.py

import face_recognition
import os
import numpy as np
from PIL import Image
import random
import cv2

def update_ind2person(ind2person, emb, person):
    ind2person[len(list(ind2person.values()))]=dict(person=person,emb=emb)
    print(f"dict ind2person update: {person}!!!")
    return ind2person
def input_an_image(image, person_name, ori_img_dir='images/ori_images',img_emb_dir='images/img_emb', save_ori_img=True):
    """
    args:
        image: PIL Image
        person_name: str
    """
    image_file_dir=os.path.join(ori_img_dir,person_name)
    emb_file_dir=os.path.join(img_emb_dir,person_name)
    if not os.path.exists(image_file_dir):
        os.mkdir(image_file_dir)
        os.mkdir(emb_file_dir)
        file_ind=0
    else:
        file_ind=len(os.listdir(image_file_dir))
    # file_ = face_recognition.load_image_file(image_file)
    if save_ori_img:
        image.save(os.path.join(image_file_dir,person_name+f'_{file_ind}.jpg'))
    file_ = np.array(image)
    emb = face_recognition.face_encodings(file_)[0]
    emb_file=person_name+f'_{file_ind}.npy'
    emb_file_out_path=os.path.join(emb_file_dir,emb_file)
    np.save(emb_file_out_path, emb)
    return emb

def init_load_embs(img_emb_dir='images/img_emb'):
    persons=os.listdir(img_emb_dir)
    i=0
    ind2person=dict()
    for oneperson in persons:
        oneperson_dir=os.path.join(img_emb_dir,oneperson)
        oneperson_list=os.listdir(oneperson_dir)
        for oneperson_j in oneperson_list:
            emb_id=i
            i+=1
            emb=np.load(os.path.join(oneperson_dir,oneperson_j))
            ind2person[emb_id]=dict(person=oneperson,emb=emb)
    return ind2person

def image_rec(image, known_face_encodings, _ind2person):
    """
    args:
        image: cv2 format
    return:
        image: cv2 format
    """
    # image = np.array(image)
    face_locations = face_recognition.face_locations(image)
    face_encodings = face_recognition.face_encodings(image, 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"

        # # If a match was found in known_face_encodings, just use the first one.
        # if True in matches:
        #     first_match_index = matches.index(True)
        #     name = known_face_names[first_match_index]

        # 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]:
            name = _ind2person[best_match_index]['person']
        print(f"rec {name}!!")
        face_names.append(name)
    nameset = list(set(face_names))
    colors=[(255,0,0),(0,255,0),(0,0,255),(0,255,255),(255,255,0),(156,102,31),(255,0,255)]
    chose_colors = random.sample(colors,len(nameset))
    name2color={_n:chose_colors[i] for i,_n in enumerate(nameset)}
    print(name2color)

    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 *= 4
        # right *= 4
        # bottom *= 4
        # left *= 4
        print("detect image")

        # Draw a box around the face
        # cv2.rectangle(image, (left, top), (right, bottom), (0, 0, 255), 2)
        cv2.rectangle(image, (left, top), (right, bottom), name2color[name], 2)

        # Draw a label with a name below the face
        # cv2.rectangle(image, (left, bottom - 35), (right, bottom), (0, 0, 255), cv2.FILLED)
        # cv2.rectangle(image, (left, bottom - 35), (right, bottom), name2color[name], cv2.FILLED)
        font = cv2.FONT_HERSHEY_DUPLEX
        cv2.putText(image, name, (left + 6, bottom - 6), font, 1.0, (255, 255, 255), 1)
        # cv2.imshow('image', image)
        # cv2.waitKey()
    return image