vlm-toolkit 0.1.0__py3-none-any.whl

vlm_toolkit/bbox.py

@@ -0,0 +1,62 @@
+import json
+import re
+
+from PIL import Image, ImageDraw
+
+
+def parse_bboxes(output: str) -> list[dict]:
+    """Parse bbox_2d/label objects from VLM output, unwrapping markdown code blocks."""
+    # strip markdown code blocks if present
+    output = output.strip()
+    if output.startswith("```"):
+        match = re.search(r"```(?:json)?\s*\n(.*?)\n```", output, re.DOTALL)
+        if match:
+            output = match.group(1)
+
+    try:
+        bboxes = json.loads(output)
+        return bboxes if isinstance(bboxes, list) else []
+    except json.JSONDecodeError:
+        return []
+
+
+def scale_bbox(
+    bbox_coords: list[int], width: int, height: int
+) -> tuple[int, int, int, int]:
+    """Scale a bbox from 1000x1000 coords to pixels, adding a 15% buffer."""
+    x1, y1, x2, y2 = bbox_coords
+    # convert from 1000-based coordinates to actual pixel coordinates
+    scaled_x1 = x1 * width / 1000
+    scaled_y1 = y1 * height / 1000
+    scaled_x2 = x2 * width / 1000
+    scaled_y2 = y2 * height / 1000
+
+    # add 15% buffer in all directions
+    bbox_width = scaled_x2 - scaled_x1
+    bbox_height = scaled_y2 - scaled_y1
+
+    buffer_x = bbox_width * 0.15
+    buffer_y = bbox_height * 0.15
+
+    # expand bbox and clamp to image boundaries
+    scaled_x1 = max(0, int(scaled_x1 - buffer_x))
+    scaled_y1 = max(0, int(scaled_y1 - buffer_y))
+    scaled_x2 = min(width, int(scaled_x2 + buffer_x))
+    scaled_y2 = min(height, int(scaled_y2 + buffer_y))
+
+    return scaled_x1, scaled_y1, scaled_x2, scaled_y2
+
+
+def create_mask_from_bboxes(width: int, height: int, bboxes: list[dict]) -> Image.Image:
+    """Create a mask from bounding boxes."""
+    mask = Image.new("RGB", (width, height), color="black")
+    draw = ImageDraw.Draw(mask)
+
+    for bbox in bboxes:
+        if "bbox_2d" in bbox:
+            # scale bbox from 1000x1000 to actual image dimensions
+            x1, y1, x2, y2 = scale_bbox(bbox["bbox_2d"], width, height)
+            # draw white rectangle where logo is
+            draw.rectangle([x1, y1, x2, y2], fill="white")
+
+    return mask

vlm_toolkit/cmd/caption_images.py

@@ -0,0 +1,44 @@
+from importlib.resources import files
+from pathlib import Path
+
+import typer
+
+from vlm_toolkit.models import DEFAULT_MODEL, SUGGESTED_MODELS
+
+app = typer.Typer()
+
+PROMPT_FILE = files("vlm_toolkit") / "prompts" / "caption.txt"
+
+
+@app.command()
+def main(
+    folder: Path = typer.Argument(..., help="Folder containing images to caption"),
+    output: Path = typer.Option(None, help="Output JSON file"),
+    model: str = typer.Option(DEFAULT_MODEL, help="Hugging Face model name"),
+    batch_size: int = typer.Option(1, help="Number of images to process in parallel"),
+) -> None:
+    """Generate captions for all images in a folder using a VLM."""
+    # deferred so `vlm --help` doesn't pay for the torch import
+    from vlm_toolkit.vlm_process import vlm_process
+
+    prompt = PROMPT_FILE.read_text()
+
+    if output is None:
+        output = folder / "captions.json"
+
+    vlm_process(
+        folder=folder,
+        output=output,
+        prompt=prompt,
+        model=model,
+        batch_size=batch_size,
+    )
+
+
+main.__doc__ = (
+    f"Generate captions for all images in a folder using a VLM.\n\n{SUGGESTED_MODELS}"
+)
+
+
+if __name__ == "__main__":
+    app()

