feat(api): improve pink/magenta auto-detection with tolerance - automatically segment pink as white (remove), rest as black (keep)

api/main.py

@@ -264,73 +264,71 @@ def inpaint_multipart(
         return resp
 
     if mask_is_painted:
-        # Derive mask from painted image - detect magenta paint (255, 0, 255) or any painted areas
-        img_rgb = cv2.cvtColor(np.array(img), cv2.COLOR_RGBA2RGB)
-        m_rgb = cv2.cvtColor(np.array(m), cv2.COLOR_RGBA2RGB)
+        # Auto-detect pink/magenta paint and convert to black/white mask
+        # White pixels = areas to remove, Black pixels = areas to keep
+        log.info("Auto-detecting pink/magenta paint from uploaded image...")
         
-        # Method 1: Detect magenta paint specifically (RGB: 255, 0, 255) - Streamlit canvas default
-        magenta_mask = np.all(m_rgb == [255, 0, 255], axis=2).astype(np.uint8) * 255
+        m_rgb = cv2.cvtColor(np.array(m), cv2.COLOR_RGBA2RGB)
         
-        # Method 2: Detect any difference between original and painted
-        diff = cv2.absdiff(img_rgb, m_rgb)
-        gray_diff = cv2.cvtColor(diff, cv2.COLOR_RGB2GRAY)
-        try:
-            _, binmask_diff = cv2.threshold(gray_diff, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-        except Exception:
-            _, binmask_diff = cv2.threshold(gray_diff, 30, 255, cv2.THRESH_BINARY)
+        # Method 1: Detect magenta/pink paint directly (RGB: 255, 0, 255)
+        # Allow some tolerance for slight variations (e.g., 250-255 for R/B, 0-10 for G)
+        magenta_detected = (
+            (m_rgb[:, :, 0] > 240) &  # Red channel: high (240-255)
+            (m_rgb[:, :, 1] < 30) &   # Green channel: low (0-30) 
+            (m_rgb[:, :, 2] > 240)    # Blue channel: high (240-255)
+        ).astype(np.uint8) * 255
         
-        # Combine both methods (magenta OR significant difference)
-        binmask = cv2.bitwise_or(magenta_mask, binmask_diff)
+        # Method 2: Also check if original image was provided to find differences
+        if img is not None:
+            img_rgb = cv2.cvtColor(np.array(img), cv2.COLOR_RGBA2RGB)
+            if img_rgb.shape == m_rgb.shape:
+                diff = cv2.absdiff(img_rgb, m_rgb)
+                gray_diff = cv2.cvtColor(diff, cv2.COLOR_RGB2GRAY)
+                # Any significant difference (>50) could be paint
+                diff_mask = (gray_diff > 50).astype(np.uint8) * 255
+                # Combine with magenta detection
+                binmask = cv2.bitwise_or(magenta_detected, diff_mask)
+            else:
+                binmask = magenta_detected
+        else:
+            # No original image provided, use magenta detection only
+            binmask = magenta_detected
         
-        # Optional: detect other bright colors that might be paint
-        # Detect any bright non-original colors (R+G+B > threshold and differs from original)
-        brightness = m_rgb.sum(axis=2)
-        diff_magnitude = diff.sum(axis=2)
-        bright_paint = ((brightness > 400) & (diff_magnitude > 50)).astype(np.uint8) * 255
-        binmask = cv2.bitwise_or(binmask, bright_paint)
+        # Clean up the mask: remove noise and fill small holes
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+        # Close small gaps in the mask
+        binmask = cv2.morphologyEx(binmask, cv2.MORPH_CLOSE, kernel, iterations=2)
+        # Remove small noise
+        binmask = cv2.morphologyEx(binmask, cv2.MORPH_OPEN, kernel, iterations=1)
         
         nonzero = int((binmask > 0).sum())
-        log.info("painted-mask detection: magenta=%d, diff=%d, combined=%d pixels", 
-                 int((magenta_mask > 0).sum()), int((binmask_diff > 0).sum()), nonzero)
+        log.info("Pink/magenta paint detected: %d pixels marked for removal (white)", nonzero)
         
         if nonzero < 50:
-            log.warning("Very few pixels detected! Trying alternative detection...")
-            # Fallback: detect any significant color change
-            color_change = np.sum(np.abs(img_rgb.astype(np.int16) - m_rgb.astype(np.int16)), axis=2)
-            _, binmask = cv2.threshold(color_change.astype(np.uint8), 30, 255, cv2.THRESH_BINARY)
+            log.warning("Very few pixels detected! Trying stricter magenta detection...")
+            # Try more strict magenta detection (exact match)
+            magenta_strict = np.all(m_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("fallback detection: %d pixels", nonzero)
-        
-        # Build RGBA mask: convert to proper black/white mask
-        # White (255) = remove, Black (0) = keep (standard convention)
-        mask_rgba = np.zeros((binmask.shape[0], binmask.shape[1], 4), dtype=np.uint8)
-        
-        # Set RGB channels: white where paint detected, black elsewhere
-        mask_rgba[:, :, 0] = binmask  # R
-        mask_rgba[:, :, 1] = binmask  # G  
-        mask_rgba[:, :, 2] = binmask  # B
-        # Set alpha to opaque so it's treated as a standard RGB mask
-        mask_rgba[:, :, 3] = 255
+            log.info("Strict magenta detection: %d pixels", nonzero)
         
-        # Also create a cleaner version: apply morphological operations to smooth edges
-        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
-        binmask_clean = cv2.morphologyEx(binmask, cv2.MORPH_CLOSE, kernel)
-        binmask_clean = cv2.morphologyEx(binmask_clean, cv2.MORPH_OPEN, kernel)
-        mask_rgba[:, :, 0] = binmask_clean
-        mask_rgba[:, :, 1] = binmask_clean
-        mask_rgba[:, :, 2] = binmask_clean
-        
-        log.info("Auto-converted painted image to black/white mask: %d white pixels (to remove)", 
-                 int((binmask_clean > 0).sum()))
-        
-        if int((binmask_clean > 0).sum()) < 50:
-            log.error("CRITICAL: Mask detection found very few pixels! Returning original image.")
-            # Return original image if mask is invalid
-            result = np.array(img.convert("RGB"))
+        if nonzero < 50:
+            log.error("CRITICAL: Could not detect pink/magenta paint! Returning original image.")
+            result = np.array(img.convert("RGB")) if img else np.array(m.convert("RGB"))
             result_name = f"output_{uuid.uuid4().hex}.png"
             result_path = os.path.join(OUTPUT_DIR, result_name)
             Image.fromarray(result).save(result_path)
-            return {"result": result_name, "error": "mask detection failed - very few pixels detected"}
+            return {"result": result_name, "error": "pink/magenta paint detection failed - very few pixels detected"}
+        
+        # Create black/white mask: white = remove (pink areas), black = keep (everything else)
+        mask_rgba = np.zeros((binmask.shape[0], binmask.shape[1], 4), dtype=np.uint8)
+        mask_rgba[:, :, 0] = binmask  # R: white where pink detected
+        mask_rgba[:, :, 1] = binmask  # G: white where pink detected
+        mask_rgba[:, :, 2] = binmask  # B: white where pink detected
+        mask_rgba[:, :, 3] = 255      # Alpha: fully opaque
+        
+        log.info("Successfully created black/white mask: %d white pixels (to remove), %d black pixels (to keep)",
+                 nonzero, binmask.shape[0] * binmask.shape[1] - nonzero)
     else:
         mask_rgba = _load_rgba_mask_from_image(m)