Add --mask-threshold, cap workers at 4, print device info, bump to 0.4.0

- Add -mt/--mask-threshold for hard/sharp edges on cartoonish images (fixes #122)
- Cap worker count at 4 to prevent hangs and GPU memory exhaustion (fixes #181)
- Print detected device (CUDA/MPS/CPU) at startup for debugging (helps #145)
- Expose mask_threshold via HTTP API server (?mt= parameter)

Author: nadermx

backgroundremover/__init__.py

@@ -4,6 +4,6 @@
 A library to remove background from videos and images
 """
 
-__version__ = "0.3.9"
+__version__ = "0.4.0"
 __author__ = 'Johnathan Nader'
 __credits__ = 'BackgroundRemoverAI.com'

backgroundremover/bg.py

@@ -27,12 +27,15 @@
 try:
     if torch.cuda.is_available():
         DEVICE = torch.device('cuda:0')
+        print(f"Device: CUDA ({torch.cuda.get_device_name(0)})")
     elif torch.backends.mps.is_available():
         DEVICE = torch.device('mps')
+        print("Device: MPS (Apple Silicon GPU)")
     else:
         DEVICE = torch.device('cpu')
+        print("Device: CPU (no GPU detected - install CUDA toolkit for GPU acceleration)")
 except Exception as e:
-    print(f"Using CPU.  Setting Cuda or MPS failed: {e}")
+    print(f"Device: CPU (Setting CUDA or MPS failed: {e})")
     DEVICE = torch.device('cpu')
 
 class Net(torch.nn.Module):
@@ -216,6 +219,7 @@ def remove(
     only_mask=False,
     background_color=None,
     background_image=None,
+    mask_threshold=None,
 ):
     model = get_model(model_name)
 
@@ -232,6 +236,10 @@ def remove(
 
     mask = detect.predict(model, np.array(img)).convert("L")
 
+    # Apply threshold for hard/sharp edges (fixes #122)
+    if mask_threshold is not None:
+        mask = mask.point(lambda p: 255 if p > mask_threshold else 0)
+
     # If only_mask is True, return just the mask
     if only_mask:
         bio = io.BytesIO()

backgroundremover/cmd/cli.py

@@ -71,6 +71,14 @@ def main():
         help="Output only the binary mask (grayscale image).",
     )
 
+    ap.add_argument(
+        "-mt",
+        "--mask-threshold",
+        default=None,
+        type=int,
+        help="Threshold (0-255) to binarize the mask for hard/sharp edges. Useful for cartoonish images. Values around 128 work well.",
+    )
+
     ap.add_argument(
         "-bc",
         "--background-color",
@@ -240,10 +248,11 @@ def main():
         print("Example: backgroundremover -i video.mp4 -tgwb -bi background.png -o output.gif")
         exit(1)
 
-    # Warn about high worker counts that may cause issues
+    # Cap worker count to prevent hangs and resource exhaustion (see issue #181)
     if args.workernodes > 4:
-        print(f"Warning: Using {args.workernodes} workers. High worker counts (>4) may cause ConnectionResetError or crashes on some systems.")
-        print("If you experience errors, try reducing workers with -wn 1 or -wn 2")
+        print(f"Warning: Requested {args.workernodes} workers, capping at 4. Higher values cause hangs and GPU memory exhaustion.")
+        print("Use -wn 1 through -wn 4 for best results.")
+        args.workernodes = 4
 
     # Parse background color if provided
     background_color = None
@@ -348,6 +357,7 @@ def is_image_file(filename):
                             only_mask=args.only_mask,
                             background_color=background_color,
                             background_image=background_image,
+                            mask_threshold=args.mask_threshold,
                         ),
                     )
         return

backgroundremover/cmd/server.py

@@ -38,6 +38,7 @@ def index():
     ab = request.values.get("ab", type=int, default=10)
     ae = request.values.get("ae", type=int, default=10)
     az = request.values.get("az", type=int, default=1000)
+    mt = request.values.get("mt", type=int, default=None)
 
     model = request.args.get("model", type=str, default="u2net")
     model_path = os.environ.get(
@@ -62,6 +63,7 @@ def index():
                     alpha_matting_background_threshold=ab,
                     alpha_matting_erode_structure_size=ae,
                     alpha_matting_base_size=az,
+                    mask_threshold=mt,
                 )
             ),
             mimetype="image/png",

setup.py

@@ -11,7 +11,7 @@
 
 setup(
     name="backgroundremover",
-    version="0.3.9",
+    version="0.4.0",
     description="Background remover from image and video using AI",
     long_description=long_description,
     long_description_content_type="text/markdown",