Create bald_processor.py

bald_processor.py

@@ -0,0 +1,241 @@
+# bald_processor.py
+import os
+import cv2
+import torch
+import numpy as np
+from PIL import Image, UnidentifiedImageError
+from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
+import io
+import traceback
+
+# Global model load with error handling
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device} | CUDA available: {torch.cuda.is_available()}")
+
+print("Loading SegFormer face-parsing model...")
+try:
+    processor = SegformerImageProcessor.from_pretrained("jonathandinu/face-parsing")
+    model = SegformerForSemanticSegmentation.from_pretrained("jonathandinu/face-parsing")
+    model.to(device)
+    model.eval()
+    print("Model loaded successfully!")
+except Exception as e:
+    print(f"CRITICAL: Model loading failed! {str(e)}")
+    traceback.print_exc()
+
+hair_class_id = 13
+ear_class_ids = [8, 9]   # l_ear=8, r_ear=9
+skin_class_id = 1
+nose_class_id = 2        # Reliable fallback for clean skin tone
+
+def make_realistic_bald(image_bytes: bytes) -> bytes:
+    try:
+        # Open image safely
+        try:
+            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        except UnidentifiedImageError:
+            raise ValueError("Invalid image format or corrupt bytes")
+        except Exception as e:
+            raise ValueError(f"Image open failed: {str(e)}")
+
+        orig_w, orig_h = image.size
+        original_np = np.array(image)
+        original_bgr = cv2.cvtColor(original_np, cv2.COLOR_RGB2BGR)
+
+        # Resize if large
+        MAX_PROCESS_DIM = 2048
+        scale_factor = 1.0
+        working_np = original_np
+        working_bgr = original_bgr
+        working_h, working_w = orig_h, orig_w
+
+        if max(orig_w, orig_h) > MAX_PROCESS_DIM:
+            scale_factor = MAX_PROCESS_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
+        )
+        parsing = upsampled_logits.argmax(dim=1).squeeze(0).cpu().numpy()
+
+        # Skin mask
+        skin_mask = (parsing == skin_class_id).astype(np.uint8)
+
+        # IMPROVED Forehead region (better pixel coverage)
+        forehead_fraction_top    = 0.25   # thoda neeche
+        forehead_fraction_bottom = 0.38   # zyada coverage
+        forehead_fraction_left   = 0.38
+        forehead_fraction_right  = 0.62   # wider center
+
+        h, w = parsing.shape
+        forehead_y_start = max(0, int(h * forehead_fraction_top))
+        forehead_y_end   = min(h, int(h * forehead_fraction_bottom))
+        forehead_x_start = max(0, int(w * forehead_fraction_left))
+        forehead_x_end   = min(w, int(w * forehead_fraction_right))
+
+        forehead_region = original_np[forehead_y_start:forehead_y_end, forehead_x_start:forehead_x_end]
+        forehead_skin_mask = skin_mask[forehead_y_start:forehead_y_end, forehead_x_start:forehead_x_end]
+
+        mean_color_rgb = np.array([210, 185, 170])  # Lighter neutral fallback
+
+        try:
+            if forehead_region.size > 0 and np.sum(forehead_skin_mask) > 80:
+                skin_pixels = forehead_region[forehead_skin_mask == 1]
+                if len(skin_pixels) > 30:
+                    brightness = np.mean(skin_pixels.astype(float), axis=1)
+                    thresh = np.percentile(brightness, 70)
+                    bright_pixels = skin_pixels[brightness > thresh]
+                    if len(bright_pixels) > 20:
+                        mean_color_rgb = np.mean(bright_pixels, axis=0).astype(int)
+                    else:
+                        mean_color_rgb = np.mean(skin_pixels, axis=0).astype(int)
+                else:
+                    mean_color_rgb = np.mean(forehead_region, axis=(0,1)).astype(int)
+            else:
+                # Fallback 1: Nose
+                nose_mask = (parsing == nose_class_id).astype(np.uint8)
+                nose_pixels = original_np[nose_mask == 1]
+                if len(nose_pixels) > 50:
+                    mean_color_rgb = np.mean(nose_pixels, axis=0).astype(int)
+                else:
+                    # Fallback 2: Full skin
+                    skin_pixels_full = original_np[skin_mask == 1]
+                    if len(skin_pixels_full) > 100:
+                        mean_color_rgb = np.mean(skin_pixels_full, axis=0).astype(int)
+        except Exception as skin_err:
+            print(f"Skin detection error (fallback used): {str(skin_err)}")
+
+        # Make detected skin color 30% brighter
+        mean_color_rgb = np.array(mean_color_rgb, dtype=float)
+        brightness_factor = 1.30   # 30% brighter (change to 1.20 / 1.40 if needed)
+        mean_color_rgb = np.clip(mean_color_rgb * brightness_factor, 0, 255).astype(int)
+
+        # Clear forehead color print (updated one)
+        hex_color = '#%02x%02x%02x' % tuple(mean_color_rgb)
+        print(f"Adjusted (30% brighter) skin color → RGB: {mean_color_rgb.tolist()} | Hex: {hex_color}")
+
+        hair_mask = (parsing == hair_class_id).astype(np.uint8)
+
+        ears_mask = np.zeros_like(hair_mask, dtype=np.uint8)
+        for cls in ear_class_ids:
+            ears_mask[parsing == cls] = 1
+
+        ears_protected = np.zeros_like(hair_mask, dtype=np.uint8)
+        ear_y, ear_x = np.where(ears_mask > 0)
+
+        left, right = 0, 0
+        if len(ear_y) > 0:
+            ear_top_y = ear_y.min()
+            ear_x_min = ear_x.min()
+            ear_x_max = ear_x.max()
+            ear_width = ear_x_max - ear_x_min + 1
+
+            kernel_protect = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 9))
+            ears_protected = cv2.dilate(ears_mask, kernel_protect, iterations=1)
+
+            if ear_top_y > 10:
+                ears_protected[:ear_top_y - 8, :] = 0
+
+            x_margin = int(ear_width * 0.25)
+            left = max(0, ear_x_min - x_margin)
+            right = min(working_w, ear_x_max + x_margin)
+
+        hair_mask_final = hair_mask.copy()
+        hair_mask_final[ears_protected == 1] = 0
+
+        top_quarter = int(working_h * 0.25)
+        if hair_mask[:top_quarter, :].sum() > 60:
+            hair_mask_final[:top_quarter, :] = np.maximum(
+                hair_mask_final[:top_quarter, :], hair_mask[:top_quarter, :]
+            )
+
+        kernel_s = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (13, 13))
+        hair_mask_final = cv2.morphologyEx(hair_mask_final, cv2.MORPH_CLOSE, kernel_s, iterations=2)
+        hair_mask_final = cv2.dilate(hair_mask_final, kernel_s, iterations=1)
+
+        blurred = cv2.GaussianBlur(hair_mask_final.astype(np.float32), (9, 9), 3)
+        hair_mask_final = (blurred > 0.28).astype(np.uint8)
+
+        kernel_edge = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+        hair_mask_final = cv2.dilate(hair_mask_final, kernel_edge, iterations=1)
+
+        hair_pixels = np.sum(hair_mask_final)
+
+        final_mask = hair_mask_final.copy()
+        use_extended_mask = False
+        if hair_pixels > 380000:
+            use_extended_mask = True
+            big_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (25, 25))
+            extended = cv2.dilate(hair_mask_final, big_kernel, iterations=1)
+            upper = np.zeros_like(hair_mask_final)
+            upper_end = int(working_h * 0.48)
+            upper[:upper_end, :] = 1
+            extended = np.logical_or(extended, upper).astype(np.uint8)
+            extended[ears_protected == 1] = 0
+            extended = cv2.morphologyEx(extended, cv2.MORPH_CLOSE, kernel_s, iterations=1)
+            extended[int(working_h * 0.75):, :] = 0
+            final_mask = extended
+
+        if use_extended_mask or hair_pixels > 420000:
+            radius = 18
+            inpaint_flag = cv2.INPAINT_TELEA
+        elif hair_pixels > 220000:
+            radius = 15
+            inpaint_flag = cv2.INPAINT_TELEA
+        else:
+            radius = 10
+            inpaint_flag = cv2.INPAINT_NS
+
+        inpainted_bgr = cv2.inpaint(working_bgr, final_mask * 255, inpaintRadius=radius, flags=inpaint_flag)
+        inpainted_rgb = cv2.cvtColor(inpainted_bgr, cv2.COLOR_BGR2RGB)
+
+        # ==================== NEW: Add realistic bald head skin texture ====================
+        pores_noise = np.random.normal(0, 12, (working_h, working_w, 3)).astype(np.float32)
+        large_kernel = cv2.getGaussianKernel(61, 20)
+        large_var = cv2.filter2D(pores_noise, -1, large_kernel) * 0.5
+        texture_noise = pores_noise * 0.7 + large_var
+        texture_noise = np.clip(texture_noise, -25, 25)
+
+        textured_area = inpainted_rgb.astype(np.float32) + texture_noise
+        textured_area = np.clip(textured_area, 0, 255).astype(np.uint8)
+
+        blend_factor = 0.75  # 75% textured, 25% smooth inpaint
+        blended_bald = (blend_factor * textured_area + (1 - blend_factor) * inpainted_rgb).astype(np.uint8)
+        # =================================================================================
+
+        result_small = working_np.copy()
+        result_small[final_mask == 1] = blended_bald[final_mask == 1]
+
+        if len(ear_x) > 0:
+            side_clean_left = max(0, left - 30)
+            side_clean_right = min(working_w, right + 30)
+            final_mask[:, side_clean_left:side_clean_right] = np.minimum(
+                final_mask[:, side_clean_left:side_clean_right],
+                1 - ears_protected[:, side_clean_left:side_clean_right]
+            )
+
+        if scale_factor < 1.0:
+            result = cv2.resize(result_small, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4)
+        else:
+            result = result_small
+
+        output_bytes = io.BytesIO()
+        Image.fromarray(result).save(output_bytes, format="JPEG")
+        output_bytes.seek(0)
+        return output_bytes.read()
+
+    except Exception as main_err:
+        print("ERROR in make_realistic_bald:")
+        traceback.print_exc()
+        raise RuntimeError(f"Bald processing failed: {str(main_err)}")