withoutbg 0.1.0__py3-none-any.whl

withoutbg/__init__.py

@@ -0,0 +1,18 @@
+"""withoutbg: AI-powered background removal with local and cloud options."""
+
+from .__version__ import __version__
+from .api import StudioAPI
+from .core import remove_background, remove_background_batch
+from .exceptions import APIError, ModelNotFoundError, WithoutBGError
+from .models import SnapModel
+
+__all__ = [
+    "remove_background",
+    "remove_background_batch",
+    "SnapModel",
+    "StudioAPI",
+    "WithoutBGError",
+    "ModelNotFoundError",
+    "APIError",
+    "__version__",
+]

withoutbg/__version__.py

@@ -0,0 +1,3 @@
+"""Version information for withoutbg."""
+
+__version__ = "0.1.0"

withoutbg/api.py

@@ -0,0 +1,241 @@
+"""Studio API client for cloud-based background removal."""
+
+import base64
+import io
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import requests
+from PIL import Image
+
+from .exceptions import APIError
+
+
+class StudioAPI:
+    """Client for withoutbg Studio API."""
+
+    def __init__(
+        self, api_key: Optional[str] = None, base_url: str = "https://api.withoutbg.com"
+    ):
+        """Initialize Studio API client.
+
+        Args:
+            api_key: API key for authentication
+            base_url: Base URL for API endpoints
+        """
+        self.api_key = api_key
+        self.base_url = base_url.rstrip("/")
+        self.session = requests.Session()
+
+        if api_key:
+            self.session.headers.update({"X-API-Key": api_key})
+
+    def _encode_image(self, image: Union[str, Path, Image.Image, bytes]) -> str:
+        """Encode image to base64 for API transmission."""
+        if isinstance(image, (str, Path)):
+            with open(image, "rb") as f:
+                image_bytes = f.read()
+        elif isinstance(image, Image.Image):
+            buffer = io.BytesIO()
+            image.save(buffer, format="PNG")
+            image_bytes = buffer.getvalue()
+        elif isinstance(image, bytes):
+            image_bytes = image
+        else:
+            raise APIError(f"Unsupported image type: {type(image)}")
+
+        return base64.b64encode(image_bytes).decode("utf-8")
+
+    def _decode_image(self, base64_string: str) -> Image.Image:
+        """Decode base64 string to PIL Image."""
+        image_bytes = base64.b64decode(base64_string)
+        with Image.open(io.BytesIO(image_bytes)) as img:
+            return img.copy()
+
+    def _resize_for_api(
+        self, image: Image.Image, max_size: int = 1024
+    ) -> tuple[Image.Image, tuple[int, int]]:
+        """Resize image for API transmission while preserving aspect ratio.
+
+        Args:
+            image: Input PIL Image
+            max_size: Maximum dimension size (default 1024)
+
+        Returns:
+            Tuple of (resized_image, original_size)
+        """
+        original_size = image.size
+        width, height = original_size
+
+        # If image is already smaller than max_size, return as-is
+        if max(width, height) <= max_size:
+            return image, original_size
+
+        # Calculate new dimensions preserving aspect ratio
+        if width > height:
+            new_width = max_size
+            new_height = int((height * max_size) / width)
+        else:
+            new_height = max_size
+            new_width = int((width * max_size) / height)
+
+        resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+        return resized_image, original_size
+
+    def _apply_alpha_channel(
+        self, original_image: Image.Image, alpha_image: Image.Image
+    ) -> Image.Image:
+        """Apply alpha channel to original image to remove background.
+
+        Args:
+            original_image: Original RGB/RGBA image
+            alpha_image: Grayscale alpha channel image
+
+        Returns:
+            RGBA image with background removed
+        """
+        # Ensure original image is in RGB mode
+        if original_image.mode != "RGB":
+            original_image = original_image.convert("RGB")
+
+        # Ensure alpha image is grayscale
+        if alpha_image.mode != "L":
+            alpha_image = alpha_image.convert("L")
+
+        # Resize alpha to match original image if needed
+        if alpha_image.size != original_image.size:
+            alpha_image = alpha_image.resize(
+                original_image.size, Image.Resampling.LANCZOS
+            )
+
+        # Create RGBA image by adding alpha channel
+        rgba_image = original_image.copy()
+        rgba_image.putalpha(alpha_image)
+
+        return rgba_image
+
+    def remove_background(
+        self, input_image: Union[str, Path, Image.Image, bytes], **kwargs: Any
+    ) -> Image.Image:
+        """Remove background using Studio API.
+
+        Args:
+            input_image: Input image
+            **kwargs: Additional API parameters
+
+        Returns:
+            PIL Image with background removed
+
+        Raises:
+            APIError: If API request fails
+        """
+        if not self.api_key:
+            raise APIError("API key required for Studio service")
+
+        try:
+            # Store original image for local alpha application
+            if isinstance(input_image, (str, Path)):
+                with Image.open(input_image) as img:
+                    original_image = img.copy()
+            elif isinstance(input_image, Image.Image):
+                original_image = input_image.copy()
+            elif isinstance(input_image, bytes):
+                with Image.open(io.BytesIO(input_image)) as img:
+                    original_image = img.copy()
+            else:
+                raise APIError(f"Unsupported image type: {type(input_image)}")
+
+            # Resize image for API transmission to optimize latency
+            api_image, original_size = self._resize_for_api(original_image)
+
+            # Encode resized image
+            encoded_image = self._encode_image(api_image)
+
+            # Prepare request for base64 endpoint
+            payload = {"image_base64": encoded_image}
+
+            # Make API request to base64 endpoint
+            response = self.session.post(
+                f"{self.base_url}/v1.0/alpha-channel-base64", json=payload, timeout=30
+            )
+
+            if response.status_code == 401:
+                raise APIError("Invalid API key")
+            elif response.status_code == 429:
+                raise APIError("Rate limit exceeded (7 requests per minute)")
+            elif response.status_code == 402:
+                raise APIError(
+                    "Insufficient credits. Get more credits at https://withoutbg.com/login"
+                )
+            elif response.status_code == 403:
+                raise APIError(
+                    "Credits expired. Top up your account to reactivate "
+                    "frozen credits at https://withoutbg.com/login"
+                )
+            elif not response.ok:
+                try:
+                    error_msg = response.json().get("error", "Unknown API error")
+                except Exception:
+                    error_msg = f"HTTP {response.status_code}: {response.text}"
+                raise APIError(f"API request failed: {error_msg}")
+
+            # Decode response
+            result_data = response.json()
+
+            if "alpha_base64" not in result_data:
+                raise APIError(
+                    f"Invalid API response format\n\n"
+                    f"More info: sample response: {result_data}"
+                )
+
+            # Decode alpha channel
+            alpha_image = self._decode_image(result_data["alpha_base64"])
+
+            # Resize alpha channel back to original dimensions if needed
+            if alpha_image.size != original_size:
+                alpha_image = alpha_image.resize(
+                    original_size, Image.Resampling.LANCZOS
+                )
+
+            # Apply alpha channel to original image locally
+            return self._apply_alpha_channel(original_image, alpha_image)
+
+        except requests.RequestException as e:
+            raise APIError(f"Network error: {str(e)}") from e
+        except Exception as e:
+            if isinstance(e, APIError):
+                raise
+            raise APIError(f"Unexpected error: {str(e)}") from e
+
+    def get_usage(self) -> dict[str, Any]:
+        """Get current API usage statistics.
+
+        Returns:
+            Dictionary with usage information
+        """
+        if not self.api_key:
+            raise APIError("API key required")
+
+        try:
+            response = self.session.get(f"{self.base_url}/available-credit")
+            response.raise_for_status()
+            result: dict[str, Any] = response.json()
+            return result
+
+        except requests.RequestException as e:
+            raise APIError(f"Failed to get usage: {str(e)}") from e
+
+    def get_models(self) -> dict[str, Any]:
+        """Get available models and their capabilities.
+
+        Returns:
+            Dictionary with model information
+        """
+        try:
+            response = self.session.get(f"{self.base_url}/v1/models")
+            response.raise_for_status()
+            result: dict[str, Any] = response.json()
+            return result
+
+        except requests.RequestException as e:
+            raise APIError(f"Failed to get models: {str(e)}") from e

