ImanAmran's picture
Update app.py
a4c5035
raw
history blame
3.65 kB
import gradio as gr
import tensorflow as tf
import numpy as np
import cv2
import os
from scipy.spatial.distance import cosine
from tensorflow.keras.models import load_model
from tensorflow.keras import layers, Model
def create_embedding_model():
# Load the model architecture
facenet_model = load_model('facenet_keras.h5', compile=False)
# Define the embedding model using FaceNet
embedding_model = Model(inputs=facenet_model.input,
outputs=facenet_model.layers[-2].output,
name="Embedding")
# Here, you don't need to set layers as trainable since it's already done in your training script
return embedding_model
# Load the embedding model
embedding_model = create_embedding_model()
embedding_model.load_weights('facenet_siamese_embedding.h5')
# Database to store embeddings and user IDs
user_embeddings = []
user_ids = []
# Threshold
RECOGNITION_THRESHOLD = 0.1 # Adjust as needed
# Preprocess the image
def preprocess_image(image):
image = cv2.resize(image, (160, 160)) # Resize image to match FaceNet input size
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert to RGB (OpenCV loads images in BGR)
image = tf.convert_to_tensor(image) # Convert to TensorFlow tensor
image = tf.image.convert_image_dtype(image, tf.float32) # Normalize pixel values
return np.expand_dims(image, axis=0) # Add batch dimension
# Generate embedding
def generate_embedding(image):
preprocessed_image = preprocess_image(image)
return embedding_model.predict(preprocessed_image)[0]
# Register new user
def register_user(image, user_id):
try:
embedding = generate_embedding(image)
user_embeddings.append(embedding)
user_ids.append(user_id)
return f"User {user_id} registered successfully."
except Exception as e:
return f"Error during registration: {str(e)}"
# Recognize user
def recognize_user(image):
try:
new_embedding = generate_embedding(image)
closest_user_id = None
closest_distance = float('inf')
for user_id, embedding in zip(user_ids, user_embeddings):
distance = cosine(new_embedding, embedding)
print(f"Distance for {user_id}: {distance}") # Debug: Print distances for each user
if distance < closest_distance:
closest_distance = distance
closest_user_id = user_id
print(f"Min distance: {closest_distance}") # Debug: Print minimum distance
if closest_distance <= RECOGNITION_THRESHOLD:
return f"Recognized User: {closest_user_id}"
else:
return f"User not recognized. Closest Distance: {closest_distance}"
except Exception as e:
return f"Error during recognition: {str(e)}"
def main():
with gr.Blocks() as demo:
gr.Markdown("Facial Recognition System")
with gr.Tab("Register"):
with gr.Row():
img_register = gr.Image()
user_id = gr.Textbox(label="User ID")
register_button = gr.Button("Register")
register_output = gr.Textbox()
register_button.click(register_user, inputs=[img_register, user_id], outputs=register_output)
with gr.Tab("Recognize"):
with gr.Row():
img_recognize = gr.Image()
recognize_button = gr.Button("Recognize")
recognize_output = gr.Textbox()
recognize_button.click(recognize_user, inputs=[img_recognize], outputs=recognize_output)
demo.launch(share=True)
if __name__ == "__main__":
main()