Crossfit Photo Filter - MediaPipe segmentation + blur background

crossfit_photo_filter.py

#!/usr/bin/env python3
"""
Script para procesar fotos de Crossfit:
- Mejora iluminación (CLAHE)
- Mejora color (saturación)
- Aplica nitidez (unsharp mask)
- Desenfoca fondo manteniendo nítida a la persona (MediaPipe segmentation)

Requiere: pip install rawpy opencv-python numpy imageio mediapipe
"""
import os
import imageio.v2 as imageio
import cv2
import numpy as np
import mediapipe as mp
from mediapipe.tasks.python.vision.image_segmenter import ImageSegmenter, ImageSegmenterOptions
from mediapipe.tasks.python.core.base_options import BaseOptions

# --- Configuración ---
ORIGEN  = os.path.expanduser("~/Documents/Crossfit_photos")
DESTINO = os.path.expanduser("~/Documents/Crossfit_photos_with_filter")
CALIDAD = 100
MODELO  = os.path.expanduser("~/models/selfie_segmenter.tflite")

os.makedirs(DESTINO, exist_ok=True)

# Inicializar MediaPipe Image Segmenter
base_options = BaseOptions(model_asset_path=MODELO)
options = ImageSegmenterOptions(base_options=base_options, output_confidence_masks=True)
segmenter = ImageSegmenter.create_from_options(options)

def mejorar_nitidez(img):
    """Unsharp mask para realzar detalles"""
    blur = cv2.GaussianBlur(img, (0, 0), 3)
    return cv2.addWeighted(img, 1.5, blur, -0.5, 0)

def mejorar_iluminacion(img):
    """CLAHE en canal L del espacio LAB"""
    lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    l = clahe.apply(l)
    lab = cv2.merge([l, a, b])
    return cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)

def mejorar_color(img):
    """Aumenta saturación levemente"""
    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV).astype(np.float32)
    hsv[:, :, 1] = np.clip(hsv[:, :, 1] * 1.2, 0, 255)
    return cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2RGB)

def desenfocar_fondo(img):
    """Usa MediaPipe ImageSegmenter para separar persona del fondo"""
    h, w = img.shape[:2]
    
    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=img)
    segmentation_result = segmenter.segment(mp_image)
    
    if not segmentation_result.confidence_masks:
        print("  ⚠ MediaPipe no detectó persona")
        return img
    
    mask = segmentation_result.confidence_masks[0].numpy_view()
    
    if mask.shape != (h, w):
        mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_LINEAR)
    
    mask_suave = cv2.GaussianBlur(mask.astype(np.float32), (31, 31), 0)
    mask_suave = np.clip(mask_suave, 0, 1)
    mask_suave = np.stack([mask_suave] * 3, axis=-1)
    
    fondo_blur = cv2.GaussianBlur(img, (51, 51), 0)
    
    resultado = (img * mask_suave + fondo_blur * (1 - mask_suave)).astype(np.uint8)
    return resultado

archivos = [f for f in os.listdir(ORIGEN) if f.lower().endswith(".jpg")]
total = len(archivos)
print(f"Procesando {total} fotos con MediaPipe segmentation...\n")

for i, archivo in enumerate(archivos, 1):
    ruta_origen  = os.path.join(ORIGEN, archivo)
    ruta_destino = os.path.join(DESTINO, archivo)
    
    # Saltar si ya existe
    if os.path.exists(ruta_destino):
        print(f"[{i}/{total}] {archivo} - ya existe, saltando")
        continue

    try:
        img = imageio.imread(ruta_origen)
        print(f"[{i}/{total}] {archivo} ({img.shape[1]}x{img.shape[0]})")

        print("  → Mejorando iluminación...")
        img = mejorar_iluminacion(img)

        print("  → Mejorando color...")
        img = mejorar_color(img)

        print("  → Aplicando nitidez...")
        img = mejorar_nitidez(img)

        print("  → Segmentando persona con MediaPipe...")
        img = desenfocar_fondo(img)

        imageio.imwrite(ruta_destino, img, quality=CALIDAD)
        print(f"  ✓ Guardado\n")

    except Exception as e:
        import traceback
        print(f"  ✗ Error: {e}")
        traceback.print_exc()

segmenter.close()
print(f"\nListo. Fotos con filtros en: {DESTINO}")