vlm_toolkit/cmd/detect_logo.py

@@ -0,0 +1,42 @@
+from importlib.resources import files
+from pathlib import Path
+
+import typer
+
+from vlm_toolkit.models import DEFAULT_MODEL, SUGGESTED_MODELS
+
+app = typer.Typer()
+
+PROMPT_FILE = files("vlm_toolkit") / "prompts" / "detect_logo.txt"
+
+
+@app.command()
+def main(
+    folder: Path = typer.Argument(..., help="Folder containing images to check"),
+    output: Path = typer.Option(None, help="Output JSON file"),
+    model: str = typer.Option(DEFAULT_MODEL, help="HF model"),
+    batch_size: int = typer.Option(8, help="Number of images to process in parallel"),
+) -> None:
+    """Detect logos with a VLM."""
+    # deferred so `vlm --help` doesn't pay for the torch import
+    from vlm_toolkit.vlm_process import vlm_process
+
+    prompt = PROMPT_FILE.read_text()
+
+    if output is None:
+        output = folder / "logo_bbox_output.json"
+
+    vlm_process(
+        folder=folder,
+        output=output,
+        prompt=prompt,
+        model=model,
+        batch_size=batch_size,
+    )
+
+
+main.__doc__ = f"Detect logos with a VLM.\n\n{SUGGESTED_MODELS}"
+
+
+if __name__ == "__main__":
+    app()

vlm_toolkit/cmd/main.py

@@ -0,0 +1,87 @@
+"""The `vlm` entry point, dispatching to the per-command Typer apps."""
+
+import importlib
+from typing import cast
+
+import click
+import typer
+
+# name -> (module under vlm_toolkit.cmd, one-line help shown by `vlm --help`)
+COMMANDS = {
+    "process": ("vlm_process", "Run a VLM over a folder of images locally"),
+    "server": ("vlm_server", "Run a folder through an OpenAI-compatible endpoint"),
+    "caption-images": ("caption_images", "Caption images with the captioning prompt"),
+    "detect-logo": ("detect_logo", "Find logos and watermarks with a VLM"),
+    "remove-logo": ("remove_logo", "Erase detected boxes by inpainting"),
+    "view-output": ("view_vlm_output", "Render a run as an HTML page"),
+    "upload-dataset": ("upload_dataset", "Publish a run to the Hugging Face Hub"),
+}
+
+# top-level module -> the extra that provides it, for a readable error when a
+# command is run against a base install
+EXTRAS = {
+    "accelerate": "local",
+    "kernels": "local",
+    "xgrammar": "local",
+    "torch": "local",
+    "torchvision": "local",
+    "transformers": "local",
+    "cv2": "logo",
+    "numpy": "logo",
+    "datasets": "hub",
+}
+
+
+class LazyGroup(click.Group):
+    """Import a subcommand's module only when that subcommand actually runs.
+
+    Keeps `vlm --help` and the light commands usable without the optional heavy
+    extras (torch, opencv, datasets) installed.
+    """
+
+    def list_commands(self, ctx: click.Context) -> list[str]:
+        return list(COMMANDS)
+
+    def get_command(self, ctx: click.Context, name: str) -> click.Command | None:
+        if name not in COMMANDS:
+            return None
+        module = importlib.import_module(f"vlm_toolkit.cmd.{COMMANDS[name][0]}")
+        command = typer.main.get_command(module.app)
+        command.short_help = COMMANDS[name][1]
+        # typer vendors its own click, so its commands aren't click.Command
+        # subclasses; click only ever duck-types what it gets back from here
+        return cast(click.Command, command)
+
+    def invoke(self, ctx: click.Context) -> object:
+        # typer's Exit and Abort are unrelated to click's, so without this a
+        # plain click group lets them escape as unhandled exceptions - which
+        # turns `--help` and every `raise typer.Exit(1)` into a traceback-ish 1
+        try:
+            return super().invoke(ctx)
+        except typer.Exit as e:
+            raise click.exceptions.Exit(e.exit_code) from e
+        except typer.Abort as e:
+            raise click.exceptions.Abort() from e
+        except ModuleNotFoundError as e:
+            extra = EXTRAS.get(e.name or "")
+            if extra is None:
+                raise
+            raise click.ClickException(
+                f"`vlm {ctx.invoked_subcommand}` needs the '{extra}' extra.\n"
+                f"Install it with: pip install 'vlm-toolkit[{extra}]'"
+            ) from e
+
+    def format_commands(self, ctx: click.Context, formatter) -> None:  # type: ignore[no-untyped-def]
+        # use the static help text, so listing commands imports nothing
+        with formatter.section("Commands"):
+            formatter.write_dl([(name, help) for name, (_, help) in COMMANDS.items()])
+
+
+@click.group(cls=LazyGroup, context_settings={"help_option_names": ["-h", "--help"]})
+@click.version_option(package_name="vlm-toolkit")
+def cli() -> None:
+    """Batch image processing with vision language models."""
+
+
+if __name__ == "__main__":
+    cli()

