How to Batch Remove Backgrounds from S3 Images

import os
import random
import time
import shutil
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path, PurePosixPath
from typing import Any, Dict, Iterator, List

import boto3
import requests
from botocore.exceptions import ClientError
from dotenv import load_dotenv
from tqdm import tqdm


#edit the following

S3_INPUT_BUCKET = "some-bucket"
S3_INPUT_PREFIX = "raw-images/"

S3_OUTPUT_BUCKET = "some-bucket"
S3_OUTPUT_PREFIX = "removed-backgrounds/"

BATCH_SIZE = 50
MAX_WORKERS = 2
SKIP_EXISTING_OUTPUTS = True
MAX_INPUT_FILE_SIZE_BYTES = 30 * 1024 * 1024

API_MAX_CONCURRENT_REQUESTS = 2
API_REQUESTS_PER_SECOND = 0.75
API_BURST_SIZE = 1

PRESIGNED_URL_EXPIRATION_SECONDS = 15 * 60
BACKGROUNDERASE_API_URL = "https://api.backgrounderase.com/v2"

TEMP_DIR = Path("./temp")

REQUEST_TIMEOUT_SECONDS = 120
MAX_RETRIES = 3
RETRY_SLEEP_SECONDS = 2

SUPPORTED_IMAGE_EXTENSIONS = {
    ".jpg",
    ".jpeg",
    ".png",
    ".webp",
    ".bmp",
    ".tif",
    ".tiff",
}


load_dotenv()

BACKGROUNDERASE_API_KEY = os.getenv("BACKGROUNDERASE_API_KEY")

if not BACKGROUNDERASE_API_KEY:
    raise RuntimeError(
        "Missing BACKGROUNDERASE_API_KEY. Add it to your .env file."
    )


s3 = boto3.client("s3")
api_connection_semaphore = threading.Semaphore(API_MAX_CONCURRENT_REQUESTS)


class TokenBucket:
    def __init__(self, rate_per_second: float, capacity: int):
        self.rate_per_second = rate_per_second
        self.capacity = capacity
        self.tokens = capacity
        self.updated_at = time.monotonic()
        self.lock = threading.Lock()

    def wait_for_token(self) -> None:
        while True:
            with self.lock:
                now = time.monotonic()
                elapsed = now - self.updated_at
                self.updated_at = now

                self.tokens = min(
                    self.capacity,
                    self.tokens + elapsed * self.rate_per_second,
                )

                if self.tokens >= 1:
                    self.tokens -= 1
                    return

                missing = 1 - self.tokens
                sleep_for = missing / self.rate_per_second

            time.sleep(sleep_for)


api_rate_limiter = TokenBucket(
    rate_per_second=API_REQUESTS_PER_SECOND,
    capacity=API_BURST_SIZE,
)


def normalize_prefix(prefix: str) -> str:
    if not prefix:
        return ""
    return prefix.strip("/") + "/"


def is_supported_image_key(key: str) -> bool:
    suffix = PurePosixPath(key).suffix.lower()
    return suffix in SUPPORTED_IMAGE_EXTENSIONS


def format_bytes(size: int) -> str:
    return f"{size / (1024 * 1024):.1f} MB"


def iter_s3_image_objects(bucket: str, prefix: str) -> Iterator[Dict[str, Any]]:
    paginator = s3.get_paginator("list_objects_v2")

    for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
        for obj in page.get("Contents", []):
            key = obj["Key"]
            size = obj.get("Size", 0)

            if size == 0:
                continue

            if not is_supported_image_key(key):
                continue

            yield obj


def count_s3_images(bucket: str, prefix: str) -> int:
    total = 0

    for _ in iter_s3_image_objects(bucket, prefix):
        total += 1

    return total


def batched(
    iterator: Iterator[Dict[str, Any]],
    batch_size: int,
) -> Iterator[List[Dict[str, Any]]]:
    batch = []

    for item in iterator:
        batch.append(item)

        if len(batch) >= batch_size:
            yield batch
            batch = []

    if batch:
        yield batch


