examples/batch_processing.py

"""Example of batch processing multiple images with PANDORA."""

import os
from pathlib import Path

from src.pandora_removal import PandoraRemoval, PandoraConfig


def main():
    """Demonstrate batch processing of multiple images."""

    # Configure the model
    config = PandoraConfig(
        model_path="stabilityai/stable-diffusion-2-1",
        device="cuda",
        max_steps=50,
        guidance_scale_ladg=7.5
    )

    # Initialize and load model
    print("Initializing PANDORA model...")
    model = PandoraRemoval(config=config)
    model.load_model()
    print("✓ Model loaded successfully!")

    # Define dataset structure
    # Expected structure:
    # dataset/
    #   ├── Images/
    #   │   ├── 001.jpg
    #   │   ├── 002.jpg
    #   │   └── ...
    #   └── Masks/
    #       ├── 001.png
    #       ├── 002.png
    #       └── ...

    dataset_path = "path/to/your/dataset"
    images_dir = os.path.join(dataset_path, "Images")
    masks_dir = os.path.join(dataset_path, "Masks")
    output_dir = "output/batch_results"

    # Check if directories exist
    if not Path(images_dir).exists():
        print(f"❌ Images directory not found: {images_dir}")
        print("Please update the dataset_path variable.")
        return

    if not Path(masks_dir).exists():
        print(f"❌ Masks directory not found: {masks_dir}")
        print("Please update the dataset_path variable.")
        return

    # Process batch
    print(f"\nProcessing images from: {images_dir}")
    print(f"Using masks from: {masks_dir}")
    print(f"Saving results to: {output_dir}\n")

    model.batch_process(
        images_dir=images_dir,
        masks_dir=masks_dir,
        output_dir=output_dir,
        border_size=17
    )

    print("\n✓ Batch processing complete!")


if __name__ == "__main__":
    main()