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)}")