Segizu's picture
sd
d771ba3
import numpy as np
from PIL import Image
import gradio as gr
from deepface import DeepFace
from datasets import load_dataset
# Cargar el dataset de Hugging Face
dataset = load_dataset("Segizu/dataset_faces")
if "train" in dataset:
dataset = dataset["train"]
# Cargar embeddings de todas las imágenes del dataset
def build_database():
database = []
for i, item in enumerate(dataset):
try:
img = item["image"]
# Convertir a RGB y np.array
img_rgb = img.convert("RGB")
img_np = np.array(img_rgb)
# Obtener representación (embedding)
representation = DeepFace.represent(img_path=img_np, model_name="Facenet", enforce_detection=False)[0]["embedding"]
database.append((f"image_{i}", img_rgb, representation))
except Exception as e:
print(f"❌ No se pudo procesar imagen {i}: {e}")
return database
# Inicializamos base de datos de embeddings
database = build_database()
# Comparar imagen cargada con la base
def find_similar_faces(uploaded_image):
try:
img_np = np.array(uploaded_image.convert("RGB"))
query_representation = DeepFace.represent(img_path=img_np, model_name="Facenet", enforce_detection=False)[0]["embedding"]
except:
return [], "⚠ No se detectó un rostro válido en la imagen."
similarities = []
for name, db_img, rep in database:
distance = np.linalg.norm(np.array(query_representation) - np.array(rep))
similarity = 1 / (1 + distance) # Normalizado
similarities.append((similarity, name, db_img))
similarities.sort(reverse=True)
top_matches = similarities[:5]
gallery_items = []
text_summary = ""
for sim, name, img in top_matches:
caption = f"{name} - Similitud: {sim:.2f}"
gallery_items.append((img, caption))
text_summary += caption + "\n"
return gallery_items, text_summary
# Interfaz Gradio
demo = gr.Interface(
fn=find_similar_faces,
inputs=gr.Image(label="📤 Sube una imagen", type="pil"),
outputs=[
gr.Gallery(label="📸 Rostros más similares"),
gr.Textbox(label="🧠 Similitud", lines=6)
],
title="🔍 Buscador de Rostros con DeepFace",
description="Sube una imagen y se comparará contra los rostros del dataset alojado en Hugging Face (`Segizu/dataset_faces`)."
)
demo.launch()