Efficient Digital Asset Management is a critical cost factor in e-commerce and digital media production. This project implements a high-performance pipeline for automated image processing using the BiRefNet model. The system is designed to standardize and remove backgrounds from large volumes of visual assets without manual intervention.

By leveraging GPU acceleration and state-of-the-art segmentation technology, this solution replaces time-consuming manual editing processes, drastically reduces time-to-market for new products, and lowers operational costs in content creation.

Business Impact & Use Cases

This automation solution addresses key bottlenecks in visual workflows:

• E-Commerce Automation: Instant, standardized preparation of product catalogs through batch processing.
• Cost Reduction: Eliminating repetitive manual tasks (background removal) frees up focus for high-value creative work.
• Data Preparation for AI: Automated cleaning and preparation of image datasets for computer vision training pipelines.
• High Scalability: Robust processing of thousands of assets per hour with consistent quality, regardless of subject complexity.

Technical Approach

The script is based on the BiRefNet model (ZhengPeng7/BiRefNet), a state-of-the-art neural network for image segmentation specifically trained for precise object delineation. The implementation uses PyTorch and Hugging Face Transformers for model inference.

1. GPU Optimization for Modern Hardware

Special attention was paid to supporting the latest NVIDIA GPUs such as the RTX 50 series. To avoid compatibility issues with CUDA kernels, I set special environment variables and PyTorch configurations:

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
os.environ['TORCH_USE_CUDA_DSA'] = '1'
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

These settings allow the model to fully utilize the tensor cores of modern GPUs while avoiding memory fragmentation. The script automatically detects available hardware and displays details such as GPU name, CUDA version, and compute capability.

2. Image Preprocessing and Normalization

Before an image is processed by the model, it goes through a standardized preprocessing pipeline:

