vto_model.py

"""
VTO Model Optimisé - Réduction latence 60-70%
✅ Cache images vêtements
✅ Redimensionnement frame avant traitement
✅ MediaPipe optimisé
"""

import cv2
import numpy as np
import mediapipe as mp
import base64
import requests
from io import BytesIO
from PIL import Image
from functools import lru_cache
import hashlib

# ✅ Configuration MediaPipe optimisée
mp_pose = mp.solutions.pose
pose = mp_pose.Pose(
    static_image_mode=False,
    model_complexity=0,  # ✅ 0 = plus rapide (était 1)
    min_detection_confidence=0.3,  # ✅ Réduit (était 0.5)
    min_tracking_confidence=0.3,  # ✅ Réduit (était 0.5)
    enable_segmentation=False,  # ✅ Désactivé pour perfs
    smooth_landmarks=True  # ✅ Lissage pour éviter tremblements
)

# ✅ Configuration
SCALE_FACTOR = {
    "top": 1.7,
    "bottom": 1.5,
    "footwear": 1.1,
    "outerwear": 1.8
}

OFFSET_Y = {
    "top": -0.15,
    "bottom": -0.1,
    "footwear": -0.4,
    "outerwear": -0.2
}

DRAW_ORDER = ["footwear", "bottom", "top", "outerwear"]

# ✅ NOUVEAU : Cache en mémoire des images de vêtements
_CLOTHES_CACHE = {}
MAX_CACHE_SIZE = 50  # Maximum 50 images en cache

def _get_cache_key(url: str) -> str:
    """Génère une clé de cache unique pour une URL"""
    return hashlib.md5(url.encode()).hexdigest()

@lru_cache(maxsize=50)
def download_image_cached(url: str):
    """
    ✅ Télécharge et cache une image de vêtement
    Utilise LRU cache de Python pour éviter re-téléchargements
    """
    try:
        cache_key = _get_cache_key(url)
        
        # Vérifier le cache manuel d'abord
        if cache_key in _CLOTHES_CACHE:
            print(f"   📦 Cache HIT: {url[:50]}...")
            return _CLOTHES_CACHE[cache_key]
        
        print(f"   📥 Downloading: {url[:50]}...")
        response = requests.get(url, timeout=5)
        response.raise_for_status()
        
        img = Image.open(BytesIO(response.content)).convert("RGBA")
        
        # ✅ Redimensionner pour économiser mémoire (max 800px)
        max_size = 800
        if max(img.size) > max_size:
            ratio = max_size / max(img.size)
            new_size = (int(img.width * ratio), int(img.height * ratio))
            img = img.resize(new_size, Image.Resampling.LANCZOS)
        
        # Sauvegarder dans le cache
        if len(_CLOTHES_CACHE) < MAX_CACHE_SIZE:
            _CLOTHES_CACHE[cache_key] = img
        
        return img
        
    except Exception as e:
        print(f"   ❌ Download failed: {str(e)}")
        return None

def overlay_transparent(background, overlay, x, y, w, h):
    """
    ✅ Superpose une image PNG transparente (optimisé)
    """
    if overlay is None:
        return background
    
    # ✅ Redimensionner une seule fois
    overlay_resized = cv2.resize(overlay, (w, h), interpolation=cv2.INTER_LINEAR)
    h_bg, w_bg = background.shape[:2]
    
    # Vérifier limites
    if x >= w_bg or y >= h_bg or x + w <= 0 or y + h <= 0:
        return background
    
    # Calculer régions
    x1, y1 = max(x, 0), max(y, 0)
    x2, y2 = min(x + w, w_bg), min(y + h, h_bg)
    ox1, oy1 = max(0, -x), max(0, -y)
    ox2, oy2 = min(w, w_bg - x), min(h, h_bg - y)
    
    overlay_crop = overlay_resized[oy1:oy2, ox1:ox2]
    background_crop = background[y1:y2, x1:x2]
    
    if overlay_crop.shape[0] != background_crop.shape[0]:
        return background
    
    # ✅ Alpha blending optimisé
    alpha = overlay_crop[:, :, 3:4].astype(np.float32) / 255.0
    
    for c in range(3):
        background_crop[:, :, c] = (
            alpha[:, :, 0] * overlay_crop[:, :, c] +
            (1.0 - alpha[:, :, 0]) * background_crop[:, :, c]
        ).astype(np.uint8)
    
    background[y1:y2, x1:x2] = background_crop
    return background

