gpu support by flag

lib/remove_bg.py

@@ -6,16 +6,13 @@ from concurrent.futures import ThreadPoolExecutor
 
 import onnxruntime as ort
 
-# Suppress ONNX Runtime logging to show only critical errors
-ort.set_default_logger_severity(3)  # 0 = verbose, 1 = info, 2 = warning, 3 = error, 4 = fatal
-
-
-def process_images_in_batch(batch):
+def process_images_in_batch(batch, use_gpu):
     """
     Process a batch of images: remove their backgrounds and save the results.
 
     Args:
     batch (list): List of tuples (input_path, output_path).
+    use_gpu (bool): Whether to enable GPU support.
 
     Returns:
     int: Number of images successfully processed in this batch.
@@ -24,14 +21,20 @@ def process_images_in_batch(batch):
     for input_path, output_path in batch:
         try:
             with Image.open(input_path) as img:
-                output = remove(img)  # This will use GPU if ONNX Runtime is GPU-enabled
+                # Initialize ONNX session options with GPU support if required
+                session_options = ort.SessionOptions()
+                providers = ["CUDAExecutionProvider"] if use_gpu else ["CPUExecutionProvider"]
+                ort.set_default_logger_severity(3)  # Suppress non-critical logging
+
+                # Initialize the rembg remove function with appropriate providers
+                output = remove(img, session_options=session_options, providers=providers)
                 output.save(output_path)
             success_count += 1
         except Exception as e:
             print(f"Error processing {input_path}: {str(e)}")
     return success_count
 
-def remove_bg_from_files_in_dir(directory, max_workers=2, batch_size=5):
+def remove_bg_from_files_in_dir(directory, max_workers=2, batch_size=3, use_gpu=False):
     """
     Process all JPG, JPEG, and PNG images in the given directory and its subfolders using parallel processing and GPU.
 
@@ -39,6 +42,7 @@ def remove_bg_from_files_in_dir(directory, max_workers=2, batch_size=5):
     directory (str): Path to the directory containing images.
     max_workers (int): Maximum number of threads to use for parallel processing.
     batch_size (int): Number of images to process per batch.
+    use_gpu (bool): Whether to enable GPU support.
 
     Returns:
     int: The number of images successfully processed.
@@ -60,8 +64,11 @@ def remove_bg_from_files_in_dir(directory, max_workers=2, batch_size=5):
     batches = [files_to_process[i:i + batch_size] for i in range(0, len(files_to_process), batch_size)]
 
     with ThreadPoolExecutor(max_workers=max_workers) as executor:
-        with tqdm(total=len(files_to_process), desc="Processing images", unit="image") as pbar:
-            futures = {executor.submit(process_images_in_batch, batch): batch for batch in batches}
+        with tqdm(total=len(files_to_process), desc="Removing Backgrounds", unit="image") as pbar:
+            futures = {
+                executor.submit(process_images_in_batch, batch, use_gpu): batch
+                for batch in batches
+            }
 
             for future in futures:
                 processed_count += future.result()