Segizu's picture
add cuda1
b24755b
import streamlit as st
import face_recognition
import cv2
import numpy as np
import os
from PIL import Image
import dlib
# Check if CUDA is available
if dlib.DLIB_USE_CUDA:
st.write("CUDA is available. Using GPU for face recognition.")
else:
st.write("CUDA is not available. Using CPU for face recognition.")
# 📂 Ruta donde están las imágenes
IMAGE_DIRECTORY = "dataset_faces/"
# 📌 Función para cargar imágenes y extraer embeddings
def load_images_and_encodings(directory):
known_encodings = []
known_images = []
known_names = []
for filename in os.listdir(directory):
if filename.endswith((".jpg", ".png", ".jpeg")):
path = os.path.join(directory, filename)
image = face_recognition.load_image_file(path)
encodings = face_recognition.face_encodings(image)
if encodings: # Si encontró una cara
known_encodings.append(encodings[0])
known_images.append(path)
known_names.append(filename)
return known_encodings, known_images, known_names
# 📌 Cargar imágenes y sus embeddings al inicio
known_encodings, known_images, known_names = load_images_and_encodings(IMAGE_DIRECTORY)
# 📌 Función para encontrar imágenes similares
def find_similar_faces(uploaded_image):
# Convertir a array de NumPy
image_np = np.array(uploaded_image)
face_encodings = face_recognition.face_encodings(image_np)
if not face_encodings:
return [], [] # Si no encuentra caras en la imagen subida
query_encoding = face_encodings[0]
distances = face_recognition.face_distance(known_encodings, query_encoding)
# Ordenar las imágenes por similitud
sorted_indices = np.argsort(distances)
return [known_images[i] for i in sorted_indices], [distances[i] for i in sorted_indices]
# 📌 Interfaz en Streamlit
st.title("🔍 Buscador de Rostros en un Directorio")
st.write("Sube una imagen y te mostraremos las fotos más similares en el directorio.")
uploaded_file = st.file_uploader("📤 Sube una imagen", type=["jpg", "png", "jpeg"])
if uploaded_file:
# Mostrar imagen subida
uploaded_image = Image.open(uploaded_file)
st.image(uploaded_image, caption="Imagen subida", use_container_width=True)
# Buscar rostros similares
similar_images, distances = find_similar_faces(uploaded_image)
# Mostrar resultados
if similar_images:
st.subheader("📸 Imágenes más similares:")
cols = st.columns(len(similar_images[:])) # Crear columnas para las imágenes
for idx, (col, image_path) in enumerate(zip(cols, similar_images[:])): # Mostrar las 5 más similares
col.image(image_path, caption=f"Similitud: {1 - distances[idx]:.2f}", use_container_width=True)
else:
st.warning("⚠ No se detectó ningún rostro en la imagen subida.")