examples/basic_removal.py

"""Basic example of using PANDORA for object removal."""

from pathlib import Path
from PIL import Image

from src.pandora_removal import PandoraRemoval, PandoraConfig


def main():
    """Demonstrate basic object removal."""

    # Configure the model
    config = PandoraConfig(
        model_path="stabilityai/stable-diffusion-2-1",
        device="cuda",  # Change to "cpu" if no GPU available
        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 input and output paths
    image_path = "path/to/your/image.jpg"
    mask_path = "path/to/your/mask.png"
    output_path = "output/result.png"

    # Check if files exist
    if not Path(image_path).exists():
        print(f"❌ Image not found: {image_path}")
        print("Please update the image_path variable with your image path.")
        return

    if not Path(mask_path).exists():
        print(f"❌ Mask not found: {mask_path}")
        print("Please update the mask_path variable with your mask path.")
        return

    # Load images
    print("\nLoading image and mask...")
    image = Image.open(image_path).convert("RGB")
    mask = Image.open(mask_path).convert("L")

    print(f"Image size: {image.size}")
    print(f"Mask size: {mask.size}")

    # Remove object
    print("\nRemoving object...")
    result = model.remove_object(
        image=image,
        mask=mask,
        border_size=17,  # Adjust based on object size
        guidance_scale=7.5,
        num_steps=50
    )

    # Save result
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    result.save(output_path)
    print(f"✓ Result saved to: {output_path}")


if __name__ == "__main__":
    main()