vlm_toolkit/cmd/remove_logo.py

@@ -0,0 +1,155 @@
+import json
+import os
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from enum import Enum
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import typer
+from PIL import Image
+from tqdm import tqdm
+
+from vlm_toolkit.bbox import parse_bboxes
+
+if TYPE_CHECKING:
+    from vlm_toolkit.logo import LogoRemover
+
+app = typer.Typer()
+
+
+class RemovalMethod(str, Enum):
+    """Logo removal methods."""
+
+    GRAY = "gray"
+    OPENCV = "opencv"
+
+
+def process_single_image(
+    image_path: Path, bboxes: list, output_path: Path, method: RemovalMethod
+) -> tuple[bool, str, Path | str]:
+    """Process a single image (used for parallel processing)."""
+    # deferred so `vlm --help` works without the logo extra installed
+    from vlm_toolkit.logo import GrayLogoRemover, OpenCVLogoRemover
+
+    try:
+        remover: "LogoRemover"
+        if method == RemovalMethod.OPENCV:
+            remover = OpenCVLogoRemover()
+        else:
+            remover = GrayLogoRemover()
+
+        image = Image.open(image_path).convert("RGB")
+        result = remover.remove(image, bboxes)
+        result.save(output_path)
+
+        return (True, image_path.name, output_path)
+    except Exception as e:
+        return (False, image_path.name, str(e))
+
+
+@app.command()
+def main(
+    bboxes: Path = typer.Argument(
+        ..., help="JSON file with bounding boxes (output from vlm-process)"
+    ),
+    output_dir: Path = typer.Argument(None, help="Output directory for cleaned images"),
+    method: RemovalMethod = typer.Option(
+        RemovalMethod.OPENCV, help="Inpainting method"
+    ),
+) -> None:
+    """Remove logos from images using inpainting.
+
+    Expects a JSON file where each entry has an 'output' field containing
+    bounding boxes in the format:
+    [{"bbox_2d": [x1, y1, x2, y2], "label": "logo"}]
+
+    Methods:
+    - gray: Replace with gray rectangles (instant, simple)
+    - opencv: Fast, lightweight inpainting using OpenCV (recommended for simple cases)
+    """
+
+    if not bboxes.exists():
+        print(f"Error: {bboxes} does not exist", flush=True)
+        raise typer.Exit(1)
+
+    # default output directory to same location as input JSON
+    if output_dir is None:
+        output_dir = bboxes.parent / f"{bboxes.stem}_cleaned"
+
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    # load detections from JSON lines file
+    detections = []
+    json_dir = bboxes.parent
+    with open(bboxes) as f:
+        for line in f:
+            if line.strip():
+                detections.append(json.loads(line))
+
+    print(f"Loaded {len(detections)} detections", flush=True)
+
+    # parse bboxes and filter out images without any
+    images_to_clean = []
+    for detection in detections:
+        bboxes_list = parse_bboxes(detection["output"])
+        if bboxes_list:
+            detection["bboxes"] = bboxes_list
+            # resolve relative path
+            rel_path = Path(detection["file_name"])
+            detection["abs_path"] = (json_dir / rel_path).resolve()
+            images_to_clean.append(detection)
+
+    print(f"Found {len(images_to_clean)} images with logos to remove", flush=True)
+
+    # prepare tasks for parallel processing
+    tasks = []
+    for detection in images_to_clean:
+        image_path = detection["abs_path"]
+        bbox_list = detection["bboxes"]
+        output_path = output_dir / f"{image_path.name}"
+        tasks.append((image_path, bbox_list, output_path, method))
+
+    # process images in parallel
+    num_workers = os.cpu_count()
+    print(
+        f"Processing {len(images_to_clean)} images with {num_workers} workers...",
+        flush=True,
+    )
+
+    results = []
+    errors = []
+    with ProcessPoolExecutor(max_workers=num_workers) as executor:
+        futures = {
+            executor.submit(
+                process_single_image, img_path, bbox_list, out_path, method
+            ): (
+                img_path,
+                len(bbox_list),
+            )
+            for img_path, bbox_list, out_path, method in tasks
+        }
+
+        # process results as they complete with progress bar
+        for future in tqdm(
+            as_completed(futures), total=len(futures), desc="Processing"
+        ):
+            img_path, num_bboxes = futures[future]
+            success, name, result_or_error = future.result()
+
+            if success:
+                results.append(result_or_error)
+            else:
+                errors.append((name, result_or_error))
+
+    # report any errors
+    if errors:
+        print(f"\nErrors encountered:", flush=True)
+        for name, error in errors:
+            print(f"  ✗ {name}: {error}", flush=True)
+
+    print(f"Processed {len(results)}/{len(images_to_clean)} images", flush=True)
+    print(f"Cleaned images saved to: {output_dir}", flush=True)
+
+
+if __name__ == "__main__":
+    app()

