jarondon82/ComputerVisionProject

1"""2Utilidades para manejar la persistencia de la base de datos de rostros.3"""4import os5import pickle6import streamlit as st7import json8import base649import numpy as np10 11# Configurar ruta para la base de datos12DATABASE_FILE = "face_database.pkl"13 14def save_face_database(database):15    """16    Guarda la base de datos de rostros en un archivo persistente.17    18    Args:19        database (dict): La base de datos de rostros a guardar20    """21    try:22        # Verificar si hay datos para guardar23        if not database:24            # Si la base de datos está vacía, eliminar el archivo si existe25            if os.path.exists(DATABASE_FILE):26                os.remove(DATABASE_FILE)27                st.sidebar.write("Database was empty - removed existing file")28            return True29        30        # Convertir numpy arrays a listas para poder serializarlas31        serializable_db = {}32        for name, info in database.items():33            serializable_db[name] = {}34            # Manejar diferentes formatos de la base de datos35            if 'embeddings' in info:36                serializable_db[name]['embeddings'] = [emb.tolist() if isinstance(emb, np.ndarray) else emb for emb in info['embeddings']]37                serializable_db[name]['models'] = info['models']38                serializable_db[name]['count'] = info['count']39                40                # Guardar imagen facial si existe41                if 'face_image' in info:42                    serializable_db[name]['face_image'] = info['face_image'].tolist() if isinstance(info['face_image'], np.ndarray) else info['face_image']43            elif 'embedding' in info:44                # Formato antiguo45                serializable_db[name]['embedding'] = info['embedding'].tolist() if isinstance(info['embedding'], np.ndarray) else info['embedding']46                serializable_db[name]['count'] = info.get('count', 1)47                48                # Guardar imagen facial si existe49                if 'face_image' in info:50                    serializable_db[name]['face_image'] = info['face_image'].tolist() if isinstance(info['face_image'], np.ndarray) else info['face_image']51        52        # Guardar en un archivo pickle53        with open(DATABASE_FILE, 'wb') as f:54            pickle.dump(serializable_db, f)55            56        # Verificar que el archivo se creó correctamente57        if os.path.exists(DATABASE_FILE):58            st.sidebar.write(f"Database saved to {DATABASE_FILE} ({len(serializable_db)} entries)")59        return True60    except Exception as e:61        st.error(f"Error al guardar la base de datos: {str(e)}")62        return False63 64def load_face_database():65    """66    Carga la base de datos de rostros desde un archivo persistente.67    68    Returns:69        dict: La base de datos de rostros cargada, o un diccionario vacío si no existe el archivo.70    """71    if not os.path.exists(DATABASE_FILE):72        return {}73    74    try:75        with open(DATABASE_FILE, 'rb') as f:76            database = pickle.load(f)77        78        # Convertir listas a numpy arrays79        for name, info in database.items():80            if 'embeddings' in info:81                database[name]['embeddings'] = [np.array(emb) if isinstance(emb, list) else emb for emb in info['embeddings']]82                # Cargar imagen facial si existe83                if 'face_image' in info:84                    database[name]['face_image'] = np.array(info['face_image']) if isinstance(info['face_image'], list) else info['face_image']85            elif 'embedding' in info:86                database[name]['embedding'] = np.array(info['embedding']) if isinstance(info['embedding'], list) else info['embedding']87                # Cargar imagen facial si existe88                if 'face_image' in info:89                    database[name]['face_image'] = np.array(info['face_image']) if isinstance(info['face_image'], list) else info['face_image']90        91        return database92    except Exception as e:93        st.error(f"Error al cargar la base de datos: {str(e)}")94        return {}95 96def export_database_json():97    """98    Exporta la base de datos a un archivo JSON para compartir o hacer backup.99    100    Returns:101        str: Ruta al archivo JSON exportado.102    """103    try:104        if 'face_database' in st.session_state and st.session_state.face_database:105            # Crear una versión serializable de la base de datos106            serializable_db = {}107            for name, info in st.session_state.face_database.items():108                serializable_db[name] = {}109                if 'embeddings' in info:110                    serializable_db[name]['embeddings'] = [111                        base64.b64encode(np.array(emb).tobytes()).decode('utf-8') 112                        for emb in info['embeddings']113                    ]114                    serializable_db[name]['models'] = info['models']115                    serializable_db[name]['count'] = info['count']116                    117                    # Incluir imagen facial si existe118                    if 'face_image' in info:119                        serializable_db[name]['face_image'] = base64.b64encode(120                            np.array(info['face_image']).tobytes()121                        ).decode('utf-8')122                        serializable_db[name]['face_image_shape'] = info['face_image'].shape123                elif 'embedding' in info:124                    serializable_db[name]['embedding'] = base64.b64encode(125                        np.array(info['embedding']).tobytes()126                    ).decode('utf-8')127                    serializable_db[name]['count'] = info.get('count', 1)128                    129                    # Incluir imagen facial si existe130                    if 'face_image' in info:131                        serializable_db[name]['face_image'] = base64.b64encode(132                            np.array(info['face_image']).tobytes()133                        ).decode('utf-8')134                        serializable_db[name]['face_image_shape'] = info['face_image'].shape135            136            # Guardar en un archivo JSON137            export_file = "face_database_export.json"138            with open(export_file, 'w') as f:139                json.dump(serializable_db, f, indent=2)140            141            return export_file142        return None143    except Exception as e:144        st.error(f"Error al exportar la base de datos: {str(e)}")145        return None146 147def import_database_json(json_file):148    """149    Importa una base de datos desde un archivo JSON.150    151    Args:152        json_file: El archivo JSON a importar153        154    Returns:155        dict: La base de datos importada.156    """157    try:158        content = json_file.read()159        imported_db = json.loads(content)160        161        # Convertir datos codificados en base64 a numpy arrays162        for name, info in imported_db.items():163            if 'embeddings' in info:164                imported_db[name]['embeddings'] = [165                    np.frombuffer(base64.b64decode(emb), dtype=np.float32) 166                    for emb in info['embeddings']167                ]168                169                # Importar imagen facial si existe170                if 'face_image' in info and 'face_image_shape' in info:171                    face_data = np.frombuffer(base64.b64decode(info['face_image']), dtype=np.uint8)172                    shape = info['face_image_shape']173                    imported_db[name]['face_image'] = face_data.reshape(shape)174            elif 'embedding' in info:175                imported_db[name]['embedding'] = np.frombuffer(176                    base64.b64decode(info['embedding']), dtype=np.float32177                )178                179                # Importar imagen facial si existe180                if 'face_image' in info and 'face_image_shape' in info:181                    face_data = np.frombuffer(base64.b64decode(info['face_image']), dtype=np.uint8)182                    shape = info['face_image_shape']183                    imported_db[name]['face_image'] = face_data.reshape(shape)184        185        return imported_db186    except Exception as e:187        st.error(f"Error al importar la base de datos: {str(e)}")188        return {}189 190def print_database_info():191    """192    Imprime información sobre la base de datos actual para depuración.193    """194    if 'face_database' in st.session_state:195        db = st.session_state.face_database196        st.sidebar.write("--- Database Debug Info ---")197        st.sidebar.write(f"Database contains {len(db)} entries")198        199        # Mostrar nombres en la base de datos200        if db:201            names = list(db.keys())202            st.sidebar.write(f"Names in database: {', '.join(names)}")203            204            # Mostrar detalles del primer elemento205            if names:206                first_entry = db[names[0]]207                st.sidebar.write(f"Sample entry for '{names[0]}':")208                if 'embeddings' in first_entry:209                    st.sidebar.write(f"- Has {len(first_entry['embeddings'])} embeddings")210                    st.sidebar.write(f"- Models: {', '.join(first_entry['models'])}")211                    st.sidebar.write(f"- Count: {first_entry['count']}")212                    213                    # Mostrar si tiene imagen214                    if 'face_image' in first_entry:215                        st.sidebar.write(f"- Has reference face image: {first_entry['face_image'].shape}")216                    else:217                        st.sidebar.write("- No reference image")218                elif 'embedding' in first_entry:219                    st.sidebar.write("- Has single embedding (old format)")220                    221                    # Mostrar si tiene imagen222                    if 'face_image' in first_entry:223                        st.sidebar.write(f"- Has reference face image: {first_entry['face_image'].shape}")224                    else:225                        st.sidebar.write("- No reference image")226        else:227            st.sidebar.write("Database is empty")