Update bald_processor.py

bald_processor.py

@@ -1,3 +1,4 @@
+# bald_processor_clean.py
 import cv2
 import torch
 import numpy as np
@@ -23,88 +24,65 @@ except Exception as e:
     raise RuntimeError("SegFormer model load failed!")
 
 hair_class_id = 13
-ear_class_ids = [8, 9]  # Corrected: 8 for left_ear, 9 for right_ear
-skin_class_id = 1  # Added for color correction reference
+ear_class_ids = [7, 8]
 
 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()
-        
+
+        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
         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
-        
+        hair_mask[ears_mask==1] = 0
+
         # Morphology to clean
-        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (13, 13))
+        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
         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)
-        
-        # Conditional color correction for large hair areas
-        if hair_pixels > 220000:
-            # Skin mask for reference
-            skin_mask = (parsing == skin_class_id).astype(np.uint8)
-            
-            # Reference mask: skin excluding hair area
-            ref_mask = skin_mask.copy()
-            ref_mask[hair_mask == 1] = 0
-            ref_mask = cv2.dilate(ref_mask, np.ones((5, 5), np.uint8), iterations=1)  # Slight expand
-            
-            # Mean colors (BGR)
-            ref_mean = cv2.mean(working_bgr, mask=ref_mask * 255)[:3]
-            inpainted_mean = cv2.mean(inpainted_bgr, mask=hair_mask * 255)[:3]
-            
-            # Color difference
-            color_diff = np.array(ref_mean) - np.array(inpainted_mean)
-            
-            # Adjust inpainted area
-            hair_mask_3ch = np.repeat(hair_mask[:, :, np.newaxis], 3, axis=2)
-            inpainted_bgr[hair_mask_3ch == 1] = np.clip(
-                inpainted_bgr[hair_mask_3ch == 1] + color_diff, 0, 255
-            ).astype(np.uint8)
-        
+        inpainted_bgr = cv2.inpaint(working_bgr, hair_mask*255, inpaintRadius=radius, flags=flag)
         inpainted_rgb = cv2.cvtColor(inpainted_bgr, cv2.COLOR_BGR2RGB)
-        
+
         # Upscale bald area if needed
         if scale_factor < 1.0:
             bald_up = cv2.resize(inpainted_rgb, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4)
@@ -112,18 +90,15 @@ def make_realistic_bald(input_image: Image.Image) -> Image.Image:
         else:
             bald_up = inpainted_rgb
             mask_up = hair_mask
-        
-        # Soft blending for edges
-        mask_up_float = cv2.GaussianBlur(mask_up.astype(np.float32) * 255, (21, 21), 0) / 255.0
-        
-        # Composite with alpha blend
-        result = (1 - mask_up_float[..., None]) * original_np + mask_up_float[..., None] * bald_up
-        result = np.clip(result, 0, 255).astype(np.uint8)
-        
+
+        # Composite only bald area, rest untouched
+        result = original_np.copy()
+        result[mask_up==1] = bald_up[mask_up==1]
+
         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)}")