Spaces:
Build error
Build error
File size: 5,094 Bytes
9fd78b1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
"""
Utilidades para manejar la persistencia de la base de datos de rostros.
"""
import os
import pickle
import streamlit as st
import json
import base64
import numpy as np
# Configurar ruta para la base de datos
DATABASE_FILE = "face_database.pkl"
def save_face_database(database):
"""
Guarda la base de datos de rostros en un archivo persistente.
Args:
database (dict): La base de datos de rostros a guardar
"""
try:
# Convertir numpy arrays a listas para poder serializarlas
serializable_db = {}
for name, info in database.items():
serializable_db[name] = {}
# Manejar diferentes formatos de la base de datos
if 'embeddings' in info:
serializable_db[name]['embeddings'] = [emb.tolist() if isinstance(emb, np.ndarray) else emb for emb in info['embeddings']]
serializable_db[name]['models'] = info['models']
serializable_db[name]['count'] = info['count']
elif 'embedding' in info:
# Formato antiguo
serializable_db[name]['embedding'] = info['embedding'].tolist() if isinstance(info['embedding'], np.ndarray) else info['embedding']
serializable_db[name]['count'] = info.get('count', 1)
# Guardar en un archivo pickle
with open(DATABASE_FILE, 'wb') as f:
pickle.dump(serializable_db, f)
return True
except Exception as e:
st.error(f"Error al guardar la base de datos: {str(e)}")
return False
def load_face_database():
"""
Carga la base de datos de rostros desde un archivo persistente.
Returns:
dict: La base de datos de rostros cargada, o un diccionario vacío si no existe el archivo.
"""
if not os.path.exists(DATABASE_FILE):
return {}
try:
with open(DATABASE_FILE, 'rb') as f:
database = pickle.load(f)
# Convertir listas a numpy arrays
for name, info in database.items():
if 'embeddings' in info:
database[name]['embeddings'] = [np.array(emb) if isinstance(emb, list) else emb for emb in info['embeddings']]
elif 'embedding' in info:
database[name]['embedding'] = np.array(info['embedding']) if isinstance(info['embedding'], list) else info['embedding']
return database
except Exception as e:
st.error(f"Error al cargar la base de datos: {str(e)}")
return {}
def export_database_json():
"""
Exporta la base de datos a un archivo JSON para compartir o hacer backup.
Returns:
str: Ruta al archivo JSON exportado.
"""
try:
if 'face_database' in st.session_state and st.session_state.face_database:
# Crear una versión serializable de la base de datos
serializable_db = {}
for name, info in st.session_state.face_database.items():
serializable_db[name] = {}
if 'embeddings' in info:
serializable_db[name]['embeddings'] = [
base64.b64encode(np.array(emb).tobytes()).decode('utf-8')
for emb in info['embeddings']
]
serializable_db[name]['models'] = info['models']
serializable_db[name]['count'] = info['count']
elif 'embedding' in info:
serializable_db[name]['embedding'] = base64.b64encode(
np.array(info['embedding']).tobytes()
).decode('utf-8')
serializable_db[name]['count'] = info.get('count', 1)
# Guardar en un archivo JSON
export_file = "face_database_export.json"
with open(export_file, 'w') as f:
json.dump(serializable_db, f, indent=2)
return export_file
return None
except Exception as e:
st.error(f"Error al exportar la base de datos: {str(e)}")
return None
def import_database_json(json_file):
"""
Importa una base de datos desde un archivo JSON.
Args:
json_file: El archivo JSON a importar
Returns:
dict: La base de datos importada.
"""
try:
content = json_file.read()
imported_db = json.loads(content)
# Convertir datos codificados en base64 a numpy arrays
for name, info in imported_db.items():
if 'embeddings' in info:
imported_db[name]['embeddings'] = [
np.frombuffer(base64.b64decode(emb), dtype=np.float32)
for emb in info['embeddings']
]
elif 'embedding' in info:
imported_db[name]['embedding'] = np.frombuffer(
base64.b64decode(info['embedding']), dtype=np.float32
)
return imported_db
except Exception as e:
st.error(f"Error al importar la base de datos: {str(e)}")
return {} |