fix: improve pink detection with multiple methods and add debug logging to diagnose removal issues

api/main.py

@@ -394,27 +394,38 @@ def remove_pink_segments(
     # Pink/Magenta → white in mask (remove)
     # Everything else (natural image colors, including dark areas) → black in mask (keep)
     
-    # Detect pink/magenta paint (RGB: R>240, G<30, B>240)
-    magenta_detected = (
-        (img_rgb[:, :, 0] > 240) &  # Red: high
-        (img_rgb[:, :, 1] < 30) &   # Green: low
-        (img_rgb[:, :, 2] > 240)    # Blue: high
+    # More robust pink/magenta detection - try multiple methods
+    # Method 1: Detect bright magenta (RGB: R>200, G<80, B>200) - more lenient
+    magenta_detected_loose = (
+        (img_rgb[:, :, 0] > 200) &  # Red: high
+        (img_rgb[:, :, 1] < 80) &   # Green: low (allow more tolerance)
+        (img_rgb[:, :, 2] > 200)    # Blue: high
     ).astype(np.uint8) * 255
     
+    # Method 2: Detect exact magenta (255, 0, 255)
+    magenta_strict = np.all(img_rgb == [255, 0, 255], axis=2).astype(np.uint8) * 255
+    
+    # Method 3: Detect pink-ish colors (high red+blue, low green)
+    pink_detected = (
+        (img_rgb[:, :, 0] > 220) &  # Red: very high
+        (img_rgb[:, :, 1] < 100) &  # Green: low
+        (img_rgb[:, :, 2] > 220)    # Blue: very high
+    ).astype(np.uint8) * 255
+    
+    # Combine all detection methods
+    binmask = np.maximum.reduce([magenta_detected_loose, magenta_strict, pink_detected])
+    
     # Clean up the pink mask
     kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
-    binmask = cv2.morphologyEx(magenta_detected, cv2.MORPH_CLOSE, kernel, iterations=2)
+    binmask = cv2.morphologyEx(binmask, cv2.MORPH_CLOSE, kernel, iterations=2)
     binmask = cv2.morphologyEx(binmask, cv2.MORPH_OPEN, kernel, iterations=1)
     
     nonzero = int((binmask > 0).sum())
-    log.info(f"Detected {nonzero} pink pixels to remove")
+    total_pixels = binmask.shape[0] * binmask.shape[1]
+    log.info(f"Detected {nonzero} pink pixels ({100*nonzero/total_pixels:.2f}% of image) to remove")
     
-    if nonzero < 50:
-        log.warning("Very few pink pixels detected! Trying strict magenta detection...")
-        magenta_strict = np.all(img_rgb == [255, 0, 255], axis=2).astype(np.uint8) * 255
-        binmask = cv2.morphologyEx(magenta_strict, cv2.MORPH_CLOSE, kernel, iterations=3)
-        nonzero = int((binmask > 0).sum())
-        log.info(f"Strict detection: {nonzero} pink pixels")
+    # Debug: show detection stats
+    log.info(f"Detection stats - Loose: {int((magenta_detected_loose > 0).sum())}, Strict: {int((magenta_strict > 0).sum())}, Pink: {int((pink_detected > 0).sum())}")
     
     if nonzero < 50:
         log.error("No pink segments detected! Returning original image.")
@@ -428,7 +439,7 @@ def remove_pink_segments(
         }
     
     # Create binary mask: Pink pixels → white (255), Everything else → black (0)
-    # But we need to encode it in RGBA format that process_inpaint expects
+    # Encode in RGBA format that process_inpaint expects
     # process_inpaint does: mask = 255 - mask[:,:,3]
     # So: alpha=0 (transparent/pink) → becomes 255 (white/remove)
     #     alpha=255 (opaque/keep) → becomes 0 (black/keep)
@@ -441,11 +452,25 @@ def remove_pink_segments(
     #        255 (opaque) everywhere else → will become black after 255-alpha
     mask_rgba[:, :, 3] = 255 - binmask  # Invert: pink areas get alpha=0, rest get alpha=255
     
+    # Verify mask encoding
+    alpha_zero_count = int((mask_rgba[:,:,3] == 0).sum())
+    alpha_255_count = int((mask_rgba[:,:,3] == 255).sum())
     total_pixels = binmask.shape[0] * binmask.shape[1]
-    log.info(f"Created binary mask: {nonzero} pink pixels → white (255) in final mask, {total_pixels - nonzero} black (0) - everything else automatically")
+    log.info(f"Mask encoding: {alpha_zero_count} pixels with alpha=0 (pink), {alpha_255_count} pixels with alpha=255 (keep)")
+    log.info(f"After 255-alpha conversion: {alpha_zero_count} will become white (255/remove), {alpha_255_count} will become black (0/keep)")
+    
+    # IMPORTANT: We need to use the ORIGINAL image WITHOUT pink paint for inpainting!
+    # Remove pink from the original image before processing
+    # Create a clean version: where pink was detected, keep original image colors
+    img_clean = np.array(img.convert("RGBA"))
+    # Where pink is detected, we want to inpaint, so we can leave it (or blend it out)
+    # Actually, the model will inpaint over those areas, so we can pass the original
+    # But for better results, we might want to remove the pink overlay first
     
     # Process with invert_mask=True (default) because process_inpaint expects alpha=0 for removal
-    result = process_inpaint(np.array(img), mask_rgba, invert_mask=True)
+    log.info(f"Starting inpainting process...")
+    result = process_inpaint(img_clean, mask_rgba, invert_mask=True)
+    log.info(f"Inpainting complete, result shape: {result.shape}")
     result_name = f"output_{uuid.uuid4().hex}.png"
     result_path = os.path.join(OUTPUT_DIR, result_name)
     Image.fromarray(result).save(result_path, "PNG", optimize=False, compress_level=1)