comfy_fm24_newgens/lib/remove_bg.py

try:
    from rembg import remove
    REMBG_AVAILABLE = True
except ImportError as e:
    print(f"Warning: rembg not available: {e}")
    REMBG_AVAILABLE = False

try:
    import onnxruntime as ort
    ONNX_AVAILABLE = True
except ImportError:
    ONNX_AVAILABLE = False

from PIL import Image
from tqdm import tqdm
import os
from concurrent.futures import ThreadPoolExecutor

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 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.
    """
    if not REMBG_AVAILABLE:
        print("Background removal not available (rembg not installed)")
        return 0

    success_count = 0
    for filename in batch:
        input_path = os.path.join(directory, f"{filename}.png")
        output_path = os.path.join(directory, f"{filename}.png")

        try:
            # Check if input file exists
            if not os.path.exists(input_path):
                print(f"Warning: Input file not found: {input_path}")
                continue

            # Check if input file is valid
            if os.path.getsize(input_path) == 0:
                print(f"Warning: Empty file: {input_path}")
                continue

            with Image.open(input_path) as img:
                # Validate image
                if img.size[0] == 0 or img.size[1] == 0:
                    print(f"Warning: Invalid image dimensions: {input_path}")
                    continue

                # Initialize ONNX session options with GPU support if required
                if ONNX_AVAILABLE:
                    try:
                        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)
                    except Exception as gpu_error:
                        print(f"GPU processing failed for {input_path}, falling back to CPU: {str(gpu_error)}")
                        # Fallback to CPU-only processing
                        output = remove(img)
                else:
                    # CPU-only processing
                    output = remove(img)

                # Validate output
                if output is None:
                    print(f"Warning: Background removal returned None for {input_path}")
                    continue

                # Save result
                output.save(output_path, 'PNG')
                success_count += 1

        except Image.UnidentifiedImageError:
            print(f"Error: Invalid image format: {input_path}")
        except PermissionError:
            print(f"Error: Permission denied accessing: {input_path}")
        except OSError as e:
            print(f"Error: OS error processing {input_path}: {str(e)}")
        except Exception as e:
            print(f"Error processing {input_path}: {str(e)}")

    return success_count

def remove_bg_from_file_list(directory, filenames, max_workers=2, batch_size=2, use_gpu=False):
    """
    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.

    Returns:
    int: The number of images successfully processed.
    """
    if not REMBG_AVAILABLE:
        print("Background removal not available (rembg not installed). Skipping background removal.")
        return 0

    print(f"Starting background removal for {len(filenames)} images in {directory}")
    print(f"Configuration: max_workers={max_workers}, batch_size={batch_size}, use_gpu={use_gpu}")

    processed_count = 0
    failed_count = 0

    # Divide filenames into batches
    batches = [filenames[i:i + batch_size] for i in range(0, len(filenames), batch_size)]
    total_batches = len(batches)

    print(f"Processing in {total_batches} batches...")

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        with tqdm(total=len(filenames), desc="Removing Backgrounds", unit="image") as pbar:
            futures = {
                executor.submit(process_images_in_batch, batch, directory, use_gpu): batch
                for batch in batches
            }

            for i, future in enumerate(futures, 1):
                try:
                    batch_result = future.result()
                    processed_count += batch_result
                    pbar.update(len(futures[future]))

                    if batch_result > 0:
                        print(f"Batch {i}/{total_batches}: {batch_result} images processed successfully")
                    else:
                        print(f"Batch {i}/{total_batches}: No images processed (possible errors)")

                except Exception as e:
                    failed_count += 1
                    print(f"Batch {i}/{total_batches}: Failed with error: {str(e)}")

    print(f"Background removal completed:")
    print(f"  - Total images: {len(filenames)}")
    print(f"  - Successfully processed: {processed_count}")
    print(f"  - Failed batches: {failed_count}")

    if processed_count == 0:
        print("WARNING: No images were processed successfully!")
        print("This might indicate:")
        print("  - All images have already been processed")
        print("  - Image files are corrupted or missing")
        print("  - Background removal is failing due to image format issues")
        print("  - GPU/CPU compatibility issues with rembg")

    return processed_count