vlm_toolkit/cmd/upload_dataset.py

@@ -0,0 +1,53 @@
+from pathlib import Path
+
+import typer
+
+app = typer.Typer()
+
+
+@app.command()
+def main(
+    results_file: Path = typer.Argument(..., help="JSONL output from vlm-process"),
+    repo_id: str = typer.Argument(
+        ..., help="Target dataset repo, e.g. user/my-dataset"
+    ),
+    private: bool = typer.Option(True, help="Create the dataset repo as private"),
+    split: str = typer.Option("train", help="Split name to upload as"),
+    token: str = typer.Option(
+        None, help="Hugging Face token (defaults to HF_TOKEN or a cached login)"
+    ),
+    card: Path = typer.Option(
+        None, help="Markdown file to attach as the dataset card (README.md)"
+    ),
+) -> None:
+    """Upload a vlm-process run to the Hugging Face Hub as an image dataset.
+
+    Images are embedded in the dataset, not referenced by path, so the result is
+    self-contained and works in the Hub's dataset viewer. --card attaches a
+    README, preserving the dataset_info metadata the upload generates.
+    """
+    # deferred so `vlm --help` works without the hub extra installed
+    from vlm_toolkit.hf_dataset import upload
+
+    if not results_file.exists():
+        print(f"Error: {results_file} does not exist", flush=True)
+        raise typer.Exit(1)
+
+    # check before uploading, so a typo doesn't surface after a long transfer
+    if card is not None and not card.exists():
+        print(f"Error: {card} does not exist", flush=True)
+        raise typer.Exit(1)
+
+    url = upload(
+        results_file,
+        repo_id,
+        private=private,
+        split=split,
+        token=token,
+        card_file=card,
+    )
+    print(f"Uploaded to {url}", flush=True)
+
+
+if __name__ == "__main__":
+    app()

vlm_toolkit/cmd/view_vlm_output.py

