From 97614a54cd48cfcb60e5fdfcffeb2cbd73112fbd Mon Sep 17 00:00:00 2001
From: limbicnation <gero@limbicnation.com>
Date: Fri, 23 May 2025 03:45:13 +0200
Subject: [PATCH] Fix alpha mask greytones by adding binary thresholding

---
 background_remover.py | 37 ++++++++++++++++++++++++++++++++++---
 nodes.py              | 16 +++++++++++++---
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/background_remover.py b/background_remover.py
index d0b7d9c..578e1dc 100644
--- a/background_remover.py
+++ b/background_remover.py
@@ -13,7 +13,8 @@ class EnhancedPixelArtProcessor:
     
     def __init__(self, tolerance: int = 30, edge_sensitivity: float = 0.8,
                  color_clusters: int = 8, foreground_bias: float = 0.7,
-                 edge_refinement: bool = True, dither_handling: bool = True):
+                 edge_refinement: bool = True, dither_handling: bool = True,
+                 binary_threshold: int = 128):
         """
         Initialize the processor with configuration parameters.
         
@@ -24,6 +25,7 @@ class EnhancedPixelArtProcessor:
             foreground_bias: Bias towards foreground preservation (0.0-1.0)
             edge_refinement: Enable edge refinement post-processing
             dither_handling: Enable dithered pattern detection
+            binary_threshold: Threshold for binary alpha mask (0-255)
         """
         self.tolerance = tolerance
         self.edge_sensitivity = edge_sensitivity
@@ -31,6 +33,7 @@ class EnhancedPixelArtProcessor:
         self.foreground_bias = foreground_bias
         self.edge_refinement = edge_refinement
         self.dither_handling = dither_handling
+        self.binary_threshold = binary_threshold
         
     def remove_background_advanced(self, image: np.ndarray) -> np.ndarray:
         """
@@ -81,6 +84,9 @@ class EnhancedPixelArtProcessor:
         # Apply foreground bias
         final_mask = self._apply_foreground_bias(final_mask, image)
         
+        # Convert to binary mask to eliminate semi-transparency
+        final_mask = self._make_binary_mask(final_mask, self.binary_threshold)
+        
         rgba_result[:, :, 3] = final_mask
         return rgba_result
     
@@ -226,9 +232,9 @@ class EnhancedPixelArtProcessor:
         # Fill small holes
         mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_small, iterations=2)
         
-        # Smooth edges
+        # Smooth edges with closing operation
         kernel_smooth = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
-        mask = cv2.morphologyEx(mask, cv2.MORPH_SMOOTH, kernel_smooth)
+        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_smooth, iterations=1)
         
         # Apply Gaussian blur for softer edges
         mask = cv2.GaussianBlur(mask, (3, 3), 0)
@@ -266,6 +272,31 @@ class EnhancedPixelArtProcessor:
         
         return (mask_float * 255).astype(np.uint8)
     
+    def _make_binary_mask(self, mask: np.ndarray, threshold: int = 128) -> np.ndarray:
+        """
+        Convert grayscale mask to binary mask to eliminate semi-transparency.
+        
+        Args:
+            mask: Input grayscale mask (0-255)
+            threshold: Threshold value for binarization (default: 128)
+            
+        Returns:
+            Binary mask with only 0 (transparent) or 255 (opaque) values
+        """
+        # Apply binary threshold
+        _, binary_mask = cv2.threshold(mask, threshold, 255, cv2.THRESH_BINARY)
+        
+        # Optional: Apply morphological operations to clean up the binary mask
+        kernel = np.ones((3, 3), np.uint8)
+        
+        # Remove small noise (opening)
+        binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_OPEN, kernel, iterations=1)
+        
+        # Fill small holes (closing)
+        binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, kernel, iterations=2)
+        
+        return binary_mask
+    
     def auto_adjust_parameters(self, image: np.ndarray) -> dict:
         """
         Automatically adjust parameters based on image analysis.
diff --git a/nodes.py b/nodes.py
index d6263e8..7e68177 100644
--- a/nodes.py
+++ b/nodes.py
@@ -60,6 +60,14 @@ class TransparencyBackgroundRemover:
                     "default": True,
                     "tooltip": "Enable dithered pattern detection and handling"
                 }),
+                "binary_threshold": ("INT", {
+                    "default": 128,
+                    "min": 0,
+                    "max": 255,
+                    "step": 1,
+                    "display": "number",
+                    "tooltip": "Threshold for binary alpha mask (0-255)"
+                }),
                 "output_format": (["RGBA", "RGB_WITH_MASK"],),
             }
         }
@@ -75,7 +83,7 @@ class TransparencyBackgroundRemover:
     def remove_background(self, image, tolerance=30, edge_sensitivity=0.8,
                          foreground_bias=0.7, color_clusters=8,
                          edge_refinement=True, dither_handling=True,
-                         output_format="RGBA"):
+                         binary_threshold=128, output_format="RGBA"):
         """
         Main processing function for background removal with error handling.
         """
@@ -95,6 +103,7 @@ class TransparencyBackgroundRemover:
                                                 color_clusters=color_clusters,
                                                 edge_refinement=edge_refinement,
                                                 dither_handling=dither_handling,
+                                                binary_threshold=binary_threshold,
                                                 output_format=output_format)
 
             return (results, masks)
@@ -109,7 +118,7 @@ class TransparencyBackgroundRemover:
     def _process_images(self, image, tolerance=30, edge_sensitivity=0.8,
                        foreground_bias=0.7, color_clusters=8,
                        edge_refinement=True, dither_handling=True,
-                       output_format="RGBA"):
+                       binary_threshold=128, output_format="RGBA"):
         """
         Internal method for processing images without error handling wrapper.
         """
@@ -130,7 +139,8 @@ class TransparencyBackgroundRemover:
                 color_clusters=color_clusters,
                 foreground_bias=foreground_bias,
                 edge_refinement=edge_refinement,
-                dither_handling=dither_handling
+                dither_handling=dither_handling,
+                binary_threshold=binary_threshold
             )
 
             # Process image