withoutbg/cli.py

@@ -0,0 +1,264 @@
+"""Command-line interface for withoutbg."""
+
+import sys
+from pathlib import Path
+from typing import Any, Optional
+
+import click
+from click._termui_impl import ProgressBar
+from PIL import Image
+
+from . import __version__, remove_background
+from .exceptions import WithoutBGError
+
+
+@click.command()
+@click.argument("input_path", type=click.Path(exists=True, path_type=Path))
+@click.option(
+    "--output",
+    "-o",
+    type=click.Path(path_type=Path),
+    help="Output file path (default: adds -withoutbg suffix)",
+)
+@click.option(
+    "--api-key",
+    envvar="WITHOUTBG_API_KEY",
+    help="API key for Studio service (or set WITHOUTBG_API_KEY env var)",
+)
+@click.option(
+    "--use-api", is_flag=True, help="Use Studio API instead of local Snap model"
+)
+@click.option(
+    "--model",
+    default="snap",
+    type=click.Choice(["snap", "studio"]),
+    help="Model to use (snap=local, studio=API)",
+)
+@click.option(
+    "--batch",
+    is_flag=True,
+    help="Process all images in directory (if input is directory)",
+)
+@click.option(
+    "--output-dir",
+    type=click.Path(path_type=Path),
+    help="Output directory for batch processing",
+)
+@click.option(
+    "--format",
+    type=click.Choice(["png", "jpg", "webp"]),
+    default="png",
+    help="Output image format",
+)
+@click.option("--quality", type=int, default=95, help="Output quality for JPEG (1-100)")
+@click.option("--verbose", "-v", is_flag=True, help="Enable verbose output")
+@click.version_option(version=__version__)
+def main(
+    input_path: Path,
+    output: Optional[Path],
+    api_key: Optional[str],
+    use_api: bool,
+    model: str,
+    batch: bool,
+    output_dir: Optional[Path],
+    format: str,
+    quality: int,
+    verbose: bool,
+) -> None:
+    """Remove background from images using AI.
+
+    Examples:
+
+        # Process single image with local Snap model
+        withoutbg image.jpg
+
+        # Use Studio API for best quality processing
+        withoutbg image.jpg --use-api --api-key sk_...
+
+        # Process all images in directory
+        withoutbg photos/ --batch --output-dir results/
+
+        # Specify output format and quality
+        withoutbg image.jpg --format jpg --quality 90
+    """
+
+    try:
+        # Determine if using API
+        using_api = use_api or api_key or model == "studio"
+
+        if using_api and not api_key:
+            click.echo("Error: API key required when using Studio service", err=True)
+            click.echo(
+                "Set WITHOUTBG_API_KEY environment variable or use --api-key", err=True
+            )
+            sys.exit(1)
+
+        # Configure model
+        actual_model = "studio" if using_api else "snap"
+
+        if verbose:
+            mode = "Studio API" if using_api else "Local Snap model"
+            click.echo(f"Using {mode} for processing...")
+
+        # Process images
+        if batch or input_path.is_dir():
+            _process_batch(
+                input_path, output_dir, api_key, actual_model, format, quality, verbose
+            )
+        else:
+            _process_single(
+                input_path, output, api_key, actual_model, format, quality, verbose
+            )
+
+        if verbose:
+            api_msg = ""
+            if not using_api:
+                api_msg = " (Want best quality? Try withoutbg.com)"
+            click.echo(f"✅ Processing complete!{api_msg}")
+
+    except WithoutBGError as e:
+        click.echo(f"Error: {e}", err=True)
+        sys.exit(1)
+    except KeyboardInterrupt:
+        click.echo("\n❌ Processing cancelled", err=True)
+        sys.exit(1)
+    except Exception as e:
+        click.echo(f"Unexpected error: {e}", err=True)
+        sys.exit(1)
+
+
+def _process_single(
+    input_path: Path,
+    output_path: Optional[Path],
+    api_key: Optional[str],
+    model: str,
+    format: str,
+    quality: int,
+    verbose: bool,
+) -> None:
+    """Process a single image."""
+
+    if not output_path:
+        # Generate output filename
+        stem = input_path.stem
+        output_path = input_path.parent / f"{stem}-withoutbg.{format}"
+
+    if verbose:
+        click.echo(f"Processing: {input_path}")
+        click.echo(f"Output: {output_path}")
+
+    # Remove background
+    bar: ProgressBar
+    with click.progressbar(length=1, label="Removing background") as bar:
+        result = remove_background(input_path, api_key=api_key, model_name=model)
+        bar.update(1)
+
+    # Save result
+    save_kwargs: dict[str, Any] = {}
+    if format.lower() == "jpg":
+        # Convert RGBA to RGB for JPEG
+        if result.mode == "RGBA":
+            background = Image.new("RGB", result.size, (255, 255, 255))
+            background.paste(result, mask=result.split()[-1])  # Use alpha as mask
+            result = background
+        save_kwargs["quality"] = quality
+        save_kwargs["optimize"] = True
+    elif format.lower() == "webp":
+        save_kwargs["quality"] = quality
+        save_kwargs["method"] = 6  # Best compression
+
+    # Map format names to PIL format strings
+    pil_format = {"jpg": "JPEG", "jpeg": "JPEG", "png": "PNG", "webp": "WEBP"}
+    result.save(
+        output_path,
+        format=pil_format.get(format.lower(), format.upper()),
+        **save_kwargs,
+    )
+
+
+def _process_batch(
+    input_dir: Path,
+    output_dir: Optional[Path],
+    api_key: Optional[str],
+    model: str,
+    format: str,
+    quality: int,
+    verbose: bool,
+) -> None:
+    """Process multiple images in a directory."""
+
+    # Find image files
+    image_extensions = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"}
+
+    if input_dir.is_file():
+        # Single file specified with --batch flag
+        image_files = [input_dir]
+        input_dir = input_dir.parent
+    else:
+        # Directory specified
+        image_files = [
+            f
+            for f in input_dir.iterdir()
+            if f.is_file() and f.suffix.lower() in image_extensions
+        ]
+
+    if not image_files:
+        click.echo("No image files found in directory", err=True)
+        sys.exit(1)
+
+    # Set up output directory
+    if not output_dir:
+        output_dir = input_dir / "withoutbg-results"
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    if verbose:
+        click.echo(f"Found {len(image_files)} images")
+        click.echo(f"Output directory: {output_dir}")
+
+    # Process images
+    with click.progressbar(image_files, label="Processing images") as bar:
+        for image_file in bar:
+            try:
+                # Generate output path
+                output_path = output_dir / f"{image_file.stem}-withoutbg.{format}"
+
+                # Remove background
+                result = remove_background(
+                    image_file, api_key=api_key, model_name=model
+                )
+
+                # Save result
+                save_kwargs: dict[str, Any] = {}
+                if format.lower() == "jpg":
+                    if result.mode == "RGBA":
+                        background = Image.new("RGB", result.size, (255, 255, 255))
+                        background.paste(result, mask=result.split()[-1])
+                        result = background
+                    save_kwargs["quality"] = quality
+                    save_kwargs["optimize"] = True
+                elif format.lower() == "webp":
+                    save_kwargs["quality"] = quality
+                    save_kwargs["method"] = 6
+
+                # Map format names to PIL format strings
+                pil_format = {
+                    "jpg": "JPEG",
+                    "jpeg": "JPEG",
+                    "png": "PNG",
+                    "webp": "WEBP",
+                }
+                result.save(
+                    output_path,
+                    format=pil_format.get(format.lower(), format.upper()),
+                    **save_kwargs,
+                )
+
+            except Exception as e:
+                if verbose:
+                    click.echo(f"\n❌ Failed to process {image_file}: {e}", err=True)
+                continue
+
+
+if __name__ == "__main__":
+    main()

