reworked to work with processed_players

lib/remove_bg.py

@@ -3,22 +3,25 @@ from PIL import Image
 from tqdm import tqdm
 import os
 from concurrent.futures import ThreadPoolExecutor
-
 import onnxruntime as ort
 
-def process_images_in_batch(batch, use_gpu):
+def process_images_in_batch(batch, directory, use_gpu):
     """
     Process a batch of images: remove their backgrounds and save the results.
 
     Args:
-    batch (list): List of tuples (input_path, output_path).
+    batch (list): List of filenames to process (without path or extension).
+    directory (str): Base directory to locate input files.
     use_gpu (bool): Whether to enable GPU support.
 
     Returns:
     int: Number of images successfully processed in this batch.
     """
     success_count = 0
-    for input_path, output_path in batch:
+    for filename in batch:
+        input_path = os.path.join(directory, f"{filename}.png")
+        output_path = os.path.join(directory, f"{filename}.png")
+
         try:
             with Image.open(input_path) as img:
                 # Initialize ONNX session options with GPU support if required
@@ -29,17 +32,20 @@ def process_images_in_batch(batch, use_gpu):
                 # 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=3, use_gpu=False):
+def remove_bg_from_file_list(directory, filenames, max_workers=2, batch_size=2, use_gpu=False):
     """
-    Process all JPG, JPEG, and PNG images in the given directory and its subfolders using parallel processing and GPU.
+    Process a list of specified filenames: remove their backgrounds and save the results.
 
     Args:
     directory (str): Path to the directory containing images.
+    filenames (list): List of filenames (without path or extension) to process.
     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.
@@ -47,26 +53,15 @@ def remove_bg_from_files_in_dir(directory, max_workers=2, batch_size=3, use_gpu=
     Returns:
     int: The number of images successfully processed.
     """
-    files_to_process = []
-
-    # Gather all the image files to process
-    for subdir, dirs, files in os.walk(directory):
-        for file in files:
-            if file.lower().endswith(('.jpg', '.jpeg', '.png')):
-                input_path = os.path.join(subdir, file)
-                output_filename = os.path.splitext(file)[0] + '.png'
-                output_path = os.path.join(subdir, output_filename)
-                files_to_process.append((input_path, output_path))
-
     processed_count = 0
 
-    # Divide files into batches
+    # Divide filenames into batches
-    batches = [files_to_process[i:i + batch_size] for i in range(0, len(files_to_process), batch_size)]
+    batches = [filenames[i:i + batch_size] for i in range(0, len(filenames), batch_size)]
 
     with ThreadPoolExecutor(max_workers=max_workers) as executor:
-        with tqdm(total=len(files_to_process), desc="Removing Backgrounds", unit="image") as pbar:
+        with tqdm(total=len(filenames), desc="Removing Backgrounds", unit="image") as pbar:
             futures = {
-                executor.submit(process_images_in_batch, batch, use_gpu): batch
+                executor.submit(process_images_in_batch, batch, directory, use_gpu): batch
                 for batch in batches
             }
 
@@ -74,4 +69,4 @@ def remove_bg_from_files_in_dir(directory, max_workers=2, batch_size=3, use_gpu=
                 processed_count += future.result()
                 pbar.update(len(futures[future]))
 
-    return processed_count
+    return processed_count