transform = transforms.Compose([
    transforms.Resize((1024, 1024)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

Normalization with ImageNet standard values ensures the model delivers optimal results, since it was trained on similarly normalized data. The resize operation to 1024×1024 pixels ensures consistent input sizes for the network.

3. Precise Mask Generation and Post-Processing

The heart of the processing is mask generation. The BiRefNet model produces a pixel-accurate segmentation mask indicating which pixels belong to the foreground:

with torch.no_grad():
    preds = model(input_tensor)[-1].sigmoid().cpu()

pred_mask = preds[0].squeeze()
pred_mask_np = pred_mask.numpy()

The mask is then upscaled to the original image size (using LANCZOS interpolation for best quality) and added as an alpha channel to the image. The result is a transparent PNG with a perfectly extracted object.

4. Robust Error Handling and Fallback Mechanisms

A special feature is the automatic CPU fallback on CUDA errors. If GPU processing of an image fails (e.g., due to memory shortage), the script automatically switches to CPU processing for that single image:

except Exception as e:
    if "CUDA" in str(e) and device.type == 'cuda':
        print(f"\nRetrying {filename} on CPU...")
        torch.cuda.empty_cache()
        model.to('cpu')
        result_image = remove_background(input_path, model, torch.device('cpu'))
        model.to(device)

This mechanism ensures that batch processing does not completely abort due to problematic images, but continues robustly.

Batch Processing with Progress Display

The script uses tqdm for detailed progress display during processing. It iterates over all images in an input folder and automatically saves the extracted versions as PNG files in the output folder:

for filename in tqdm(image_files, desc="Removing backgrounds"):
    try:
        input_path = os.path.join(input_folder, filename)
        result_image = remove_background(input_path, model, device)
        
        output_filename = Path(filename).stem + ".png"
        output_path = os.path.join(output_folder, output_filename)
        result_image.save(output_path, "PNG")
        
        successful += 1
    except Exception as e:
        print(f"\nError processing {filename}: {str(e)}")
        failed += 1

At the end of processing, the script outputs a clear summary showing successful and failed operations.

Technologies Used

• Language: Python 3
• Deep Learning Framework: PyTorch with CUDA support
• Model: BiRefNet (ZhengPeng7/BiRefNet) from Hugging Face
• Image Processing: PIL (Pillow), torchvision transforms
• API Access: Hugging Face Transformers with token authentication
• Progress Display: tqdm
• Configuration: python-dotenv for environment variables

Security Aspects and Best Practices

The script follows modern best practices for machine learning applications:

• Token Security: The Hugging Face API token is loaded via a .env file, never hardcoded.
• Memory Management: Automatic GPU cache clearing (torch.cuda.empty_cache()) on errors.
• Evaluation Mode: The model is set to evaluation mode with model.eval() to disable dropout and batch normalization.
• Transparency Preservation: All outputs are saved as PNG to preserve the alpha channel.

The Complete Script

import os
from pathlib import Path
from PIL import Image
import torch
from transformers import AutoModelForImageSegmentation
from torchvision import transforms
import numpy as np
from tqdm import tqdm
from dotenv import load_dotenv

# Force PyTorch to recompile CUDA kernels for RTX 50-series
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
os.environ['TORCH_COMPILE_DEBUG'] = '0'
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
os.environ['TORCH_USE_CUDA_DSA'] = '1'

# Load environment variables
load_dotenv()

# Configuration
INPUT_FOLDER = r"C:\Users\grisc\Desktop\py\shoes\training\boots"
OUTPUT_FOLDER = r"C:\Users\grisc\Desktop\py\shoes\training\boots_no_bg"
MODEL_NAME = "ZhengPeng7/BiRefNet"
HF_TOKEN = os.getenv("HF_TOKEN")

# Image transformation
image_size = (1024, 1024)

def load_model():
    """Load the BiRefNet model."""
    print(f"Loading model: {MODEL_NAME}")
    
    if not HF_TOKEN:
        raise ValueError("HF_TOKEN not found in .env file. Please add your Hugging Face token.")
    
    # Configure CUDA for RTX 50-series
    if torch.cuda.is_available():
        # Use memory efficient settings
        torch.backends.cudnn.benchmark = False
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
    
    model = AutoModelForImageSegmentation.from_pretrained(
        MODEL_NAME, 
        trust_remote_code=True,
        token=HF_TOKEN
    )
    
    # Use GPU if available
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if torch.cuda.is_available():
        print(f"✓ Using GPU: {torch.cuda.get_device_name(0)} (CUDA {torch.version.cuda})")
        print(f"  Compute capability: sm_{torch.cuda.get_device_capability(0)[0]}{torch.cuda.get_device_capability(0)[1]}")
    else:
        print("⚠ GPU not available, using CPU")
    
    model.to(device)
    model.eval()
    
    print(f"Model loaded on: {device}")
    return model, device

def preprocess_image(image):
    """Preprocess image for the model."""
    # Resize and normalize
    transform = transforms.Compose([
        transforms.Resize(image_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    return transform(image).unsqueeze(0)

def remove_background(image_path, model, device):
    """Remove background from a single image."""
    # Load image
    original_image = Image.open(image_path).convert("RGB")
    original_size = original_image.size
    
    # Preprocess
    input_tensor = preprocess_image(original_image).to(device)
    
    # Generate mask
    with torch.no_grad():
        if device.type == 'cuda':
            # Try running without autocast (float32) to avoid kernel issues
            preds = model(input_tensor)[-1].sigmoid().cpu()
        else:
            preds = model(input_tensor)[-1].sigmoid().cpu()
    
    # Post-process mask
    pred_mask = preds[0].squeeze()
    pred_mask_np = pred_mask.numpy()
    
    # Resize mask to original image size
    mask_pil = Image.fromarray((pred_mask_np * 255).astype(np.uint8))
    mask_pil = mask_pil.resize(original_size, Image.LANCZOS)
    
    # Apply mask to original image
    original_image.putalpha(mask_pil)
    
    return original_image

def process_folder(input_folder, output_folder, model, device):
    """Process all images in a folder."""
    # Create output folder
    os.makedirs(output_folder, exist_ok=True)
    
    # Get all image files
    image_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp']
    image_files = []
    
    for file in os.listdir(input_folder):
        if any(file.lower().endswith(ext) for ext in image_extensions):
            image_files.append(file)
    
    print(f"\nFound {len(image_files)} images to process")
    print(f"Output folder: {output_folder}\n")
    
    # Process images with progress bar
    successful = 0
    failed = 0
    
    for filename in tqdm(image_files, desc="Removing backgrounds"):
        try:
            input_path = os.path.join(input_folder, filename)
            
            # Remove background
            result_image = remove_background(input_path, model, device)
            
            # Save as PNG (to preserve transparency)
            output_filename = Path(filename).stem + ".png"
            output_path = os.path.join(output_folder, output_filename)
            result_image.save(output_path, "PNG")
            
            successful += 1
            
        except Exception as e:
            if "CUDA" in str(e) and device.type == 'cuda':
                print(f"\nCUDA error for {filename}: {str(e)}")
                print(f"Retrying {filename} on CPU...")
                try:
                    torch.cuda.empty_cache()
                    model.to('cpu')
                    result_image = remove_background(input_path, model, torch.device('cpu'))
                    
                    # Save as PNG
                    output_filename = Path(filename).stem + ".png"
                    output_path = os.path.join(output_folder, output_filename)
                    result_image.save(output_path, "PNG")
                    
                    successful += 1
                    
                    # Move model back to GPU
                    model.to(device)
                    continue
                except Exception as cpu_e:
                    print(f"Failed on CPU retry: {str(cpu_e)}")
                    # Ensure model is back on GPU
                    model.to(device)
            
            print(f"\nError processing {filename}: {str(e)}")
            failed += 1
    
    return successful, failed

def main():
    """Main function."""
    print("="*60)
    print("BACKGROUND REMOVAL - BiRefNet")
    print("="*60)
    print(f"Input folder: {INPUT_FOLDER}")
    print(f"Output folder: {OUTPUT_FOLDER}")
    print("="*60)
    
    # Check if CUDA is available
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
    else:
        print("GPU: Not available, using CPU")
    
    # Load model
    model, device = load_model()
    
    # Process images
    successful, failed = process_folder(INPUT_FOLDER, OUTPUT_FOLDER, model, device)
    
    # Summary
    print("\n" + "="*60)
    print("PROCESSING COMPLETE")
    print("="*60)
    print(f"✓ Successful: {successful}")
    print(f"✗ Failed: {failed}")
    print(f"Total: {successful + failed}")
    print("="*60)

if __name__ == "__main__":
    main()

Possible Extensions

• Web Interface: Integration into a Flask or FastAPI web application for browser-based uploads.
• Batch API: RESTful API endpoint for integration into automated workflows.
• Quality Control: Automatic evaluation of mask quality and flagging of problematic images.
• Multi-Model Support: Comparison of different segmentation models and selection of the best result.
• Cloud Deployment: Scaling processing on AWS Lambda or Google Cloud Functions for serverless batch jobs.