finalproject / utils /download_models.py
jarondon82's picture
Initial commit for EmotionMirror finalproject
f7e620e
"""
Script para descargar modelos adicionales necesarios para EmotionMirror
"""
import os
import sys
import logging
import requests
import bz2
import shutil
from pathlib import Path
# Configurar logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# URL del modelo de landmarks faciales
LANDMARKS_MODEL_URL = "http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2"
def get_project_root():
"""Obtiene la ruta raíz del proyecto"""
path = Path(__file__).parent.parent
return path
def create_models_dir():
"""Crea el directorio de modelos si no existe"""
root = get_project_root()
models_dir = root / "models"
models_dir.mkdir(exist_ok=True)
return models_dir
def download_file(url, destination):
"""Descarga un archivo desde una URL"""
try:
logger.info(f"Descargando {url} en {destination}")
response = requests.get(url, stream=True)
response.raise_for_status() # Verificar si hay errores
with open(destination, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
logger.info(f"Descarga completada: {destination}")
return True
except Exception as e:
logger.error(f"Error descargando archivo: {e}")
return False
def extract_bz2(source, destination):
"""Extrae un archivo .bz2 a un destino"""
try:
logger.info(f"Descomprimiendo {source} en {destination}")
with bz2.BZ2File(source, 'rb') as source_file:
with open(destination, 'wb') as dest_file:
shutil.copyfileobj(source_file, dest_file)
logger.info(f"Descompresión completada: {destination}")
return True
except Exception as e:
logger.error(f"Error descomprimiendo archivo: {e}")
return False
def download_landmarks_model():
"""Descarga y configura el modelo de landmarks faciales"""
models_dir = create_models_dir()
model_compressed = models_dir / "shape_predictor_68_face_landmarks.dat.bz2"
model_path = models_dir / "shape_predictor_68_face_landmarks.dat"
# Verificar si el modelo ya existe
if model_path.exists():
logger.info(f"El modelo ya existe en {model_path}")
return True
# Descargar el modelo comprimido
if not download_file(LANDMARKS_MODEL_URL, model_compressed):
return False
# Extraer el modelo
if not extract_bz2(model_compressed, model_path):
return False
# Eliminar el archivo comprimido después de extraerlo
try:
os.remove(model_compressed)
logger.info(f"Archivo comprimido eliminado: {model_compressed}")
except Exception as e:
logger.warning(f"No se pudo eliminar el archivo comprimido: {e}")
return True
if __name__ == "__main__":
logger.info("Iniciando descarga de modelos adicionales...")
if download_landmarks_model():
logger.info("Modelo de landmarks faciales descargado correctamente.")
else:
logger.error("Error descargando el modelo de landmarks faciales.")
sys.exit(1)
logger.info("Todos los modelos han sido descargados correctamente.")