fix: correctly preprocess input image with LineartDetector before passing to Stable Diffusion ControlNet to eliminate ripples

faceless_processing.py

@@ -1,26 +1,19 @@
-"""
-faceless_lineart_pro.py - VERSÃO CORRIGIDA E SANITIZADA
-"""
-
-from __future__ import annotations
 import base64
 import io
-import logging
+import os
+import tempfile
+from typing import Tuple, Optional
+from dataclasses import dataclass
+
 import cv2
 import numpy as np
 import torch
-from dataclasses import dataclass
-from functools import lru_cache
-from pathlib import Path
-from typing import Optional
-from PIL import Image, ImageFilter
+from PIL import Image, ImageChops, ImageFilter, ImageOps, ImageStat
+from rembg import remove, new_session
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
 
-# ── Logging ──────────────────────────────────────────────────────────────────
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# ── Variável Global ──────────────────────────────────────────────────────────
+_SESSION_ISNET = new_session("isnet-general-use")
+_LINEART_DETECTOR = None
 _SD_PIPELINE = None
 
 @dataclass(frozen=True)
@@ -36,7 +29,6 @@ class LineartConfig:
     prompt: str = "technical line art, clean vector, simple continuous lines, white background"
     negative_prompt: str = "shading, realism, messy lines, noisy, double lines, background clutter"
 
-# ── Carregamento Otimizado ───────────────────────────────────────────────────
 def get_sd_pipeline(cfg: LineartConfig):
     global _SD_PIPELINE
     if _SD_PIPELINE is not None:
@@ -44,7 +36,7 @@ def get_sd_pipeline(cfg: LineartConfig):
         
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float16 if device == "cuda" else torch.float32
-    logger.info(f"Carregando IA no dispositivo: {device} com precisao: {dtype}...")
+    print(f"[AI] Loading SD pipeline on {device} with {dtype}...")
     try:
         controlnet = ControlNetModel.from_pretrained(
             cfg.controlnet_model_id, 
@@ -61,64 +53,351 @@ def get_sd_pipeline(cfg: LineartConfig):
             _SD_PIPELINE.enable_attention_slicing()
         return _SD_PIPELINE
     except Exception as exc:
-        logger.error(f"Falha ao carregar IA no dispositivo {device}: {exc}")
+        print(f"[AI] Failed to load SD pipeline on {device}: {exc}")
         return None
 
-# ── Módulos Extras ───────────────────────────────────────────────────────────
-try:
-    from rembg import remove, new_session as _rembg_new_session
-    _REMBG_AVAILABLE = True
-except ImportError:
-    remove, _REMBG_AVAILABLE = None, False
 
-@lru_cache(maxsize=2)
-def _get_rembg_session(model_name: str):
-    if not _REMBG_AVAILABLE: return None
-    try: return _rembg_new_session(model_name)
-    except Exception: return None
+def decode_base64_image(image_data: str) -> Image.Image:
+    payload = image_data.split(",", 1)[1] if "," in image_data else image_data
+    decoded = base64.b64decode(payload)
+    return Image.open(io.BytesIO(decoded)).convert("RGBA")
 
-# ── Lógica de Processamento ──────────────────────────────────────────────────
-def make_faceless_lineart(img: Image.Image, cfg: LineartConfig = LineartConfig()) -> Image.Image:
-    # 1. Background
+
+def encode_png_data_url(image: Image.Image) -> str:
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    return f"data:image/png;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
+
+
+def _has_real_transparency(img: Image.Image) -> bool:
     rgba = img.convert("RGBA")
-    session = _get_rembg_session(cfg.rembg_model)
-    if session and remove:
-        mask = remove(rgba, session=session, only_mask=True)
-        rgba.putalpha(mask.convert("L"))
-    
-    # 2. Resize para 512x512 para o SD
-    rgb_input = Image.new("RGB", (512, 512), (255, 255, 255))
-    resized_subject = rgba.resize((512, 512), Image.Resampling.LANCZOS)
-    rgb_input.paste(resized_subject, mask=resized_subject.split()[3])
-
-    # 3. IA Inference
-    pipe = get_sd_pipeline(cfg)
-    if pipe:
-        logger.info("Executando IA...")
-        generated = pipe(
-            cfg.prompt,
-            negative_prompt=cfg.negative_prompt,
-            image=rgb_input,
-            controlnet_conditioning_scale=cfg.controlnet_conditioning_scale,
-            num_inference_steps=cfg.num_inference_steps
-        ).images[0]
-    else:
-        generated = rgb_input
-
-    # 4. Binário e Afinamento
-    gray = cv2.cvtColor(np.array(generated), cv2.COLOR_RGB2GRAY)
-    _, binary = cv2.threshold(gray, 120, 255, cv2.THRESH_BINARY_INV)
-    
+    alpha = rgba.getchannel("A")
+    sample = list(alpha.getdata())
+    total = 0
+    visible = 0
+    for index in range(0, len(sample), 12):
+        total += 1
+        if sample[index] < 245:
+            visible += 1
+    return total > 0 and (visible / total) >= 0.03
+
+
+def remove_background_subject(img: Image.Image) -> Image.Image:
+    rgba = img.convert("RGBA")
+    if _has_real_transparency(rgba):
+        return rgba
+
+    mask = remove(rgba, session=_SESSION_ISNET, only_mask=True)
+    alpha = mask.convert("L") if isinstance(mask, Image.Image) else Image.open(io.BytesIO(mask)).convert("L")
+    if alpha.size != rgba.size:
+        alpha = alpha.resize(rgba.size, Image.LANCZOS)
+    alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.8))
+
+    r, g, b, _a = rgba.split()
+    return Image.merge("RGBA", (r, g, b, alpha))
+
+
+def _get_lineart_detector():
+    global _LINEART_DETECTOR
+    if _LINEART_DETECTOR is not None:
+        return _LINEART_DETECTOR
+
     try:
-        binary = cv2.ximgproc.thinning(binary, thinningType=cv2.ximgproc.THINNING_ZHANGSUEN)
-    except: pass
+        from controlnet_aux import LineartDetector
+
+        _LINEART_DETECTOR = LineartDetector.from_pretrained("lllyasviel/Annotators")
+    except Exception as exc:
+        print(f"[AI] LineartDetector unavailable: {exc}")
+        _LINEART_DETECTOR = False
+    return _LINEART_DETECTOR or None
+
+
+def _box_blur_rgba_region(image: Image.Image, top_ratio: float, height_ratio: float, blur_radius: int) -> Image.Image:
+    rgba = image.convert("RGBA")
+    width, height = rgba.size
+    top = max(0, int(height * top_ratio))
+    region_height = max(1, int(height * height_ratio))
+    bottom = min(height, top + region_height)
+    if bottom <= top:
+        return rgba
+
+    blurred = rgba.filter(ImageFilter.GaussianBlur(radius=blur_radius))
+    result = rgba.copy()
+    region = blurred.crop((0, top, width, bottom))
+    result.paste(region, (0, top))
+    return result
+
+
+def _mask_to_transparent_lineart(lineart: Image.Image) -> Image.Image:
+    rgba = lineart.convert("RGBA")
+    gray = ImageOps.grayscale(rgba)
+    gray = ImageOps.autocontrast(gray)
+    # Increased MedianFilter and GaussianBlur to smooth out tiny clothing wrinkles and details
+    gray = gray.filter(ImageFilter.MedianFilter(5))
+    gray = gray.filter(ImageFilter.GaussianBlur(radius=0.75))
+
+    stat = ImageStat.Stat(gray)
+    mean = stat.mean[0] if stat.mean else 0.0
+    threshold = max(160, min(240, int(mean + 35)))
+    alpha = gray.point(lambda p: 255 if p < threshold else 0, mode="L")
+    alpha = alpha.filter(ImageFilter.MaxFilter(3))
+    alpha = alpha.filter(ImageFilter.MinFilter(3))
+
+    if alpha.getbbox() is None:
+        alpha = gray.point(lambda p: 255 if p < 245 else 0, mode="L")
+
+    # Clean up small busy lines/details using OpenCV contours
+    alpha_np = np.array(alpha)
+    contours, _ = cv2.findContours(alpha_np, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+    mask = np.zeros_like(alpha_np)
+    for cnt in contours:
+        # Keep only contours with perimeter larger than 35 pixels (keeps main lines only)
+        if cv2.arcLength(cnt, False) > 35:
+            cv2.drawContours(mask, [cnt], -1, 255, -1)
     
-    # 5. Final
-    final = Image.fromarray(binary).resize(img.size, Image.Resampling.NEAREST)
-    out = Image.new("RGBA", img.size, (0, 0, 0, 0))
-    out.putalpha(final)
+    alpha_filtered = cv2.bitwise_and(alpha_np, mask)
+    alpha = Image.fromarray(alpha_filtered)
+
+    out = Image.new("RGBA", alpha.size, (0, 0, 0, 0))
+    out.putalpha(alpha)
     return out
 
-def build_faceless_embroidery_assets(img: Image.Image, cfg: Optional[LineartConfig] = None) -> tuple[Image.Image, Image.Image]:
+
+def _merge_outline_and_detail(outline: Image.Image, detail: Image.Image) -> Image.Image:
+    outline_l = outline.convert("RGBA")
+    detail_l = detail.convert("RGBA")
+    outline_a = outline_l.getchannel("A")
+    detail_a = detail_l.getchannel("A")
+    combined = ImageChops.lighter(outline_a, detail_a)
+    combined = combined.filter(ImageFilter.MaxFilter(3))
+    combined = combined.filter(ImageFilter.MinFilter(3))
+    result = Image.new("RGBA", combined.size, (0, 0, 0, 0))
+    result.putalpha(combined)
+    return result
+
+
+def _alpha_outline(alpha: Image.Image) -> Image.Image:
+    alpha = alpha.convert("L")
+    alpha = ImageOps.autocontrast(alpha)
+    alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.4))
+    alpha = alpha.filter(ImageFilter.MedianFilter(3))
+
+    outer = alpha.filter(ImageFilter.MaxFilter(5))
+    inner = alpha.filter(ImageFilter.MinFilter(5))
+    outline = ImageChops.difference(outer, inner)
+    outline = ImageOps.autocontrast(outline)
+    outline = outline.point(lambda p: 255 if p >= 22 else 0, mode="L")
+    outline = outline.filter(ImageFilter.MaxFilter(3))
+    outline = outline.filter(ImageFilter.MinFilter(3))
+
+    if outline.getbbox() is None:
+        outline = alpha.point(lambda p: 255 if p >= 40 else 0, mode="L")
+        outline = outline.filter(ImageFilter.MaxFilter(3))
+
+    return outline
+
+
+def blur_faces_smooth(image: Image.Image) -> tuple[Image.Image, Image.Image, list]:
+    rgba = image.convert("RGBA")
+    width, height = rgba.size
+    open_cv_image = np.array(rgba.convert("RGB"))
+    open_cv_image = open_cv_image[:, :, ::-1].copy()
+    gray = cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2GRAY)
+    
+    face_cascade_default = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
+    face_cascade_alt = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_alt2.xml')
+    profile_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_profileface.xml')
+    
+    faces_default = face_cascade_default.detectMultiScale(gray, scaleFactor=1.05, minNeighbors=3, minSize=(20, 20))
+    faces_alt = face_cascade_alt.detectMultiScale(gray, scaleFactor=1.05, minNeighbors=3, minSize=(20, 20))
+    faces_profile = profile_cascade.detectMultiScale(gray, scaleFactor=1.05, minNeighbors=3, minSize=(20, 20))
+    
+    faces = []
+    if len(faces_default) > 0: faces.extend(faces_default)
+    if len(faces_alt) > 0: faces.extend(faces_alt)
+    if len(faces_profile) > 0: faces.extend(faces_profile)
+    
+    mask = Image.new("L", (width, height), 0)
+    if len(faces) == 0:
+        mask_np = np.array(mask)
+        cx, cy = width // 2, int(height * 0.32)
+        rx, ry = int(width * 0.25), int(height * 0.18)
+        cv2.ellipse(mask_np, (cx, cy), (rx, ry), 0, 0, 360, 255, -1)
+        mask = Image.fromarray(mask_np).filter(ImageFilter.GaussianBlur(radius=20))
+        blurred_img = rgba.filter(ImageFilter.GaussianBlur(radius=25))
+        return Image.composite(blurred_img, rgba, mask), mask, []
+
+    mask_np = np.array(mask)
+    for (x, y, w, h) in faces:
+        cx = x + w // 2
+        rx = int(w * 0.48)
+        ry = int(h * 0.32)
+        cy_shifted = y + int(h * 0.65)
+        cv2.ellipse(mask_np, (cx, cy_shifted), (rx, ry), 0, 0, 360, 255, -1)
+        
+    mask = Image.fromarray(mask_np)
+    mask_feathered = mask.filter(ImageFilter.GaussianBlur(radius=10))
+    blurred_img = rgba.filter(ImageFilter.GaussianBlur(radius=35))
+    return Image.composite(blurred_img, rgba, mask_feathered), mask, list(faces)
+
+
+def make_faceless_cutout(img: Image.Image) -> Image.Image:
+    return remove_background_subject(img)
+
+
+def make_faceless_lineart(img: Image.Image, cfg: LineartConfig = LineartConfig()) -> Image.Image:
+    subject = remove_background_subject(img)
+    
+    # 1. Get face erase stencil mask
+    _pre, face_mask, faces = blur_faces_smooth(subject)
+    
+    detector = _get_lineart_detector()
+    lineart_out = None
+
+    if detector is not None:
+        try:
+            # Blur the face/hair region to prevent ControlNet from seeing face details
+            preprocessed = _box_blur_rgba_region(subject, top_ratio=0.16, height_ratio=0.52, blur_radius=18)
+            
+            # Extract lineart preprocessed map (black and white lines)
+            detected = detector(preprocessed.convert("RGB"), coarse=False)
+            if not isinstance(detected, Image.Image):
+                detected = Image.open(io.BytesIO(detected))
+                
+            # Resize detected lines to 512x512 for Stable Diffusion ControlNet
+            detected_resized = detected.resize((512, 512), Image.Resampling.LANCZOS)
+            
+            # Run Stable Diffusion ControlNet to clean up and stylize the lineart
+            pipe = get_sd_pipeline(cfg)
+            if pipe:
+                print("[AI] Running Stable Diffusion to generate lineart...")
+                generated = pipe(
+                    cfg.prompt,
+                    negative_prompt=cfg.negative_prompt,
+                    image=detected_resized,
+                    controlnet_conditioning_scale=cfg.controlnet_conditioning_scale,
+                    num_inference_steps=cfg.num_inference_steps
+                ).images[0]
+                
+                # Resize generated back to original size
+                generated_resized = generated.resize(subject.size, Image.Resampling.LANCZOS)
+                cleaned = _mask_to_transparent_lineart(generated_resized)
+            else:
+                cleaned = _mask_to_transparent_lineart(detected)
+                
+            silhouette = _alpha_outline(subject.getchannel("A"))
+            lineart_out = _merge_outline_and_detail(silhouette, cleaned)
+        except Exception as exc:
+            print(f"[AI] SD/LineartDetector pipeline failed, using local outline fallback: {exc}")
+
+    if lineart_out is None:
+        try:
+            open_cv_image = np.array(subject.convert("RGB"))
+            open_cv_image = open_cv_image[:, :, ::-1].copy()
+            gray = cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2GRAY)
+            smoothed = cv2.bilateralFilter(gray, d=9, sigmaColor=75, sigmaSpace=75)
+            # Higher thresholds to ignore fabric textures and faint details
+            edges = cv2.Canny(smoothed, 45, 120)
+            
+            # Prune small busy lines using contour filtering
+            contours, _ = cv2.findContours(edges, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+            filtered_edges = np.zeros_like(edges)
+            for cnt in contours:
+                if cv2.arcLength(cnt, False) > 35:
+                    cv2.drawContours(filtered_edges, [cnt], -1, 255, 1)
+            edges = filtered_edges
+
+            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
+            edges_dilated = cv2.dilate(edges, kernel, iterations=1)
+            details_alpha = Image.fromarray(edges_dilated).convert("L")
+            
+            silhouette = _alpha_outline(subject.getchannel("A"))
+            combined = ImageChops.lighter(silhouette, details_alpha)
+            
+            subject_alpha = subject.getchannel("A")
+            combined = ImageChops.darker(combined, subject_alpha)
+            
+            lineart_out = Image.new("RGBA", combined.size, (0, 0, 0, 0))
+            lineart_out.putalpha(combined)
+        except Exception as local_err:
+            print(f"[AI] Local Canny fallback failed: {local_err}")
+            outline = _alpha_outline(subject.getchannel("A"))
+            lineart_out = Image.new("RGBA", outline.size, (0, 0, 0, 0))
+            lineart_out.putalpha(outline)
+
+    # 2. Erase face details in lineart using stencil mask
+    lineart_rgba = lineart_out.convert("RGBA")
+    lineart_alpha = lineart_rgba.getchannel("A")
+    black_img = Image.new("L", lineart_alpha.size, 0)
+    clean_alpha = Image.composite(black_img, lineart_alpha, face_mask)
+    
+    r, g, b, _ = lineart_rgba.split()
+    return Image.merge("RGBA", (r, g, b, clean_alpha))
+
+
+def try_huggingface_faceless(img: Image.Image, prompt: str, negative_prompt: str) -> Image.Image | None:
+    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
+    model_id = os.environ.get("HF_FACELESS_MODEL") or os.environ.get("HUGGINGFACE_FACELESS_MODEL")
+    space_id = os.environ.get("HF_SPACE_ID") or os.environ.get("HUGGINGFACE_SPACE_ID")
+
+    if not token or (not model_id and not space_id):
+        return None
+
+    try:
+        if space_id:
+            from gradio_client import Client, handle_file
+
+            api_name = os.environ.get("HF_SPACE_API_NAME") or os.environ.get("HUGGINGFACE_SPACE_API_NAME") or "/predict"
+            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+                img.save(temp_file.name, format="PNG")
+                temp_path = temp_file.name
+            try:
+                client = Client(space_id, token=token)
+                generated = client.predict(
+                    handle_file(temp_path),
+                    prompt,
+                    negative_prompt,
+                    api_name=api_name,
+                )
+            finally:
+                try:
+                    os.unlink(temp_path)
+                except OSError:
+                    pass
+            if isinstance(generated, str) and generated.startswith("data:image"):
+                return Image.open(io.BytesIO(base64.b64decode(generated.split(",", 1)[1]))).convert("RGBA")
+            if isinstance(generated, Image.Image):
+                return generated.convert("RGBA")
+            if isinstance(generated, (list, tuple)) and generated:
+                first = generated[0]
+                if isinstance(first, str) and first.startswith("data:image"):
+                    return Image.open(io.BytesIO(base64.b64decode(first.split(",", 1)[1]))).convert("RGBA")
+                if isinstance(first, Image.Image):
+                    return first.convert("RGBA")
+        else:
+            from huggingface_hub import InferenceClient
+
+            client = InferenceClient(model=model_id, token=token, timeout=float(os.environ.get("HF_TIMEOUT_SECONDS", "45")))
+            generated = client.image_to_image(
+                image=img,
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=20,
+                guidance_scale=7.0,
+            )
+            if isinstance(generated, Image.Image):
+                return generated.convert("RGBA")
+    except Exception as exc:
+        print(f"[AI] Hugging Face faceless processing failed: {exc}")
+    return None
+
+
+def build_faceless_embroidery_assets(
+    img: Image.Image,
+    cfg: Optional[LineartConfig] = None
+) -> Tuple[Image.Image, Image.Image]:
+    subject = remove_background_subject(img)
     cfg = cfg or LineartConfig()
-    return make_faceless_lineart(img, cfg), img.convert("RGBA")
+    lineart = make_faceless_lineart(subject, cfg)
+    cutout = make_faceless_cutout(subject)
+    return lineart, cutout