withoutbg/core.py

@@ -0,0 +1,109 @@
+"""Core background removal functionality."""
+
+from pathlib import Path
+from typing import Any, Optional, Union
+
+from PIL import Image
+
+from .api import StudioAPI
+from .exceptions import WithoutBGError
+from .models import SnapModel
+
+
+def remove_background(
+    input_image: Union[str, Path, Image.Image, bytes],
+    output_path: Optional[Union[str, Path]] = None,
+    api_key: Optional[str] = None,
+    model_name: str = "snap",
+    **kwargs: Any,
+) -> Image.Image:
+    """Remove background from an image.
+
+    Args:
+        input_image: Input image as file path, PIL Image, or bytes
+        output_path: Optional path to save the result
+        api_key: API key for Studio service (uses local Snap model if None)
+        model_name: Model to use ("snap" for local, "studio" for API)
+        **kwargs: Additional arguments passed to the model/API
+
+    Returns:
+        PIL Image with background removed
+
+    Examples:
+        >>> # Local processing with Snap model
+        >>> result = remove_background("input.jpg")
+
+        >>> # Cloud processing with Studio API
+        >>> result = remove_background("input.jpg", api_key="sk_...")
+
+        >>> # Save result directly
+        >>> remove_background("input.jpg", output_path="output.png")
+    """
+    try:
+        if api_key or model_name == "studio":
+            # Use Studio API
+            api = StudioAPI(api_key)
+            result = api.remove_background(input_image, **kwargs)
+        else:
+            # Use local Snap model
+            model = SnapModel()
+            result = model.remove_background(input_image, **kwargs)
+
+        if output_path:
+            result.save(output_path)
+
+        return result
+
+    except Exception as e:
+        raise WithoutBGError(f"Background removal failed: {str(e)}") from e
+
+
+def remove_background_batch(
+    input_images: list[Union[str, Path, Image.Image, bytes]],
+    output_dir: Optional[Union[str, Path]] = None,
+    api_key: Optional[str] = None,
+    model_name: str = "snap",
+    **kwargs: Any,
+) -> list[Image.Image]:
+    """Remove background from multiple images.
+
+    Args:
+        input_images: List of input images
+        output_dir: Directory to save results (optional)
+        api_key: API key for Studio service
+        model_name: Model to use
+        **kwargs: Additional arguments
+
+    Returns:
+        List of PIL Images with backgrounds removed
+    """
+    results = []
+
+    for i, input_image in enumerate(input_images):
+        output_path = None
+        if output_dir:
+            output_dir = Path(output_dir)
+            output_dir.mkdir(parents=True, exist_ok=True)
+
+            # Try to get original filename
+            if isinstance(input_image, (str, Path)):
+                input_path = Path(input_image)
+                stem = input_path.stem
+                suffix = input_path.suffix or ".png"
+                output_filename = f"{stem}-withoutbg{suffix}"
+            else:
+                # For PIL Images or bytes, use numbered fallback
+                output_filename = f"output_{i:04d}-withoutbg.png"
+
+            output_path = output_dir / output_filename
+
+        result = remove_background(
+            input_image,
+            output_path=output_path,
+            api_key=api_key,
+            model_name=model_name,
+            **kwargs,
+        )
+        results.append(result)
+
+    return results

withoutbg/exceptions.py

@@ -0,0 +1,31 @@
+"""Custom exceptions for withoutbg package."""
+
+
+class WithoutBGError(Exception):
+    """Base exception for withoutbg package."""
+
+    pass
+
+
+class ModelNotFoundError(WithoutBGError):
+    """Raised when a model file cannot be found or loaded."""
+
+    pass
+
+
+class APIError(WithoutBGError):
+    """Raised when API requests fail."""
+
+    pass
+
+
+class InvalidImageError(WithoutBGError):
+    """Raised when image input is invalid or corrupted."""
+
+    pass
+
+
+class ConfigurationError(WithoutBGError):
+    """Raised when configuration is invalid."""
+
+    pass