snipbg 2.0.0__py3-none-any.whl

snipbg/__init__.py

@@ -0,0 +1,60 @@
+"""snipbg - remove image backgrounds with ONNX segmentation models.
+
+    from snipbg import remove
+    png_bytes = remove(open("photo.jpg", "rb").read())
+
+Attributes are resolved lazily so that ``snipbg.models`` and
+``snipbg.download`` stay importable without numpy, Pillow or onnxruntime
+installed (the Docker model stage relies on this).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+__version__ = "2.0.0"
+
+__all__ = [
+    "DEFAULT_MODEL",
+    "MODELS",
+    "DownloadError",
+    "ModelSpec",
+    "Remover",
+    "__version__",
+    "cache_dir",
+    "ensure_model",
+    "get_remover",
+    "list_models",
+    "remove",
+]
+
+if TYPE_CHECKING:
+    from .download import DownloadError, cache_dir, ensure_model
+    from .models import DEFAULT_MODEL, MODELS, ModelSpec, list_models
+    from .session import Remover, get_remover, remove
+
+_LAZY: dict[str, str] = {
+    "DEFAULT_MODEL": "snipbg.models",
+    "MODELS": "snipbg.models",
+    "ModelSpec": "snipbg.models",
+    "list_models": "snipbg.models",
+    "DownloadError": "snipbg.download",
+    "cache_dir": "snipbg.download",
+    "ensure_model": "snipbg.download",
+    "Remover": "snipbg.session",
+    "get_remover": "snipbg.session",
+    "remove": "snipbg.session",
+}
+
+
+def __getattr__(name: str) -> Any:
+    module_name = _LAZY.get(name)
+    if module_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    import importlib
+
+    return getattr(importlib.import_module(module_name), name)
+
+
+def __dir__() -> list[str]:
+    return sorted(__all__)

snipbg/cli.py

@@ -0,0 +1,290 @@
+"""Command line interface: ``snipbg remove|download|models|serve``."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from . import __version__
+from .models import DEFAULT_MODEL, MODELS, get_spec
+
+_EPILOG = """\
+examples:
+  snipbg remove photo.jpg
+  snipbg remove photo.jpg -o out.png --refine --trim
+  snipbg remove *.jpg -o cutouts/ -m isnet-general-use
+  cat photo.jpg | snipbg remove - > out.png
+  snipbg remove photo.jpg --bg white -o flat.jpg
+  snipbg download u2net isnet-general-use
+  snipbg serve --port 3001
+"""
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="snipbg",
+        description="Remove image backgrounds with ONNX segmentation models.",
+        epilog=_EPILOG,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--version", action="version", version=f"snipbg {__version__}")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    # remove ----------------------------------------------------------------
+    remove_parser = subparsers.add_parser(
+        "remove",
+        help="remove the background from one or more images",
+        description="Remove the background from one or more images.",
+    )
+    remove_parser.add_argument("inputs", nargs="+", help="image files, or - for stdin")
+    remove_parser.add_argument(
+        "-o",
+        "--output",
+        help="output file, output directory, or - for stdout "
+        "(default: <name>_nobg.png next to each input)",
+    )
+    remove_parser.add_argument(
+        "-m", "--model", default=DEFAULT_MODEL, help=f"model to use (default: {DEFAULT_MODEL})"
+    )
+    remove_parser.add_argument(
+        "--only-mask", action="store_true", help="write the grayscale mask instead of the cutout"
+    )
+    remove_parser.add_argument(
+        "--refine", action="store_true", help="harden and feather the mask edges"
+    )
+    remove_parser.add_argument("--foreground-threshold", type=int, default=240, metavar="N")
+    remove_parser.add_argument("--background-threshold", type=int, default=15, metavar="N")
+    remove_parser.add_argument(
+        "--erode", type=int, default=0, metavar="N", help="shrink the mask edge by N pixels"
+    )
+    remove_parser.add_argument(
+        "--blur", type=float, default=1.0, metavar="R", help="feather radius used with --refine"
+    )
+    remove_parser.add_argument(
+        "--bg",
+        "--background",
+        dest="background",
+        metavar="COLOR",
+        help="composite onto a colour (name, #rrggbb) instead of transparency",
+    )
+    remove_parser.add_argument(
+        "--trim", action="store_true", help="crop fully transparent borders away"
+    )
+    remove_parser.add_argument(
+        "-f", "--format", dest="image_format", help="output format (default: from the suffix, PNG)"
+    )
+    remove_parser.add_argument("-q", "--quality", type=int, help="JPEG/WEBP quality")
+    remove_parser.set_defaults(func=_cmd_remove)
+
+    # download --------------------------------------------------------------
+    download_parser = subparsers.add_parser(
+        "download", help="pre-fetch model weights into the cache"
+    )
+    download_parser.add_argument("models", nargs="*", help="model names (default: the default one)")
+    download_parser.add_argument("--all", action="store_true", help="download every known model")
+    download_parser.add_argument(
+        "--check", action="store_true", help="re-hash cached files and report mismatches"
+    )
+    download_parser.set_defaults(func=_cmd_download)
+
+    # models ----------------------------------------------------------------
+    models_parser = subparsers.add_parser("models", help="list the available models")
+    models_parser.add_argument("--json", action="store_true", help="machine readable output")
+    models_parser.set_defaults(func=_cmd_models)
+
+    # serve -----------------------------------------------------------------
+    serve_parser = subparsers.add_parser("serve", help="run the HTTP API")
+    serve_parser.add_argument("--host", default="0.0.0.0", help="bind address (default: 0.0.0.0)")
+    serve_parser.add_argument("--port", type=int, default=3001, help="port (default: 3001)")
+    serve_parser.add_argument("--reload", action="store_true", help="reload on code changes")
+    serve_parser.add_argument("--workers", type=int, default=1, help="worker processes")
+    serve_parser.set_defaults(func=_cmd_serve)
+
+    return parser
+
+
+def _output_path(source: Path, output: str | None, suffix: str, many: bool) -> Path:
+    if output is None:
+        return source.with_name(f"{source.stem}_nobg{suffix}")
+    target = Path(output)
+    if many or target.is_dir() or output.endswith(("/", "\\")):
+        target.mkdir(parents=True, exist_ok=True)
+        return target / f"{source.stem}_nobg{suffix}"
+    if target.parent != Path(""):
+        target.parent.mkdir(parents=True, exist_ok=True)
+    return target
+
+
+def _cmd_remove(args: argparse.Namespace) -> int:
+    from .images import encode_image, format_for_path, load_image
+    from .session import get_remover
+
+    remover = get_remover(args.model)
+    options = dict(
+        only_mask=args.only_mask,
+        refine=args.refine,
+        foreground_threshold=args.foreground_threshold,
+        background_threshold=args.background_threshold,
+        erode_size=args.erode,
+        blur_radius=args.blur,
+        background=args.background,
+        trim_border=args.trim,
+    )
+
+    inputs = list(args.inputs)
+    many = len(inputs) > 1
+    to_stdout = args.output == "-"
+    failures = 0
+
+    for name in inputs:
+        try:
+            if name == "-":
+                data = sys.stdin.buffer.read()
+                image = load_image(data)
+                source = Path("stdin.png")
+            else:
+                source = Path(name)
+                image = load_image(source)
+
+            result = remover.remove(image, **options)
+
+            if name == "-" and args.output is None:
+                to_stdout = True
+            image_format = (args.image_format or "PNG").upper()
+            if not to_stdout:
+                suffix = f".{(args.image_format or 'png').lower().replace('jpeg', 'jpg')}"
+                destination = _output_path(source, args.output, suffix, many)
+                if args.image_format is None:
+                    image_format = format_for_path(destination)
+
+            payload = encode_image(result, image_format, quality=args.quality)
+            if to_stdout:
+                sys.stdout.buffer.write(payload)
+            else:
+                destination.write_bytes(payload)
+                print(f"{name} -> {destination}", file=sys.stderr)
+        except (ValueError, OSError) as exc:
+            print(f"error: {name}: {exc}", file=sys.stderr)
+            failures += 1
+
+    return 1 if failures else 0
+
+
+def _cmd_download(args: argparse.Namespace) -> int:
+    from .download import DownloadError, download_model, is_cached, model_path, verify_model
+
+    names = list(MODELS) if args.all else (args.models or [DEFAULT_MODEL])
+    failures = 0
+
+    for name in names:
+        try:
+            spec = get_spec(name)
+        except ValueError as exc:
+            print(f"error: {exc}", file=sys.stderr)
+            failures += 1
+            continue
+
+        if args.check:
+            if not model_path(spec).is_file():
+                print(f"{spec.name}: not cached")
+                continue
+            ok = verify_model(spec)
+            print(f"{spec.name}: {'ok' if ok else 'CORRUPT - re-download with snipbg download'}")
+            failures += 0 if ok else 1
+            continue
+
+        if is_cached(spec):
+            print(f"{spec.name}: already cached at {model_path(spec)}")
+            continue
+
+        print(f"{spec.name}: downloading {spec.size_mb:.1f} MB from {spec.url}", file=sys.stderr)
+        try:
+            path = download_model(spec, progress=True, force=True)
+            print(f"{spec.name}: saved to {path}")
+        except DownloadError as exc:
+            print(f"error: {exc}", file=sys.stderr)
+            failures += 1
+
+    return 1 if failures else 0
+
+
+def _cmd_models(args: argparse.Namespace) -> int:
+    from .download import is_cached, model_path
+
+    specs = list(MODELS.values())
+    if args.json:
+        import json
+
+        print(
+            json.dumps(
+                [
+                    {
+                        "name": s.name,
+                        "filename": s.filename,
+                        "url": s.url,
+                        "sha256": s.sha256,
+                        "size_bytes": s.size_bytes,
+                        "input_size": list(s.input_size),
+                        "license": s.license,
+                        "description": s.description,
+                        "cached": is_cached(s),
+                        "path": str(model_path(s)),
+                    }
+                    for s in specs
+                ],
+                indent=2,
+            )
+        )
+        return 0
+
+    width = max(len(s.name) for s in specs)
+    print(f"{'MODEL'.ljust(width)}  {'SIZE':>8}  {'INPUT':>9}  {'CACHED':>6}  DESCRIPTION")
+    for spec in specs:
+        size = f"{spec.size_mb:.0f} MB"
+        shape = f"{spec.input_size[0]}x{spec.input_size[1]}"
+        cached = "yes" if is_cached(spec) else "no"
+        default = " (default)" if spec.name == DEFAULT_MODEL else ""
+        print(
+            f"{spec.name.ljust(width)}  {size:>8}  {shape:>9}  {cached:>6}  "
+            f"{spec.description}{default}"
+        )
+    return 0
+
+
+def _cmd_serve(args: argparse.Namespace) -> int:
+    try:
+        import uvicorn
+    except ImportError:
+        print(
+            'error: the API needs the extra dependencies: pip install "snipbg[api]"',
+            file=sys.stderr,
+        )
+        return 1
+
+    uvicorn.run(
+        "snipbg.server.app:app",
+        host=args.host,
+        port=args.port,
+        reload=args.reload,
+        workers=args.workers if not args.reload else None,
+    )
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    """Entry point. Returns a process exit code."""
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+    try:
+        return int(args.func(args))
+    except KeyboardInterrupt:
+        print("interrupted", file=sys.stderr)
+        return 130
+    except ValueError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

