Update bald_processor.py

bald_processor.py

@@ -1,4 +1,3 @@
-# bald_processor_clean.py
 import cv2
 import torch
 import numpy as np
@@ -24,81 +23,114 @@ except Exception as e:
     raise RuntimeError("SegFormer model load failed!")
 
 hair_class_id = 13
-ear_class_ids = [7, 8]
+ear_class_ids = [8, 9]   # l_ear=8, r_ear=9 (confirmed from model card)
+skin_class_id = 1        # skin
 
 def make_realistic_bald(input_image: Image.Image) -> Image.Image:
     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, working_h = int(orig_w*scale_factor), int(orig_h*scale_factor)
+            working_w, working_h = int(orig_w * scale_factor), 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 = torch.nn.functional.interpolate(
-            logits, size=(working_h, working_w),
-            mode="bilinear", align_corners=False
-        )
-        parsing = upsampled.argmax(dim=1).squeeze(0).cpu().numpy()
-
-        # Hair mask
+            upsampled = torch.nn.functional.interpolate(
+                logits, size=(working_h, working_w), mode="bilinear", align_corners=False
+            )
+            parsing = upsampled.argmax(dim=1).squeeze(0).cpu().numpy()
+        
+        # Masks
         hair_mask = (parsing == hair_class_id).astype(np.uint8)
         ears_mask = np.zeros_like(hair_mask)
         for cls in ear_class_ids:
             ears_mask[parsing == cls] = 1
-        hair_mask[ears_mask==1] = 0
-
-        # Morphology to clean
-        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (13,13))
+        hair_mask[ears_mask == 1] = 0
+        
+        # Clean mask
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (13, 13))
         hair_mask = cv2.morphologyEx(hair_mask, cv2.MORPH_CLOSE, kernel, iterations=2)
         hair_mask = cv2.dilate(hair_mask, kernel, iterations=1)
-        hair_mask = (cv2.GaussianBlur(hair_mask.astype(np.float32), (5,5), 0) > 0.28).astype(np.uint8)
-
+        hair_mask = (cv2.GaussianBlur(hair_mask.astype(np.float32), (5, 5), 0) > 0.28).astype(np.uint8)
+        
         hair_pixels = np.sum(hair_mask)
         if hair_pixels < 50:
             raise ValueError("NO_HAIR_DETECTED")
-
-        # Inpainting (no extra noise, no blur)
-        radius = 15 if hair_pixels > 220000 else 10
+        
+        # Inpainting
+        radius = 20 if hair_pixels > 220000 else 12   # Slightly larger for big areas
         flag = cv2.INPAINT_TELEA if hair_pixels > 220000 else cv2.INPAINT_NS
-        inpainted_bgr = cv2.inpaint(working_bgr, hair_mask*255, inpaintRadius=radius, flags=flag)
+        inpainted_bgr = cv2.inpaint(working_bgr, hair_mask * 255, inpaintRadius=radius, flags=flag)
+        
+        # Large hair: Color correction + Seamless blending
+        if hair_pixels > 220000:
+            skin_mask = (parsing == skin_class_id).astype(np.uint8)
+            ref_mask = skin_mask.copy()
+            ref_mask[hair_mask == 1] = 0
+            ref_mask = cv2.erode(ref_mask, np.ones((3,3), np.uint8), iterations=1)  # Avoid border artifacts
+            
+            ref_mean = cv2.mean(working_bgr, mask=ref_mask * 255)[:3]
+            inp_mean = cv2.mean(inpainted_bgr, mask=hair_mask * 255)[:3]
+            
+            if ref_mean != (0,0,0):  # Valid reference
+                color_diff = np.array(ref_mean) - np.array(inp_mean)
+                color_diff = np.clip(color_diff * 0.7, -40, 40)  # Softer adjustment to avoid over-correction
+                
+                hair_3ch = np.repeat(hair_mask[..., None], 3, axis=2)
+                inpainted_bgr[hair_3ch == 1] = np.clip(
+                    inpainted_bgr[hair_3ch == 1].astype(np.float32) + color_diff, 0, 255
+                ).astype(np.uint8)
+            
+            # Seamless clone for better merge (MIXED_CLONE good for skin)
+            moments = cv2.moments(hair_mask * 255)
+            if moments["m00"] != 0:
+                cx = int(moments["m10"] / moments["m00"])
+                cy = int(moments["m01"] / moments["m00"])
+            else:
+                cx, cy = working_w // 2, working_h // 2
+            
+            inpainted_bgr = cv2.seamlessClone(
+                inpainted_bgr, working_bgr, hair_mask * 255, (cx, cy), cv2.MIXED_CLONE
+            )
+        
         inpainted_rgb = cv2.cvtColor(inpainted_bgr, cv2.COLOR_BGR2RGB)
-
-        # Upscale bald area if needed
+        
+        # Upscale if downscaled
         if scale_factor < 1.0:
             bald_up = cv2.resize(inpainted_rgb, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4)
-            mask_up = cv2.resize(hair_mask.astype(np.uint8), (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
+            mask_up = cv2.resize(hair_mask, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
         else:
             bald_up = inpainted_rgb
             mask_up = hair_mask
-
-        # Composite only bald area, rest untouched
-        result = original_np.copy()
-        result[mask_up==1] = bald_up[mask_up==1]
-
+        
+        # Soft alpha blend (always, for smooth edges)
+        mask_float = cv2.GaussianBlur(mask_up.astype(np.float32) * 255, (25, 25), 0) / 255.0
+        result = (1 - mask_float[..., None]) * original_np.astype(np.float32) + \
+                 mask_float[..., None] * bald_up.astype(np.float32)
+        result = np.clip(result, 0, 255).astype(np.uint8)
+        
         return Image.fromarray(result)
-
+    
     except UnidentifiedImageError:
         raise ValueError("Invalid image format or corrupt image!")
     except Exception as e:
         logger.error(f"Bald processing failed: {str(e)}", exc_info=True)
-        raise RuntimeError(f"Bald processing failed: {str(e)}")
+        raise RuntimeError(f"Bald processing failed: {str(e)}")