diff --git a/lib/remove_bg.py b/lib/remove_bg.py index 9dddc5f..0e815c7 100644 --- a/lib/remove_bg.py +++ b/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 - batches = [files_to_process[i:i + batch_size] for i in range(0, len(files_to_process), batch_size)] + # Divide filenames into batches + 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 \ No newline at end of file