Update bald_processor.py

bald_processor.py

@@ -10,11 +10,9 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# ---------------- DEVICE ----------------
 device = "cuda" if torch.cuda.is_available() else "cpu"
 logger.info(f"Using device: {device}")
 
-# ---------------- MODEL ----------------
 try:
     logger.info("Loading SegFormer face-parsing model...")
     processor = SegformerImageProcessor.from_pretrained("jonathandinu/face-parsing")
@@ -26,39 +24,38 @@ except Exception as e:
     logger.error(f"Failed to load model: {e}", exc_info=True)
     raise RuntimeError("SegFormer model load failed!")
 
-# ---------------- CLASS IDS ----------------
 hair_class_id = 13
 ear_class_ids = [7, 8]
 skin_class_id = 1
 nose_class_id = 2  # fallback
 
-# ---------------- CORE FUNCTION ----------------
 def make_realistic_bald(input_image: Image.Image) -> Image.Image:
     """
     Takes PIL Image, returns PIL Image bald version.
+    Only bald area is modified; rest of image stays sharp.
     """
     if input_image is None:
         raise ValueError("No input image provided!")
 
     try:
-        # -------- ORIGINAL IMAGE & RESIZE --------
         orig_w, orig_h = input_image.size
         original_np = np.array(input_image)
         original_bgr = cv2.cvtColor(original_np, cv2.COLOR_RGB2BGR)
 
+        # ---------------- RESIZE FOR PROCESSING ----------------
         MAX_DIM = 2048
         scale_factor = 1.0
-        working_np = original_np
-        working_bgr = original_bgr
+        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_np = cv2.resize(original_np, (working_w, working_h), cv2.INTER_AREA)
+            working_np = cv2.resize(original_np, (working_w, working_h), interpolation=cv2.INTER_AREA)
             working_bgr = cv2.cvtColor(working_np, cv2.COLOR_RGB2BGR)
 
-        # -------- SEGMENTATION --------
+        # ---------------- SEGMENTATION ----------------
         pil_working = Image.fromarray(working_np)
         inputs = processor(images=pil_working, return_tensors="pt").to(device)
         with torch.no_grad():
@@ -71,15 +68,15 @@ def make_realistic_bald(input_image: Image.Image) -> Image.Image:
         )
         parsing = upsampled.argmax(dim=1).squeeze(0).cpu().numpy()
 
-        # -------- HAIR & EARS MASK --------
+        # ---------------- 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
 
-        # Smooth & clean hair mask
+        # Morphology
         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)
@@ -89,30 +86,27 @@ def make_realistic_bald(input_image: Image.Image) -> Image.Image:
         if hair_pixels < 50:
             raise ValueError("NO_HAIR_DETECTED")
 
-        # -------- INPAINTING --------
+        # ---------------- INPAINT ----------------
         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)
         inpainted_rgb = cv2.cvtColor(inpainted_bgr, cv2.COLOR_BGR2RGB)
 
-        # -------- ADD SUBTLE SKIN TEXTURE --------
+        # ---------------- ADD SUBTLE SKIN TEXTURE ----------------
         noise = np.random.normal(0,12,(working_h, working_w,3)).astype(np.float32)
-        blended = np.clip(inpainted_rgb + noise*0.7, 0,255).astype(np.uint8)
+        bald_area = np.clip(inpainted_rgb + noise*0.7, 0,255).astype(np.uint8)
 
-        # -------- PREPARE RESULT SMALL & FINAL MASK --------
-        result_small = working_np.copy()
-        final_mask = hair_mask.copy()
-        result_small[final_mask == 1] = blended[final_mask == 1]
-
-        # -------- FINAL COMPOSITING (NO BLUR FIX) --------
+        # ---------------- COMPOSITE BACK ON ORIGINAL IMAGE ----------------
         if scale_factor < 1.0:
-            # Upscale ONLY bald area
-            bald_up = cv2.resize(result_small, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4)
-            mask_up = cv2.resize(final_mask.astype(np.uint8), (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
-            result = original_np.copy()
-            result[mask_up == 1] = bald_up[mask_up == 1]
+            # Upscale bald area mask and content separately
+            bald_up = cv2.resize(bald_area, (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)
         else:
-            result = result_small
+            bald_up = bald_area
+            mask_up = hair_mask
+
+        result = original_np.copy()
+        result[mask_up==1] = bald_up[mask_up==1]
 
         return Image.fromarray(result)