bald_processor.py

import cv2
import torch
import numpy as np
from PIL import Image
import torch._dynamo
from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
from scipy.ndimage import label
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

device = "cuda" if torch.cuda.is_available() else "cpu"
logger.info(f"Using device: {device}")
torch._dynamo.config.suppress_errors = True
try:
    logger.info("Loading SegFormer face-parsing model...")
    processor = SegformerImageProcessor.from_pretrained("jonathandinu/face-parsing")
    model = SegformerForSemanticSegmentation.from_pretrained("jonathandinu/face-parsing")
    model.to(device)
    model.eval()

            
    logger.info("Model loaded successfully!")
except Exception as e:
    logger.error(f"Failed to load model: {e}", exc_info=True)
    raise RuntimeError("Model loading failed!")

hair_class_id = 13
ear_class_ids = [8, 9]

def make_realistic_bald(input_image: Image.Image) -> Image.Image:
    """
    Main function for Hugging Face Space / single image processing
    Input: PIL Image (RGB)
    Output: PIL Image (bald version)
    """
    if input_image is None:
        raise ValueError("No input image provided!")
    
    try:
        orig_w, orig_h = input_image.size
        original_np = np.array(input_image)
        original_bgr = cv2.cvtColor(original_np, cv2.COLOR_RGB2BGR)
        
        MAX_DIM = 2048
        scale_factor = 1.0
        working_np = original_np.copy()
        working_bgr = original_bgr.copy()
        working_h, working_w = orig_h, orig_w
        
        if max(orig_w, orig_h) > MAX_DIM:
            scale_factor = MAX_DIM / max(orig_w, orig_h)
            working_w = int(orig_w * scale_factor)
            working_h = int(orig_h * scale_factor)
            working_np = cv2.resize(original_np, (working_w, working_h), interpolation=cv2.INTER_AREA)
            working_bgr = cv2.cvtColor(working_np, cv2.COLOR_RGB2BGR)
        
        # ── Segmentation ────────────────────────────────────────
        pil_working = Image.fromarray(working_np)
        inputs = processor(images=pil_working, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
        
        upsampled_logits = torch.nn.functional.interpolate(
            logits, size=(working_h, working_w), mode="bilinear", align_corners=False
        )
        probs = torch.softmax(upsampled_logits, dim=1)
        hair_prob = probs[0, hair_class_id].cpu().numpy()
        parsing = upsampled_logits.argmax(dim=1).squeeze(0).cpu().numpy()
        
        hair_mask = (hair_prob > 0.55).astype(np.uint8)
        
        # ── Smart Ear Protection ───────────────────────────────
        ears_mask = np.zeros_like(hair_mask)
        for cls in ear_class_ids:
            ears_mask[parsing == cls] = 1
        
        ear_y, ear_x = np.where(ears_mask > 0)
        ears_protected = np.zeros_like(hair_mask)
        
        if len(ear_y) > 0:
            ear_top_y = ear_y.min()
            kernel_v = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 18))
            ears_protected = cv2.dilate(ears_mask, kernel_v, iterations=1)
            
            top_margin = 4
            top_start = max(0, ear_top_y - top_margin)
            
            ear_x_min, ear_x_max = ear_x.min(), ear_x.max()
            ear_width = ear_x_max - ear_x_min + 1
            x_margin = int(ear_width * 0.25)
            protected_left = max(0, ear_x_min - x_margin)
            protected_right = min(working_w, ear_x_max + x_margin)
            
            limited_top_mask = np.zeros_like(ears_mask)
            limited_top_mask[top_start : ear_top_y + 5, protected_left:protected_right] = 1
            kernel_h = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (13, 3))
            limited_top_mask = cv2.dilate(limited_top_mask, kernel_h, iterations=1)
            
            ears_protected = np.logical_or(ears_protected, limited_top_mask).astype(np.uint8)
            
            hair_above_ears = np.zeros_like(hair_mask)
            above_ear_line = max(0, ear_top_y - 6)
            hair_above_ears[:above_ear_line, :] = hair_mask[:above_ear_line, :]
            ears_protected[hair_above_ears == 1] = 0
        
        hair_mask[ears_protected == 1] = 0
        
        # Forehead boost (if hair detected high)
        if hair_mask[:int(working_h * 0.25), :].sum() > 60:
            hair_mask[:int(working_h * 0.25), :] = np.maximum(
                hair_mask[:int(working_h * 0.25), :],
                (hair_prob[:int(working_h * 0.25), :] > 0.35).astype(np.uint8)
            )
        
        # Cleanup small noise
        kernel_clean = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        hair_mask = cv2.morphologyEx(hair_mask, cv2.MORPH_OPEN, kernel_clean, iterations=2)
        
        # Keep only largest hair component
        labeled, num_features = label(hair_mask)
        if num_features > 0:
            sizes = np.bincount(labeled.ravel())[1:]
            if len(sizes) > 0:
                largest_label = sizes.argmax() + 1
                hair_mask = (labeled == largest_label).astype(np.uint8)
        
        # Final mask refinement
        kernel_s = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9))
        hair_mask = cv2.dilate(hair_mask, kernel_s, iterations=1)
        blurred = cv2.GaussianBlur(hair_mask.astype(np.float32), (9, 9), 3)
        hair_mask_final = (blurred > 0.28).astype(np.uint8)
        
        # Extra fine hair catch
        kernel_tiny = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        hair_mask_final = cv2.dilate(hair_mask_final, kernel_tiny, iterations=1)
        blurred_extra = cv2.GaussianBlur(hair_mask_final.astype(np.float32), (7, 7), 2)
        hair_mask_final = (blurred_extra > 0.22).astype(np.uint8)
        
        hair_pixels = np.sum(hair_mask_final)
        if hair_pixels < 50:
            raise ValueError("NO_HAIR_DETECTED")
        
        logger.info(f"Hair pixels detected: {hair_pixels:,}")
        
        # Extended mask for very dense hair
        final_mask = hair_mask_final.copy()
        use_extended = False
        density = hair_pixels / (working_h * working_w) if working_h * working_w > 0 else 0
        
        if density > 0.18 or hair_pixels > 350000:
            use_extended = True
            big_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
            extended = cv2.dilate(hair_mask_final, big_kernel, iterations=1)
            
            upper = np.zeros_like(hair_mask_final)
            upper_end = int(working_h * 0.40)
            upper[:upper_end, :] = 1
            extended = np.logical_or(extended, upper).astype(np.uint8)
            
            extended[ears_protected == 1] = 0
            
            # Avoid sky/background
            hsv = cv2.cvtColor(working_np, cv2.COLOR_RGB2HSV)
            if np.mean(working_np) > 120:
                skyish = (hsv[:,:,0] > 70) & (hsv[:,:,0] < 160) & (hsv[:,:,1] < 80) & (hsv[:,:,2] > 160)
                extended[skyish] = 0
            
            extended = cv2.morphologyEx(extended, cv2.MORPH_CLOSE, kernel_s, iterations=1)
            extended[int(working_h * 0.70):, :] = 0
            extended = cv2.erode(extended, kernel_clean, iterations=2)
            final_mask = extended
        
        # Adaptive inpainting
        if use_extended or hair_pixels > 250000:
            radius, flag = 18, cv2.INPAINT_TELEA
        elif hair_pixels > 150000:
            radius, flag = 16, cv2.INPAINT_TELEA
        else:
            radius, flag = 12, cv2.INPAINT_NS
        
        inpainted_bgr = cv2.inpaint(working_bgr, final_mask * 255, inpaintRadius=radius, flags=flag)
        inpainted_rgb = cv2.cvtColor(inpainted_bgr, cv2.COLOR_BGR2RGB)
        
        result = working_np.copy()
        result[final_mask == 1] = inpainted_rgb[final_mask == 1]
        
        # Color correction for large areas
        if use_extended or hair_pixels > 250000:
            regions = [(0.20, 0.32, 0.35, 0.65), (0.35, 0.50, 0.30, 0.70)]
            colors = []
            for y1r, y2r, x1r, x2r in regions:
                y1 = int(working_h * y1r)
                y2 = int(working_h * y2r)
                x1 = int(working_w * x1r)
                x2 = int(working_w * x2r)
                if y2 > y1 + 50 and x2 > x1 + 100:
                    crop = working_np[y1:y2, x1:x2]
                    if crop.size > 0:
                        colors.append(np.median(crop, axis=(0,1)).astype(np.float32))
            
            if colors:
                target_color = np.mean(colors, axis=0)
                brightness = np.mean(target_color)
                strength = 0.9 if brightness > 140 else 0.7 if brightness < 90 else 0.8
                
                bald_area = result[final_mask == 1].astype(np.float32)
                if len(bald_area) > 500:
                    current_mean = bald_area.mean(axis=0)
                    diff = target_color - current_mean
                    corrected = np.clip(bald_area + diff * strength, 0, 255).astype(np.uint8)
                    result[final_mask == 1] = corrected
        
        # Edge feathering for smooth transition
        if hair_pixels > 80000:
            mask_feather = cv2.GaussianBlur(final_mask.astype(np.float32)*255, (21, 21), 12)
            mask_feather = (mask_feather > 35).astype(np.uint8)
            blurred_result = cv2.GaussianBlur(result, (9, 9), 3)
            result[mask_feather == 1] = cv2.addWeighted(
                result[mask_feather == 1], 0.5,
                blurred_result[mask_feather == 1], 0.5, 0
            )
        
        # Upscale back if needed
        if scale_factor < 1.0:
            result = cv2.resize(result, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4)
        
        return Image.fromarray(result)
    
    except Exception as e:
        logger.error(f"Processing failed: {str(e)}", exc_info=True)
        raise RuntimeError(f"Bald processing failed: {str(e)}")