no cache embeddings baches

app.py

@@ -2,12 +2,11 @@ import numpy as np
 from PIL import Image
 import gradio as gr
 from deepface import DeepFace
-from datasets import load_dataset, DownloadConfig
+from datasets import load_dataset, DownloadConfig, Image as HfImage
 import os
 import pickle
 from pathlib import Path
 import gc
-import io
 
 # 🔑 Configurar token de Hugging Face
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -19,33 +18,23 @@ EMBEDDINGS_DIR = Path("embeddings")
 EMBEDDINGS_DIR.mkdir(exist_ok=True)
 EMBEDDINGS_FILE = EMBEDDINGS_DIR / "embeddings.pkl"
 
-os.system("rm -rf ~/.cache/huggingface/hub/datasets--Segizu--facial-recognition")
-
-# ✅ Cargar el dataset de Hugging Face forzando la descarga limpia
-download_config = DownloadConfig(
-    force_download=True,
-    token=HF_TOKEN
-)
+# ✅ Cargar el dataset de Hugging Face con imágenes decodificadas
+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"]
 
+# Asegurar que la columna 'image' sea del tipo imagen
+dataset = dataset.cast_column("image", HfImage())
+
 # 🔄 Preprocesar imagen para Facenet
-def preprocess_image(img):
-    if isinstance(img, str):
-        # Si es una ruta de archivo o bytes en string
-        img = Image.open(io.BytesIO(img.encode() if isinstance(img, str) else img))
-    elif isinstance(img, bytes):
-        # Si son bytes directos
-        img = Image.open(io.BytesIO(img))
-    
+def preprocess_image(img: Image.Image) -> np.ndarray:
     img_rgb = img.convert("RGB")
     img_resized = img_rgb.resize((160, 160), Image.Resampling.LANCZOS)
     return np.array(img_resized)
 
 # 📦 Construir base de datos de embeddings
 def build_database():
-    # Intentar cargar embeddings desde el archivo
     if EMBEDDINGS_FILE.exists():
         print("📂 Cargando embeddings desde el archivo...")
         with open(EMBEDDINGS_FILE, 'rb') as f:
@@ -53,61 +42,45 @@ def build_database():
     
     print("🔄 Calculando embeddings (esto puede tomar unos minutos)...")
     database = []
-    batch_size = 10  # Procesar 10 imágenes a la vez
+    batch_size = 10
     
     for i in range(0, len(dataset), batch_size):
         batch = dataset[i:i + batch_size]
-        print(f"📦 Procesando lote {i//batch_size + 1}/{(len(dataset) + batch_size - 1)//batch_size}")
+        print(f"📦 Procesando lote {i // batch_size + 1}/{(len(dataset) + batch_size - 1) // batch_size}")
         
         for j, item in enumerate(batch):
             try:
-                # Debug: Imprimir la estructura del item
-                print(f"Estructura del item {i+j}:", type(item), item.keys() if hasattr(item, 'keys') else "No tiene keys")
-                
-                # Intentar diferentes formas de acceder a la imagen
-                if isinstance(item, dict):
-                    if 'image' in item:
-                        img = item['image']
-                    elif 'bytes' in item:
-                        img = item['bytes']
-                    else:
-                        print(f"❌ No se encontró la imagen en el item {i+j}")
-                        continue
-                else:
-                    print(f"❌ Formato de item no reconocido: {type(item)}")
-                    continue
-
+                img = item["image"]  # Ya es un objeto PIL.Image
                 img_processed = preprocess_image(img)
                 embedding = DeepFace.represent(
                     img_path=img_processed,
                     model_name="Facenet",
                     enforce_detection=False
                 )[0]["embedding"]
+                
                 database.append((f"image_{i+j}", img, embedding))
                 print(f"✅ Procesada imagen {i+j+1}/{len(dataset)}")
                 
                 # Liberar memoria
                 del img_processed
                 gc.collect()
-                
+            
             except Exception as e:
                 print(f"❌ No se pudo procesar imagen {i+j}: {str(e)}")
-                print(f"Tipo de error: {type(e)}")
                 continue
         
-        # Guardar progreso después de cada lote
-        if database:  # Solo guardar si hay datos
+        # Guardar progreso
+        if database:
             print("💾 Guardando progreso...")
             with open(EMBEDDINGS_FILE, 'wb') as f:
                 pickle.dump(database, f)
         
-        # Liberar memoria después de cada lote
         gc.collect()
     
     return database
 
 # 🔍 Buscar rostros similares
-def find_similar_faces(uploaded_image):
+def find_similar_faces(uploaded_image: Image.Image):
     try:
         img_processed = preprocess_image(uploaded_image)
         query_embedding = DeepFace.represent(
@@ -115,11 +88,8 @@ def find_similar_faces(uploaded_image):
             model_name="Facenet",
             enforce_detection=False
         )[0]["embedding"]
-        
-        # Liberar memoria
         del img_processed
         gc.collect()
-        
     except Exception as e:
         print(f"Error al procesar imagen de consulta: {str(e)}")
         return [], "⚠ No se detectó un rostro válido en la imagen."