@@ -0,0 +1,135 @@
+"""Generate HTML visualization of VLM output."""
+
+import functools
+import http.server
+import os
+import socket
+import socketserver
+import webbrowser
+from html import escape
+from importlib.resources import files
+from pathlib import Path
+from string import Template
+from typing import Any
+from urllib.parse import quote
+
+import typer
+
+from vlm_toolkit.results import load_results
+
+app = typer.Typer()
+
+
+class _ViewerHandler(http.server.SimpleHTTPRequestHandler):
+    """Serve the generated page at / and images from the served root."""
+
+    page = b""
+
+    def do_GET(self) -> None:
+        if self.path in ("/", "/index.html"):
+            self.send_response(200)
+            self.send_header("Content-Type", "text/html; charset=utf-8")
+            self.send_header("Content-Length", str(len(self.page)))
+            self.end_headers()
+            self.wfile.write(self.page)
+            return
+        super().do_GET()
+
+    def log_message(self, fmt: str, *args: Any) -> None:
+        pass
+
+
+def _render(items: list[dict], srcs: list[str]) -> str:
+    """Render the viewer page, using srcs as the img src for each item."""
+    templates = files("vlm_toolkit") / "templates"
+    page = Template((templates / "viewer.html").read_text())
+    item_template = Template((templates / "item.html").read_text())
+
+    # model output is untrusted text, so escape everything going into the markup
+    blocks = [
+        item_template.substitute(
+            src=escape(src, quote=True),
+            name=escape(item["abs_path"].name),
+            output=escape(item["output"]),
+        )
+        for item, src in zip(items, srcs)
+    ]
+    return page.substitute(items="".join(blocks))
+
+
+def _web_root(items: list[dict], json_dir: Path) -> tuple[Path, list[str]]:
+    """Pick a server root containing every image, and URLs relative to it."""
+    paths = [item["abs_path"] for item in items]
+    root = Path(os.path.commonpath([json_dir.resolve(), *paths]))
+    return root, [quote(str(p.relative_to(root))) for p in paths]
+
+
+def _lan_ip() -> str:
+    """Best guess at this machine's outward-facing address, for the printed URL."""
+    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+        try:
+            s.connect(("8.8.8.8", 80))  # no packets sent, just picks the route
+            return s.getsockname()[0]
+        except OSError:
+            return socket.gethostname()
+
+
+def _serve(items: list[dict], input_json: Path, host: str, port: int) -> None:
+    """Serve the viewer, rooted at the common ancestor of the JSON and its images."""
+    root, srcs = _web_root(items, input_json.parent)
+
+    _ViewerHandler.page = _render(items, srcs).encode()
+    handler = functools.partial(_ViewerHandler, directory=str(root))
+
+    socketserver.TCPServer.allow_reuse_address = True
+    with socketserver.TCPServer((host, port), handler) as httpd:
+        local = f"http://{'127.0.0.1' if host in ('0.0.0.0', '') else host}:{port}"
+        print(f"Serving {root} (Ctrl-C to stop)", flush=True)
+        print(f"  local:   {local}", flush=True)
+        if host in ("0.0.0.0", ""):
+            print(f"  network: http://{_lan_ip()}:{port}", flush=True)
+            print("  reachable by anyone who can route to this host", flush=True)
+        webbrowser.open(local)
+        try:
+            httpd.serve_forever()
+        except KeyboardInterrupt:
+            print("\nStopped", flush=True)
+
+
+@app.command()
+def main(
+    input_json: Path = typer.Argument(..., help="JSON file with VLM output"),
+    serve: bool = typer.Option(
+        False, "--serve", help="Serve over HTTP instead of writing an HTML file"
+    ),
+    host: str = typer.Option(
+        "127.0.0.1", help="Address to bind; use 0.0.0.0 to accept external connections"
+    ),
+    port: int = typer.Option(8000, help="Port to serve on"),
+) -> None:
+    """Generate HTML visualization of VLM output.
+
+    Writes an HTML file next to the input and opens it, or with --serve hosts it
+    on localhost instead, which is what you want when the images live on a
+    remote machine. --host 0.0.0.0 accepts external connections, which serves the
+    image directory to the network with no authentication.
+    """
+    if not input_json.exists():
+        print(f"Error: {input_json} does not exist", flush=True)
+        raise typer.Exit(1)
+
+    items = load_results(input_json)
+
+    if serve:
+        _serve(items, input_json, host, port)
+        return
+
+    html = _render(items, [str(item["abs_path"]) for item in items])
+    output_path = Path(str(input_json) + ".html")
+    output_path.write_text(html)
+    print(f"Generated {output_path}", flush=True)
+    webbrowser.open(f"file://{output_path.absolute()}")
+
+
+if __name__ == "__main__":
+    app()