snipbg/download.py

@@ -0,0 +1,184 @@
+"""Model cache: where weights live, and how they get there.
+
+Only the standard library may be imported here: the Docker model stage installs
+no third-party packages and runs this module together with :mod:`models`.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import shutil
+import sys
+import urllib.request
+from collections.abc import Callable
+from pathlib import Path
+
+from .models import ModelSpec, get_spec
+
+_CHUNK = 1024 * 256
+_USER_AGENT = "snipbg (+https://github.com/asimmakhmudov/snipbg)"
+
+ProgressFn = Callable[[int, int], None]
+
+
+class DownloadError(RuntimeError):
+    """A model could not be downloaded, or arrived corrupted."""
+
+
+def cache_dir() -> Path:
+    """Directory holding downloaded weights.
+
+    ``SNIPBG_HOME`` wins; otherwise ``$XDG_CACHE_HOME/snipbg``, falling
+    back to ``~/.cache/snipbg``.
+    """
+    env = os.environ.get("SNIPBG_HOME")
+    if env:
+        return Path(env).expanduser()
+    xdg = os.environ.get("XDG_CACHE_HOME")
+    base = Path(xdg).expanduser() if xdg else Path.home() / ".cache"
+    return base / "snipbg"
+
+
+def model_path(spec: ModelSpec | str) -> Path:
+    """Path the given model's weights are cached at (may not exist yet)."""
+    if isinstance(spec, str):
+        spec = get_spec(spec)
+    return cache_dir() / spec.filename
+
+
+def sha256_file(path: Path) -> str:
+    """SHA-256 of a file, read in chunks."""
+    digest = hashlib.sha256()
+    with open(path, "rb") as fh:
+        for chunk in iter(lambda: fh.read(_CHUNK), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def is_cached(spec: ModelSpec | str) -> bool:
+    """True if the weights are already on disk with the expected size."""
+    if isinstance(spec, str):
+        spec = get_spec(spec)
+    path = model_path(spec)
+    if not path.is_file():
+        return False
+    if spec.size_bytes and path.stat().st_size != spec.size_bytes:
+        return False
+    return True
+
+
+def _stderr_progress(done: int, total: int) -> None:
+    if total <= 0:
+        sys.stderr.write(f"\r  {done / 1048576:.1f} MB")
+    else:
+        pct = 100.0 * done / total
+        sys.stderr.write(f"\r  {pct:5.1f}%  {done / 1048576:7.1f} / {total / 1048576:.1f} MB")
+    sys.stderr.flush()
+
+
+def download_model(
+    spec: ModelSpec | str,
+    *,
+    progress: ProgressFn | bool = False,
+    force: bool = False,
+) -> Path:
+    """Download a model's weights into the cache and return the path.
+
+    The body is streamed to a temporary file next to the destination, hashed
+    while it is written, and only renamed into place once the digest matches —
+    so a cached path is never a half-written file.
+
+    Raises:
+        DownloadError: on a transport error or a digest mismatch.
+    """
+    if isinstance(spec, str):
+        spec = get_spec(spec)
+
+    dest = model_path(spec)
+    if dest.is_file() and not force:
+        return dest
+
+    on_progress: ProgressFn | None
+    if progress is True:
+        on_progress = _stderr_progress
+    elif progress is False:
+        on_progress = None
+    else:
+        on_progress = progress
+
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    tmp = dest.with_suffix(dest.suffix + ".part")
+    digest = hashlib.sha256()
+    request = urllib.request.Request(spec.url, headers={"User-Agent": _USER_AGENT})
+
+    try:
+        with urllib.request.urlopen(request) as response:  # noqa: S310 - fixed https base
+            total = int(response.headers.get("Content-Length") or spec.size_bytes or 0)
+            done = 0
+            with open(tmp, "wb") as fh:
+                while True:
+                    chunk = response.read(_CHUNK)
+                    if not chunk:
+                        break
+                    fh.write(chunk)
+                    digest.update(chunk)
+                    done += len(chunk)
+                    if on_progress is not None:
+                        on_progress(done, total)
+    except OSError as exc:  # URLError and friends are OSError subclasses
+        tmp.unlink(missing_ok=True)
+        raise DownloadError(f"could not download {spec.name} from {spec.url}: {exc}") from exc
+    finally:
+        if on_progress is _stderr_progress:
+            sys.stderr.write("\n")
+
+    actual = digest.hexdigest()
+    if spec.sha256 and actual != spec.sha256:
+        tmp.unlink(missing_ok=True)
+        raise DownloadError(
+            f"{spec.name}: SHA-256 mismatch (expected {spec.sha256}, got {actual}). "
+            "The download was discarded."
+        )
+
+    os.replace(tmp, dest)
+    return dest
+
+
+def ensure_model(spec: ModelSpec | str, *, progress: ProgressFn | bool = False) -> Path:
+    """Return the cached path for a model, downloading it if it is missing."""
+    if isinstance(spec, str):
+        spec = get_spec(spec)
+    if is_cached(spec):
+        return model_path(spec)
+    return download_model(spec, progress=progress, force=True)
+
+
+def verify_model(spec: ModelSpec | str) -> bool:
+    """Hash a cached file and compare it with the expected digest.
+
+    Returns ``False`` if the file is missing or corrupt, ``True`` if it matches
+    (or if the registry carries no digest for it).
+    """
+    if isinstance(spec, str):
+        spec = get_spec(spec)
+    path = model_path(spec)
+    if not path.is_file():
+        return False
+    if not spec.sha256:
+        return True
+    return sha256_file(path) == spec.sha256
+
+
+def remove_cached(spec: ModelSpec | str) -> bool:
+    """Delete a cached model. True if a file was removed."""
+    path = model_path(spec)
+    if path.is_file():
+        path.unlink()
+        return True
+    return False
+
+
+def clear_cache() -> None:
+    """Delete the whole model cache directory."""
+    shutil.rmtree(cache_dir(), ignore_errors=True)

snipbg/images.py

@@ -0,0 +1,148 @@
+"""Getting images in and out: decoding inputs, encoding results."""
+
+from __future__ import annotations
+
+import os
+from io import BytesIO
+from pathlib import Path
+from typing import Literal
+
+import numpy as np
+from PIL import Image, ImageOps
+
+#: Anything :func:`load_image` accepts.
+ImageInput = bytes | bytearray | memoryview | str | os.PathLike | Image.Image | np.ndarray
+
+InputKind = Literal["bytes", "path", "pil", "numpy"]
+
+#: Formats that cannot carry an alpha channel, so a cutout has to be flattened.
+_OPAQUE_FORMATS = frozenset({"JPEG", "JPG", "BMP", "PCX", "PPM", "EPS"})
+
+_SUFFIX_FORMATS = {
+    ".png": "PNG",
+    ".jpg": "JPEG",
+    ".jpeg": "JPEG",
+    ".jpe": "JPEG",
+    ".webp": "WEBP",
+    ".tif": "TIFF",
+    ".tiff": "TIFF",
+    ".bmp": "BMP",
+    ".gif": "GIF",
+}
+
+
+def input_kind(data: ImageInput) -> InputKind:
+    """Classify an input so the caller's result can match its type."""
+    if isinstance(data, Image.Image):
+        return "pil"
+    if isinstance(data, np.ndarray):
+        return "numpy"
+    if isinstance(data, bytes | bytearray | memoryview):
+        return "bytes"
+    if isinstance(data, str | os.PathLike):
+        return "path"
+    raise TypeError(f"unsupported image input type: {type(data).__name__}")
+
+
+def _from_numpy(array: np.ndarray) -> Image.Image:
+    if array.ndim not in (2, 3):
+        raise ValueError(f"expected a 2-D or 3-D array, got shape {array.shape}")
+    if array.dtype != np.uint8:
+        working = array.astype(np.float32)
+        # Floating point arrays are taken as 0..1 unless they clearly are not.
+        if working.size and working.max() <= 1.0:
+            working = working * 255.0
+        array = np.clip(working, 0, 255).astype(np.uint8)
+    if array.ndim == 3 and array.shape[2] == 1:
+        array = array[:, :, 0]
+    if array.ndim == 3 and array.shape[2] not in (3, 4):
+        raise ValueError(f"expected 1, 3 or 4 channels, got {array.shape[2]}")
+    return Image.fromarray(array)
+
+
+def load_image(data: ImageInput) -> Image.Image:
+    """Decode any supported input into a PIL image with EXIF rotation applied.
+
+    Raises:
+        ValueError: if the bytes or file are not a readable image.
+        TypeError: if the input type is not supported.
+    """
+    kind = input_kind(data)
+
+    if kind == "pil":
+        assert isinstance(data, Image.Image)
+        image = data
+    elif kind == "numpy":
+        assert isinstance(data, np.ndarray)
+        image = _from_numpy(data)
+    else:
+        source: BytesIO | Path
+        if kind == "bytes":
+            source = BytesIO(bytes(data))  # type: ignore[arg-type]
+        else:
+            source = Path(os.fspath(data))  # type: ignore[arg-type]
+            if not source.is_file():
+                raise ValueError(f"no such image file: {source}")
+        try:
+            image = Image.open(source)
+            image.load()
+        except OSError as exc:
+            # PIL's message for a buffer embeds an object repr (and its memory
+            # address), which would end up in an HTTP error body.
+            if kind == "bytes":
+                raise ValueError(
+                    "could not decode an image from the given bytes: "
+                    "the data is not a supported image format"
+                ) from exc
+            raise ValueError(f"could not decode an image from {source}: {exc}") from exc
+
+    rotated = ImageOps.exif_transpose(image)
+    if rotated is None:
+        return image
+    # exif_transpose hands back a copy, which does not carry .format.
+    rotated.format = rotated.format or image.format
+    return rotated
+
+
+def format_for_path(path: str | os.PathLike, default: str = "PNG") -> str:
+    """Pick an output format from a filename suffix."""
+    return _SUFFIX_FORMATS.get(Path(os.fspath(path)).suffix.lower(), default)
+
+
+def flatten_if_opaque(image: Image.Image, image_format: str, color: str = "white") -> Image.Image:
+    """Composite away the alpha channel for formats that cannot store it."""
+    if image_format.upper() not in _OPAQUE_FORMATS or image.mode not in ("RGBA", "LA", "P"):
+        return image
+    rgba = image.convert("RGBA")
+    background = Image.new("RGBA", rgba.size, color)
+    return Image.alpha_composite(background, rgba).convert("RGB")
+
+
+def encode_image(
+    image: Image.Image,
+    image_format: str = "PNG",
+    *,
+    quality: int | None = None,
+) -> bytes:
+    """Serialise an image to bytes.
+
+    Alpha is flattened onto white for formats that cannot carry it, so asking
+    for a JPEG never fails on an RGBA cutout.
+    """
+    image_format = image_format.upper()
+    prepared = flatten_if_opaque(image, image_format)
+
+    options: dict[str, object] = {}
+    if quality is not None and image_format in ("JPEG", "WEBP"):
+        options["quality"] = int(quality)
+    if image_format == "PNG":
+        options["optimize"] = True
+
+    buffer = BytesIO()
+    prepared.save(buffer, format=image_format, **options)
+    return buffer.getvalue()
+
+
+def to_numpy(image: Image.Image) -> np.ndarray:
+    """RGBA/RGB image as a uint8 array."""
+    return np.asarray(image)