Add hair color changer with segmentation model

Implement hair color changing functionality using semantic segmentation.

app.py

@@ -0,0 +1,283 @@
+import os
+import re
+import numpy as np
+import cv2
+import torch
+import torch.nn.functional as F
+from PIL import Image
+import gradio as gr
+from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation
+
+# =========================
+# CONFIG
+# =========================
+MODEL_ID = "jonathandinu/face-parsing"
+HAIR_ID = 13  # clase "hair" en este modelo
+
+# Presets de color (puedes editar)
+COLOR_PRESETS = {
+    "Personalizado (picker)": None,
+    "Negro": "#121212",
+    "Castaño": "#4b2e1f",
+    "Rubio": "#d8c27a",
+    "Platinado": "#d9d9d9",
+    "Rojo": "#c1121f",
+    "Azul": "#0077b6",
+    "Verde": "#2a9d8f",
+    "Morado": "#7209b7",
+    "Rosa": "#ff4d8d",
+}
+
+def get_device():
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+DEVICE = get_device()
+
+# Recomendado en Spaces para evitar timeouts raros al bajar modelos
+os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
+os.environ.setdefault("HF_HUB_READ_TIMEOUT", "60")
+os.environ.setdefault("HF_HUB_CONNECT_TIMEOUT", "30")
+
+# Limita threads en CPU (opcional, mejora estabilidad)
+try:
+    torch.set_num_threads(min(4, os.cpu_count() or 1))
+except Exception:
+    pass
+
+# =========================
+# LOAD MODEL (una sola vez)
+# =========================
+processor = AutoImageProcessor.from_pretrained(MODEL_ID)
+model = AutoModelForSemanticSegmentation.from_pretrained(MODEL_ID).to(DEVICE)
+model.eval()
+
+# =========================
+# UTIL: color parsing robusto
+# =========================
+def parse_color_to_rgb(color):
+    """
+    Acepta:
+      - "#RRGGBB"
+      - "#RRGGBBAA" (ignora AA)
+      - "#RGB"
+      - "rgb(r,g,b)" / "rgba(r,g,b,a)"
+      - (r,g,b) o [r,g,b]
+      - dict con {"hex": "..."} (por si acaso)
+    Devuelve (r,g,b) en 0..255
+    """
+    if color is None:
+        return (255, 0, 0)
+
+    if isinstance(color, dict):
+        color = color.get("hex") or color.get("value") or color.get("color")
+
+    if isinstance(color, (tuple, list)) and len(color) >= 3:
+        return (int(color[0]), int(color[1]), int(color[2]))
+
+    if not isinstance(color, str):
+        raise ValueError(f"Formato de color no soportado: {type(color)} -> {color}")
+
+    s = color.strip()
+
+    # rgb/rgba(...)
+    m = re.match(r"rgba?\(\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)", s.lower())
+    if m:
+        r, g, b = map(int, m.groups())
+        r = max(0, min(255, r))
+        g = max(0, min(255, g))
+        b = max(0, min(255, b))
+        return (r, g, b)
+
+    # hex
+    if s.startswith("#"):
+        h = s[1:]
+        if len(h) == 3:      # #RGB -> #RRGGBB
+            h = "".join([c * 2 for c in h])
+        if len(h) == 8:      # #RRGGBBAA -> ignora AA
+            h = h[:6]
+        if len(h) != 6:
+            raise ValueError(f"HEX inválido: {s} (usa #RRGGBB)")
+        r = int(h[0:2], 16)
+        g = int(h[2:4], 16)
+        b = int(h[4:6], 16)
+        return (r, g, b)
+
+    raise ValueError(f"Color inválido: {color}")
+
+# =========================
+# IMAGE UTILS
+# =========================
+def resize_keep_aspect(pil: Image.Image, max_side: int) -> Image.Image:
+    w, h = pil.size
+    m = max(w, h)
+    if m <= max_side:
+        return pil
+    scale = max_side / float(m)
+    nw, nh = max(1, int(w * scale)), max(1, int(h * scale))
+    return pil.resize((nw, nh), Image.BILINEAR)
+
+@torch.inference_mode()
+def get_hair_mask(image: Image.Image, max_side: int = 640) -> Image.Image:
+    """
+    Devuelve una máscara L (0..255) del cabello, al tamaño original.
+    """
+    image = image.convert("RGB")
+    ow, oh = image.size
+
+    infer_img = resize_keep_aspect(image, max_side=max_side)
+
+    inputs = processor(images=infer_img, return_tensors="pt")
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+
+    outputs = model(**inputs)
+    logits = outputs.logits  # (B,C,h,w)
+
+    up = F.interpolate(
+        logits,
+        size=infer_img.size[::-1],  # (H,W)
+        mode="bilinear",
+        align_corners=False,
+    )
+
+    labels = up.argmax(dim=1)[0]  # (H,W)
+    hair = (labels == HAIR_ID).to(torch.uint8).cpu().numpy() * 255
+
+    mask = Image.fromarray(hair, mode="L")
+
+    if mask.size != (ow, oh):
+        mask = mask.resize((ow, oh), Image.NEAREST)
+
+    return mask
+
+def refine_mask(mask: Image.Image, close_kernel: int = 9, feather: int = 9) -> Image.Image:
+    m = np.array(mask.convert("L"))
+
+    # binariza
+    _, mb = cv2.threshold(m, 127, 255, cv2.THRESH_BINARY)
+
+    # close
+    k = max(3, int(close_kernel) | 1)  # impar
+    kernel = np.ones((k, k), np.uint8)
+    mb = cv2.morphologyEx(mb, cv2.MORPH_CLOSE, kernel, iterations=1)
+
+    # feather (blur)
+    f = max(1, int(feather))
+    if f % 2 == 0:
+        f += 1
+    mb = cv2.GaussianBlur(mb, (f, f), 0)
+
+    return Image.fromarray(mb, mode="L")
+
+def recolor_hair_lab(
+    image: Image.Image,
+    mask: Image.Image,
+    color_input,
+    strength: float = 0.85,
+    brighten: float = 0.0,
+) -> Image.Image:
+    """
+    Recolor en LAB para mantener sombras/luces.
+    strength: 0..1 confirmando cuánto entra el color
+    brighten: -0.3..0.3 (opcional, solo en cabello)
+    """
+    image_rgb = np.array(image.convert("RGB"))
+    mask_f = np.array(mask.convert("L")).astype(np.float32) / 255.0
+    alpha = np.clip(mask_f * float(strength), 0.0, 1.0)[..., None]  # (H,W,1)
+
+    # RGB -> BGR -> LAB
+    bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
+    lab = cv2.cvtColor(bgr, cv2.COLOR_BGR2LAB).astype(np.float32)
+
+    # color objetivo -> LAB
+    r, g, b = parse_color_to_rgb(color_input)
+    target_bgr = np.array([[[b, g, r]]], dtype=np.uint8)
+    target_lab = cv2.cvtColor(target_bgr, cv2.COLOR_BGR2LAB).astype(np.float32)[0, 0]
+
+    # Mezcla a/b hacia el objetivo
+    lab[:, :, 1] = lab[:, :, 1] * (1.0 - alpha[:, :, 0]) + target_lab[1] * alpha[:, :, 0]
+    lab[:, :, 2] = lab[:, :, 2] * (1.0 - alpha[:, :, 0]) + target_lab[2] * alpha[:, :, 0]
+
+    # Ajuste de brillo en cabello
+    if abs(brighten) > 1e-6:
+        lab[:, :, 0] = np.clip(lab[:, :, 0] + (brighten * 255.0) * alpha[:, :, 0], 0, 255)
+
+    lab_u8 = np.clip(lab, 0, 255).astype(np.uint8)
+    out_bgr = cv2.cvtColor(lab_u8, cv2.COLOR_LAB2BGR)
+    out_rgb = cv2.cvtColor(out_bgr, cv2.COLOR_BGR2RGB)
+
+    return Image.fromarray(out_rgb)
+
+# =========================
+# GRADIO RUN
+# =========================
+def run(image, preset, picked_color, strength, brighten, max_side, close_kernel, feather):
+    try:
+        if image is None:
+            return None, None, "Sube una imagen primero."
+
+        # color final
+        preset_hex = COLOR_PRESETS.get(preset)
+        final_color = (picked_color or "#ff0000") if preset_hex is None else preset_hex
+
+        # máscara
+        raw_mask = get_hair_mask(image, max_side=int(max_side))
+        mask = refine_mask(raw_mask, close_kernel=int(close_kernel), feather=int(feather))
+
+        # si la máscara salió vacía
+        if np.mean(np.array(mask)) < 2.0:
+            return image, mask, "No detecté cabello en esta foto. Prueba otra (mejor luz/frente)."
+
+        # recolor
+        result = recolor_hair_lab(
+            image=image,
+            mask=mask,
+            color_input=final_color,
+            strength=float(strength),
+            brighten=float(brighten),
+        )
+
+        return result, mask, f"OK ✅ Color aplicado: {final_color}"
+
+    except Exception as e:
+        # devuelve el error visible en la app
+        return None, None, f"ERROR: {type(e).__name__}: {e}"
+
+DESCRIPTION = """
+Sube una foto y cambia el color del cabello.
+- Segmentación de cabello (hair mask)
+- Recolor en LAB para conservar sombras/luces
+"""
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 🎨 Cambiar color de cabello")
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row():
+        inp = gr.Image(label="Tu foto", type="pil")
+        out = gr.Image(label="Resultado", type="pil")
+
+    with gr.Accordion("Controles", open=True):
+        preset = gr.Dropdown(
+            label="Preset",
+            choices=list(COLOR_PRESETS.keys()),
+            value="Personalizado (picker)",
+        )
+        picked_color = gr.ColorPicker(label="Color personalizado", value="#ff0000")
+        strength = gr.Slider(0.0, 1.0, value=0.85, step=0.05, label="Intensidad")
+        brighten = gr.Slider(-0.3, 0.3, value=0.0, step=0.05, label="Brillo cabello (opcional)")
+        max_side = gr.Slider(384, 1024, value=640, step=64, label="Resolución segmentación")
+        close_kernel = gr.Slider(3, 21, value=9, step=2, label="Cerrar huecos (máscara)")
+        feather = gr.Slider(1, 31, value=9, step=2, label="Suavizado bordes (máscara)")
+
+    btn = gr.Button("Aplicar")
+    mask_out = gr.Image(label="Máscara (debug)", type="pil")
+    status = gr.Textbox(label="Estado", value="Listo.")
+
+    btn.click(
+        fn=run,
+        inputs=[inp, preset, picked_color, strength, brighten, max_side, close_kernel, feather],
+        outputs=[out, mask_out, status],
+    )
+
+if __name__ == "__main__":
+    demo.launch()