Spaces:
Runtime error
Runtime error
File size: 4,130 Bytes
f5f3b58 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
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 |