Spaces:
Sleeping
Sleeping
File size: 4,475 Bytes
4794a48 24afedb 4794a48 | 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 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import cv2
from pinecone import Pinecone
import os
from dotenv import load_dotenv , dotenv_values
import numpy as np
from mtcnn import MTCNN
import matplotlib.pyplot as plt
import scipy # for MTCNN Dependencies
from keras_facenet import FaceNet
detector = MTCNN()
embedder = FaceNet()
load_dotenv()
pc = Pinecone(os.getenv("PINECONE_KEY"))
index = pc.Index(os.getenv("pinecone_index"))
def get_embedding(file):
img = plt.imread(file)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
try:
faces = detector.detect_faces(img)
if len(faces) > 1:
return print("Multiple Faces")
else:
faces = faces[0]
x1, y1, width, height = faces['box']
x1, y1 = abs(x1), abs(y1)
img = img[y1:y1+height, x1:x1+width]
img = cv2.resize(img, (224, 224))
img = np.expand_dims(img, axis=0)
print(f"Detected Face in {file}")
return embedder.embeddings(img)[0]
except:
print(f"Failed to detect faces in {file}")
return None
def make_embeddings(file , name):
embeddings = {}
emb = get_embedding(file)
if emb is not None:
embeddings[name] = emb
return embeddings
def make_meta(file , i , name):
meta = {'image_id':f"Image_{i}"}
meta['label'] = name
return meta
def upsert_embeddings(face_embeddings , metadata):
assert len(face_embeddings) == len(metadata) , f"Dimesnion mismatch, got embeddings len {len(face_embeddings)} and metadata length {len(metadata)}"
assert face_embeddings is not None and metadata is not None , 'Input vectors cannot be None'
assert face_embeddings[0].shape[0] == index.describe_index_stats()['dimension'] , 'Dimension mismatch'
a = index.describe_index_stats()['total_vector_count'] + 1
upsert_data = [(str(i + a), face_embeddings[i].tolist(), metadata[i]) for i in range(len(face_embeddings))] # index starts from last vector in database
index.upsert(vectors=upsert_data)
return True
def fetch_embeddings(embeddings):
assert embeddings.shape[0] == index.describe_index_stats()['dimension'] , f"Expected{embeddings.shape[0]} got {index.describe_index_stats()['dimension']} instead"
assert embeddings is not None , 'Input vectors cannot be None'
out = index.query(
vector=embeddings.tolist(),
top_k=1,
include_metadata=True
)
name , confidence = out['matches'][0]['metadata']['label'] , out['matches'][0]['score']
return name , confidence
def multiple_faces(file):
imga = plt.imread(file)
#imga = cv2.cvtColor(imga, cv2.COLOR_BGR2RGB)
try:
faces = detector.detect_faces(imga)
embs = []
boxes = []
for face in faces:
x1, y1, width, height = face['box']
x1, y1 = abs(x1), abs(y1)
img = imga[y1:y1+height, x1:x1+width]
img = cv2.resize(img, (224, 224))
img = np.expand_dims(img, axis=0)
embs.append(embedder.embeddings(img)[0])
boxes.append([x1, y1, width, height])
return embs, boxes
except Exception as e:
print(f"Failed to detect faces in {file} due to {e}")
return None
def draw_boxes(file, boxes ,labels, probs):
img = plt.imread(file)
#img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.copy(img)
for box , label , prob in zip(boxes,labels, probs):
x1, y1, width, height = box
img = cv2.rectangle(img, (x1, y1), (x1+width, y1+height), (0, 255, 0), 1)
if label == 'Unknown':
img = cv2.putText(img, f"{label}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 0, 255), 1)
else:
img = cv2.putText(img, f"{label} {np.round(prob , decimals = 2)}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 255, 0), 1)
#img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img
def write_and_upsert(dir , Name, upsert = True):
embeddings = make_embeddings(dir , Name)
metadata = [make_meta(dir , 1 , Name)]
if upsert:
upsert_embeddings(list(embeddings.values()), metadata)
return True
def detect_and_fetch(dir , directory = False , min_confidence = 0):
files = []
files.append(dir)
for file in files:
try:
embs , boxes = multiple_faces(file)
labels = []
probs = []
for emb in embs:
name , confidence = fetch_embeddings(emb)
if confidence > min_confidence:
labels.append(name)
probs.append(confidence)
else:
labels.append('Unknown')
probs.append('unk')
img = draw_boxes(file, boxes, labels, probs)
except Exception as e:
print(f"Failed to detect faces in {file} due to {e}")
return img , labels
|