def generate_presigned_url(bucket: str, key: str) -> str:
    return s3.generate_presigned_url(
        ClientMethod="get_object",
        Params={
            "Bucket": bucket,
            "Key": key,
        },
        ExpiresIn=PRESIGNED_URL_EXPIRATION_SECONDS,
    )


def get_relative_key(input_key: str, input_prefix: str) -> str:
    if input_prefix and input_key.startswith(input_prefix):
        return input_key[len(input_prefix):].lstrip("/")

    return PurePosixPath(input_key).name


def make_output_relative_key(relative_input_key: str) -> str:
    return str(PurePosixPath(relative_input_key).with_suffix(".png"))


def make_temp_output_path(relative_output_key: str) -> Path:
    return TEMP_DIR / relative_output_key


def make_s3_output_key(relative_output_key: str) -> str:
    output_prefix = normalize_prefix(S3_OUTPUT_PREFIX)
    return output_prefix + relative_output_key.lstrip("/")


def reset_temp_dir() -> None:
    if TEMP_DIR.exists():
        shutil.rmtree(TEMP_DIR)

    TEMP_DIR.mkdir(parents=True, exist_ok=True)


def delete_temp_dir() -> None:
    if TEMP_DIR.exists():
        shutil.rmtree(TEMP_DIR)


def s3_object_exists(bucket: str, key: str) -> bool:
    try:
        s3.head_object(Bucket=bucket, Key=key)
        return True
    except ClientError as exc:
        status = exc.response.get("ResponseMetadata", {}).get("HTTPStatusCode")
        if status == 404:
            return False
        raise


def call_backgrounderase_api(image_url: str) -> bytes:
    headers = {
        "x-api-key": BACKGROUNDERASE_API_KEY,
    }

    
    files = {
        "image_url": (None, image_url),
        "format": (None, "png"),
        "channels": (None, "rgba"),
    }

    last_error = None
    attempts_made = 0

    for attempt in range(1, MAX_RETRIES + 1):
        attempts_made = attempt
        response = None
        should_retry = False

        try:
            api_rate_limiter.wait_for_token()

            with api_connection_semaphore:
                response = requests.post(
                    BACKGROUNDERASE_API_URL,
                    headers=headers,
                    files=files,
                    timeout=REQUEST_TIMEOUT_SECONDS,
                )

            content_type = response.headers.get("Content-Type", "")

            if response.status_code == 200 and content_type.startswith("image/"):
                return response.content

            body_preview = response.text[:500] if response.text else ""
            last_error = (
                f"API returned status={response.status_code}, "
                f"content_type={content_type}, body={body_preview}"
            )
            should_retry = should_retry_response(response)

        except requests.RequestException as exc:
            last_error = str(exc)
            should_retry = True

        if not should_retry:
            break

        if attempt < MAX_RETRIES:
            time.sleep(get_retry_sleep_seconds(response, attempt))

    raise RuntimeError(
        f"BackgroundErase request failed after {attempts_made} attempts: {last_error}"
    )


def should_retry_response(response: Any) -> bool:
    if response is None:
        return True

    if response.status_code in {408, 429}:
        return True

    if 500 <= response.status_code < 600:
        return True

    return False


def get_retry_sleep_seconds(response: Any, attempt: int) -> float:
    if response is not None and response.status_code == 429:
        retry_after = response.headers.get("Retry-After")

        if retry_after:
            try:
                return max(float(retry_after), 0)
            except ValueError:
                pass

        return min(60, 2 ** attempt) + random.uniform(0, 1.5)

    return RETRY_SLEEP_SECONDS * attempt + random.uniform(0, 0.5)


def upload_file_to_s3(local_path: Path, bucket: str, key: str) -> None:
    s3.upload_file(
        Filename=str(local_path),
        Bucket=bucket,
        Key=key,
        ExtraArgs={
            "ContentType": "image/png",
        },
    )


