sd

app.py

@@ -1,52 +1,53 @@
-import os
 import numpy as np
 from PIL import Image
 import gradio as gr
 from deepface import DeepFace
 from datasets import load_dataset
 
-# Cargamos el dataset sin aplicar formato especial
+# 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 filename in os.listdir(IMAGE_DIRECTORY):
-        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
-            path = os.path.join(IMAGE_DIRECTORY, filename)
-            try:
-                representation = DeepFace.represent(img_path=path, model_name="Facenet")[0]["embedding"]
-                database.append((filename, path, representation))
-            except:
-                print(f"❌ No se pudo procesar: {filename}")
+    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
+# Inicializamos base de datos de embeddings
 database = build_database()
 
-# Comparar imagen cargada con las del dataset
+# Comparar imagen cargada con la base
 def find_similar_faces(uploaded_image):
     try:
-        uploaded_image = np.array(uploaded_image)
-        query_representation = DeepFace.represent(img_path=uploaded_image, model_name="Facenet")[0]["embedding"]
+        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, path, rep in database:
+    for name, db_img, rep in database:
         distance = np.linalg.norm(np.array(query_representation) - np.array(rep))
-        similarity = 1 / (1 + distance)  # Normalizamos para que 1 = muy similar
-        similarities.append((similarity, name, path))
+        similarity = 1 / (1 + distance)  # Normalizado
+        similarities.append((similarity, name, db_img))
 
-    # Ordenar por similitud
     similarities.sort(reverse=True)
     top_matches = similarities[:5]
 
-    # Formatear salida para gradio
     gallery_items = []
     text_summary = ""
-    for sim, name, path in top_matches:
-        img = Image.open(path)
+    for sim, name, img in top_matches:
         caption = f"{name} - Similitud: {sim:.2f}"
         gallery_items.append({"image": img, "caption": caption})
         text_summary += caption + "\n"
@@ -62,9 +63,7 @@ demo = gr.Interface(
         gr.Textbox(label="🧠 Similitud", lines=6)
     ],
     title="🔍 Buscador de Rostros con DeepFace",
-    description="Sube una imagen y te mostrará los rostros más similares desde el directorio `dataset_faces/`."
+    description="Sube una imagen y se comparará contra los rostros del dataset alojado en Hugging Face (`Segizu/dataset_faces`)."
 )
 
 demo.launch()
-
-