cache embeddings

app.py

@@ -4,11 +4,27 @@ import gradio as gr
 from deepface import DeepFace
 from datasets import load_dataset, DownloadConfig
 import os
+import pickle
+from pathlib import Path
+
+# 🔑 Configurar token de Hugging Face
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("⚠️ Por favor, configura la variable de entorno HF_TOKEN para acceder al dataset privado")
+
+# 📁 Configurar directorio de caché
+CACHE_DIR = Path("cache")
+CACHE_DIR.mkdir(exist_ok=True)
+EMBEDDINGS_CACHE = CACHE_DIR / "embeddings.pkl"
+
 os.system("rm -rf ~/.cache/huggingface/hub/datasets--Segizu--dataset_faces")
 
 # ✅ Cargar el dataset de Hugging Face forzando la descarga limpia
-download_config = DownloadConfig(force_download=True)
-dataset = load_dataset("Segizu/dataset_faces", download_config=download_config)
+download_config = DownloadConfig(
+    force_download=True,
+    token=HF_TOKEN
+)
+dataset = load_dataset("Segizu/facial-recognition", download_config=download_config)
 if "train" in dataset:
     dataset = dataset["train"]
 
@@ -20,6 +36,13 @@ def preprocess_image(img):
 
 # 📦 Construir base de datos de embeddings
 def build_database():
+    # Intentar cargar embeddings desde caché
+    if EMBEDDINGS_CACHE.exists():
+        print("📂 Cargando embeddings desde caché...")
+        with open(EMBEDDINGS_CACHE, 'rb') as f:
+            return pickle.load(f)
+    
+    print("🔄 Calculando embeddings (esto puede tomar unos minutos)...")
     database = []
     for i, item in enumerate(dataset):
         try:
@@ -31,8 +54,15 @@ def build_database():
                 enforce_detection=False
             )[0]["embedding"]
             database.append((f"image_{i}", img, embedding))
+            print(f"✅ Procesada imagen {i+1}/{len(dataset)}")
         except Exception as e:
             print(f"❌ No se pudo procesar imagen {i}: {e}")
+    
+    # Guardar embeddings en caché
+    print("💾 Guardando embeddings en caché...")
+    with open(EMBEDDINGS_CACHE, 'wb') as f:
+        pickle.dump(database, f)
+    
     return database
 
 # 🔍 Buscar rostros similares
@@ -54,7 +84,7 @@ def find_similar_faces(uploaded_image):
         similarities.append((sim_score, name, db_img))
 
     similarities.sort(reverse=True)
-    top_matches = similarities[:]
+    top_matches = similarities[:5]
 
     gallery_items = []
     text_summary = ""
@@ -66,7 +96,9 @@ def find_similar_faces(uploaded_image):
     return gallery_items, text_summary
 
 # ⚙️ Inicializar base
+print("🚀 Iniciando aplicación...")
 database = build_database()
+print(f"✅ Base de datos cargada con {len(database)} imágenes")
 
 # 🎛️ Interfaz Gradio
 demo = gr.Interface(
@@ -77,7 +109,7 @@ demo = gr.Interface(
         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`)."
+    description="Sube una imagen y se comparará contra los rostros del dataset alojado en Hugging Face (`Segizu/facial-recognition`)."
 )
 
 demo.launch()