def process_one(s3_object: Dict[str, Any], input_prefix: str) -> Dict[str, Any]:
    started_at = time.time()

    input_key = s3_object["Key"]
    input_size = s3_object.get("Size", 0)

    relative_input_key = get_relative_key(input_key, input_prefix)
    relative_output_key = make_output_relative_key(relative_input_key)
    temp_output_path = make_temp_output_path(relative_output_key)
    s3_output_key = make_s3_output_key(relative_output_key)

    if input_size > MAX_INPUT_FILE_SIZE_BYTES:
        return {
            "status": "skipped",
            "reason": (
                "larger than the 30 MB API limit "
                f"({format_bytes(input_size)})"
            ),
            "input_key": input_key,
            "output_key": s3_output_key,
            "elapsed": time.time() - started_at,
        }

    if SKIP_EXISTING_OUTPUTS and s3_object_exists(S3_OUTPUT_BUCKET, s3_output_key):
        return {
            "status": "skipped",
            "input_key": input_key,
            "output_key": s3_output_key,
            "elapsed": time.time() - started_at,
        }

    presigned_url = generate_presigned_url(S3_INPUT_BUCKET, input_key)
    output_bytes = call_backgrounderase_api(presigned_url)

    temp_output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(temp_output_path, "wb") as f:
        f.write(output_bytes)

    upload_file_to_s3(
        local_path=temp_output_path,
        bucket=S3_OUTPUT_BUCKET,
        key=s3_output_key,
    )

    return {
        "status": "success",
        "input_key": input_key,
        "output_key": s3_output_key,
        "elapsed": time.time() - started_at,
    }




def process_batch(
    batch: List[Dict[str, Any]],
    input_prefix: str,
    progress: tqdm,
) -> Dict[str, int]:
    counts = {
        "success": 0,
        "skipped": 0,
        "failed": 0,
    }

    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = {
            executor.submit(process_one, obj, input_prefix): obj["Key"]
            for obj in batch
        }

        for future in as_completed(futures):
            input_key = futures[future]

            try:
                result = future.result()
                status = result["status"]
                counts[status] += 1

                if status == "success":
                    tqdm.write(
                        "[SUCCESS] "
                        f"{result['input_key']} -> {result['output_key']} "
                        f"({result['elapsed']:.3f}s)"
                    )
                else:
                    reason = result.get("reason")

                    if reason:
                        tqdm.write(f"[SKIPPED] {result['input_key']} - {reason}")
                    else:
                        tqdm.write(
                            f"[SKIPPED] {result['input_key']} -> {result['output_key']}"
                        )

            except Exception as exc:
                counts["failed"] += 1
                tqdm.write(f"[FAILED] {input_key} - {exc}")

            progress.update(1)

    return counts


def main() -> None:
    input_prefix = normalize_prefix(S3_INPUT_PREFIX)
    output_prefix = normalize_prefix(S3_OUTPUT_PREFIX)

    reset_temp_dir()

    try:
        print("Counting input images...")
        total_images = count_s3_images(S3_INPUT_BUCKET, input_prefix)

        if total_images == 0:
            print("No supported images found.")
            return

        print(f"Found {total_images:,} images.")
        print(f"Input:  s3://{S3_INPUT_BUCKET}/{input_prefix}")
        print(f"Output: s3://{S3_OUTPUT_BUCKET}/{output_prefix}")
        print(f"Workers: {MAX_WORKERS}")

        totals = {
            "success": 0,
            "skipped": 0,
            "failed": 0,
        }

        image_iterator = iter_s3_image_objects(S3_INPUT_BUCKET, input_prefix)

        with tqdm(total=total_images, unit="img") as progress:
            for batch in batched(image_iterator, BATCH_SIZE):
                batch_counts = process_batch(
                    batch=batch,
                    input_prefix=input_prefix,
                    progress=progress,
                )

                for key, value in batch_counts.items():
                    totals[key] += value

        print("")
        print("Done.")
        print(f"Successful: {totals['success']}")
        print(f"Skipped:    {totals['skipped']}")
        print(f"Failed:     {totals['failed']}")
        print(f"Outputs uploaded to: s3://{S3_OUTPUT_BUCKET}/{output_prefix}")

    finally:
        delete_temp_dir()


if __name__ == "__main__":
    main()