Improve character cutout prompts and offline candidate evaluation

segment.py

@@ -9,6 +9,7 @@ import argparse
 import json
 import os
 import sys
+import numpy as np
 
 try:
     import cv2
@@ -25,8 +26,17 @@ def segment(atlas_path: str, output_dir: str, min_area: int, padding: int, bg_th
     elif image.shape[2] == 3:
         image = cv2.cvtColor(image, cv2.COLOR_BGR2BGRA)
 
-    gray = cv2.cvtColor(image[:, :, :3], cv2.COLOR_BGR2GRAY)
-    _, mask = cv2.threshold(gray, bg_threshold, 255, cv2.THRESH_BINARY_INV)
+    # Removing every bright pixel also erased eye whites, teeth and pale
+    # clothing. Treat white connected to the canvas boundary as background;
+    # enclosed highlights remain part of their surrounding cutout.
+    white = np.all(image[:, :, :3] >= bg_threshold, axis=2)
+    transparent = image[:, :, 3] == 0
+    bg_candidates = (white | transparent).astype(np.uint8)
+    _, bg_labels, _, _ = cv2.connectedComponentsWithStats(bg_candidates, 8)
+    border_labels = np.unique(np.concatenate((bg_labels[0], bg_labels[-1], bg_labels[:, 0], bg_labels[:, -1])))
+    border_labels = border_labels[border_labels != 0]
+    background = np.isin(bg_labels, border_labels) | transparent
+    mask = (~background).astype(np.uint8) * 255
     count, labels, stats, _ = cv2.connectedComponentsWithStats(mask, 8)
     os.makedirs(output_dir, exist_ok=True)
     height, width = image.shape[:2]