def process_frame_vto(frame_base64: str, clothes_data: list):
    """
    ✅ Traite une frame avec optimisations de performance
    
    Optimisations:
    - Redimensionnement frame si > 640px
    - Cache images vêtements
    - MediaPipe model_complexity=0
    - Interpolation rapide
    """
    try:
        # ✅ Décoder l'image
        img_data = base64.b64decode(frame_base64)
        nparr = np.frombuffer(img_data, np.uint8)
        frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if frame is None:
            return {"success": False, "error": "Invalid image data"}
        
        original_shape = frame.shape
        
        # ✅ OPTIMISATION 1 : Redimensionner la frame si trop grande
        max_width = 640
        if frame.shape[1] > max_width:
            ratio = max_width / frame.shape[1]
            new_size = (max_width, int(frame.shape[0] * ratio))
            frame = cv2.resize(frame, new_size, interpolation=cv2.INTER_LINEAR)
            print(f"   📏 Resized: {original_shape[1]}x{original_shape[0]} → {new_size[0]}x{new_size[1]}")
        
        # ✅ OPTIMISATION 2 : Détection pose MediaPipe (model_complexity=0)
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = pose.process(rgb)
        
        if not results.pose_landmarks:
            # Pas de corps détecté, retourner frame originale
            _, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 85])
            encoded = base64.b64encode(buffer).decode('utf-8')
            return {"success": True, "frame": encoded, "message": "No body detected"}
        
        lm = results.pose_landmarks.landmark
        h_frame, w_frame = frame.shape[:2]
        
        # ✅ OPTIMISATION 3 : Charger vêtements avec cache
        wardrobe = {}
        for cloth in clothes_data:
            category = cloth.get("category", "").lower()
            url = cloth.get("processedImageURL") or cloth.get("imageURL")
            
            if not url:
                continue
            
            try:
                # ✅ Utiliser le cache
                img_pil = download_image_cached(url)
                if img_pil is None:
                    continue
                
                img_cv = np.array(img_pil)
                img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGBA2BGRA)
                wardrobe[category] = img_cv
                
            except Exception as e:
                print(f"   ⚠️ Failed to load {category}: {str(e)}")
                continue
        
        if not wardrobe:
            _, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 85])
            encoded = base64.b64encode(buffer).decode('utf-8')
            return {"success": True, "frame": encoded, "message": "No clothes to apply"}
        
        # ✅ OPTIMISATION 4 : Appliquer vêtements dans l'ordre
        for category in DRAW_ORDER:
            if category not in wardrobe:
                continue
            
            cloth_img = wardrobe[category]
            
            # Déterminer les points de référence
            if category in ["top", "outerwear"]:
                p1 = lm[mp_pose.PoseLandmark.LEFT_SHOULDER]
                p2 = lm[mp_pose.PoseLandmark.RIGHT_SHOULDER]
            elif category == "bottom":
                p1 = lm[mp_pose.PoseLandmark.LEFT_HIP]
                p2 = lm[mp_pose.PoseLandmark.RIGHT_HIP]
            elif category == "footwear":
                p1 = lm[mp_pose.PoseLandmark.LEFT_ANKLE]
                p2 = lm[mp_pose.PoseLandmark.RIGHT_ANKLE]
            else:
                continue
            
            x1 = int(p1.x * w_frame)
            y1 = int(p1.y * h_frame)
            x2 = int(p2.x * w_frame)
            y2 = int(p2.y * h_frame)
            
            body_width = int(np.hypot(x1 - x2, y1 - y2))
            
            if body_width > 20:
                scale = SCALE_FACTOR.get(category, 1.5)
                cloth_w = int(body_width * scale)
                cloth_h = int(cloth_w * cloth_img.shape[0] / cloth_img.shape[1])
                
                center_x = (x1 + x2) // 2
                center_y = (y1 + y2) // 2
                
                pos_x = center_x - cloth_w // 2
                pos_y = center_y + int(cloth_h * OFFSET_Y.get(category, 0))
                
                frame = overlay_transparent(frame, cloth_img, pos_x, pos_y, cloth_w, cloth_h)
        
        # ✅ OPTIMISATION 5 : Encoder avec qualité modérée
        _, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 85])
        encoded = base64.b64encode(buffer).decode('utf-8')
        
        return {"success": True, "frame": encoded}
        
    except Exception as e:
        print(f"   ❌ VTO Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return {"success": False, "error": str(e)}

def clear_cache():
    """Vide le cache des vêtements"""
    global _CLOTHES_CACHE
    _CLOTHES_CACHE.clear()
    download_image_cached.cache_clear()
    print("✅ Cache cleared")