vlm_toolkit/cmd/vlm_process.py

@@ -0,0 +1,51 @@
+from pathlib import Path
+
+import typer
+
+from vlm_toolkit.models import DEFAULT_MODEL, DEFAULT_PROMPT, SUGGESTED_MODELS
+
+app = typer.Typer()
+
+
+@app.command()
+def main(
+    folder: Path = typer.Argument(..., help="Folder containing images to process"),
+    output: Path = typer.Option(None, help="Output JSON file"),
+    prompt: str = typer.Option(DEFAULT_PROMPT, help="Prompt for the VLM"),
+    prompt_file: Path = typer.Option(
+        None, help="File containing prompt (overrides --prompt)"
+    ),
+    model: str = typer.Option(DEFAULT_MODEL, help="HF model"),
+    batch_size: int = typer.Option(8, help="Number of images to process in parallel"),
+    max_dim: int = typer.Option(
+        None, help="Maximum dimension for image resizing (default: 1024)"
+    ),
+    schema: Path = typer.Option(
+        None, help="JSON schema file for constrained decoding (uses xgrammar)"
+    ),
+) -> None:
+    """Process images using a Vision Language Model."""
+    # deferred so `vlm --help` doesn't pay for the torch import
+    from vlm_toolkit.vlm_process import vlm_process
+
+    if prompt_file is not None:
+        prompt = prompt_file.read_text().strip()
+
+    schema_str = schema.read_text() if schema is not None else None
+
+    vlm_process(
+        folder=folder,
+        output=output,
+        prompt=prompt,
+        model=model,
+        batch_size=batch_size,
+        max_dim=max_dim,
+        schema=schema_str,
+    )
+
+
+main.__doc__ = f"Process images using a Vision Language Model.\n\n{SUGGESTED_MODELS}"
+
+
+if __name__ == "__main__":
+    app()

vlm_toolkit/cmd/vlm_server.py

@@ -0,0 +1,48 @@
+from pathlib import Path
+
+import typer
+
+from vlm_toolkit.vlm_server import (
+    DEFAULT_BASE_URL,
+    DEFAULT_MODEL,
+    DEFAULT_PROMPT,
+    vlm_server_process,
+)
+
+app = typer.Typer()
+
+
+@app.command()
+def main(
+    folder: Path = typer.Argument(..., help="Folder containing images to process"),
+    output: Path = typer.Option(None, help="Output JSONL file"),
+    prompt: str = typer.Option(DEFAULT_PROMPT, help="Prompt for the VLM"),
+    prompt_file: Path = typer.Option(
+        None, help="File containing prompt (overrides --prompt)"
+    ),
+    base_url: str = typer.Option(DEFAULT_BASE_URL, help="Server base URL"),
+    model: str = typer.Option(DEFAULT_MODEL, help="Model name to send in requests"),
+    max_dim: int = typer.Option(
+        None, help="Maximum dimension for image resizing (default: 1024)"
+    ),
+    max_tokens: int = typer.Option(512, help="Max tokens to generate"),
+    concurrency: int = typer.Option(8, help="Number of concurrent requests"),
+) -> None:
+    """Process images via an OpenAI-compatible vision endpoint."""
+    if prompt_file is not None:
+        prompt = prompt_file.read_text().strip()
+
+    vlm_server_process(
+        folder=folder,
+        output=output,
+        prompt=prompt,
+        base_url=base_url,
+        model=model,
+        max_dim=max_dim,
+        max_tokens=max_tokens,
+        concurrency=concurrency,
+    )
+
+
+if __name__ == "__main__":
+    app()