Spaces:
Build error
Build error
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.") |