OpenAI互換APIサーバ＋VLMを用いて画像をキャプショニングする

image_caption_cli.py

from __future__ import annotations

import argparse
import base64
import json
import mimetypes
import sys
from pathlib import Path
from typing import Iterable
from urllib import error, parse, request


DEFAULT_PROMPT = "Please describe the image."
DEFAULT_SERVER_URL = "http://127.0.0.1:8080/v1"
SUPPORTED_EXTENSIONS = {".jpg", ".jpeg", ".png"}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate captions for JPEG/PNG images via an OpenAI-compatible local server."
    )
    parser.add_argument(
        "image_dir",
        type=Path,
        help="Directory that contains target images.",
    )
    parser.add_argument(
        "-p",
        "--prompt-file",
        type=Path,
        help="UTF-8 text file containing the prompt to send.",
    )
    parser.add_argument(
        "-s",
        "--server-url",
        default=DEFAULT_SERVER_URL,
        help=f"Base URL of the OpenAI-compatible server. Default: {DEFAULT_SERVER_URL}",
    )
    parser.add_argument(
        "-m",
        "--model",
        help="Model ID to use. If omitted, the script fetches the first model from /models.",
    )
    parser.add_argument(
        "-r",
        "--recursive",
        action="store_true",
        help="Recursively scan subdirectories for images.",
    )
    return parser.parse_args()


def load_prompt(prompt_file: Path | None) -> str:
    if prompt_file is None:
        return DEFAULT_PROMPT

    try:
        prompt = prompt_file.read_text(encoding="utf-8").strip()
    except OSError as exc:
        raise RuntimeError(f"Failed to read prompt file: {prompt_file}") from exc

    if not prompt:
        raise RuntimeError(f"Prompt file is empty: {prompt_file}")

    return prompt


def iter_image_files(root: Path, recursive: bool) -> Iterable[Path]:
    # Use generator so that processing and traversing a massive directory 
    # of thousands of images do not consume significant upfront memory.
    iterator = root.rglob("*") if recursive else root.iterdir()
    for path in iterator:
        if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS:
            yield path


def build_url(base_url: str, endpoint: str) -> str:
    return parse.urljoin(base_url.rstrip("/") + "/", endpoint.lstrip("/"))


def http_json(
    method: str,
    url: str,
    payload: dict | None = None,
) -> dict:
    # Use built-in urllib to avoid external dependencies like 'requests', 
    # keeping this script lightweight and completely standalone.
    headers = {"Accept": "application/json"}
    data = None
    if payload is not None:
        headers["Content-Type"] = "application/json"
        data = json.dumps(payload).encode("utf-8")

    req = request.Request(url=url, method=method, headers=headers, data=data)
    try:
        with request.urlopen(req) as response:
            body = response.read().decode("utf-8")
    except error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {exc.code} for {url}: {detail}") from exc
    except error.URLError as exc:
        raise RuntimeError(f"Failed to connect to server: {url}") from exc

    try:
        return json.loads(body)
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"Server did not return valid JSON: {url}") from exc


def resolve_model(server_url: str) -> str:
    # Automatically determine the model ID if not provided by the user.
    # Local servers often host a single model natively, so we can just grab the first available one.
    models_url = build_url(server_url, "/models")
    response = http_json("GET", models_url)
    models = response.get("data")
    if not isinstance(models, list) or not models:
        raise RuntimeError(f"No models returned from {models_url}")

    model_id = models[0].get("id")
    if not isinstance(model_id, str) or not model_id:
        raise RuntimeError(f"Invalid model list returned from {models_url}")

    return model_id


def image_to_data_url(image_path: Path) -> str:
    # Vision APIs require image data to be passed inline.
    # The image file is base64 encoded and formatted as a Data URI scheme.
    mime_type, _ = mimetypes.guess_type(image_path.name)
    if mime_type not in {"image/jpeg", "image/png"}:
        raise RuntimeError(f"Unsupported MIME type for {image_path}")

    encoded = base64.b64encode(image_path.read_bytes()).decode("ascii")
    return f"data:{mime_type};base64,{encoded}"


def extract_caption(response: dict) -> str:
    # Navigate the nested JSON structure mandated by OpenAI-compatible endpoints.
    try:
        content = response["choices"][0]["message"]["content"]
    except (KeyError, IndexError, TypeError) as exc:
        raise RuntimeError("Unexpected response format from chat/completions") from exc

    # Vision models may return 'content' as either a simple string or a list of 
    # message part objects. This handles cases where the text and an image 
    # (or multiple text segments) are conceptually mixed together.
    if isinstance(content, str):
        caption = content.strip()
    elif isinstance(content, list):
        text_parts: list[str] = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                text = item.get("text")
                if isinstance(text, str):
                    text_parts.append(text)
        caption = "\n".join(part.strip() for part in text_parts if part.strip())
    else:
        caption = ""

    if not caption:
        raise RuntimeError("Caption was empty in the server response")

    return caption


def generate_caption(server_url: str, model: str, prompt: str, image_path: Path) -> str:
    data_url = image_to_data_url(image_path)
    payload = {
        "model": model,
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": data_url}},
                ],
            }
        ],
    }
    response = http_json("POST", build_url(server_url, "/chat/completions"), payload)
    return extract_caption(response)


def save_caption(image_path: Path, caption: str) -> Path:
    # Save a text file beside the image using the same basename.
    # This acts as paired data, widely adopted for datasets and tagging flows.
    output_path = image_path.with_suffix(".txt")
    output_path.write_text(caption + "\n", encoding="utf-8")
    return output_path


def main() -> int:
    args = parse_args()
    image_dir = args.image_dir.resolve()

    if not image_dir.exists():
        print(f"Image directory does not exist: {image_dir}", file=sys.stderr)
        return 1
    if not image_dir.is_dir():
        print(f"Image directory is not a directory: {image_dir}", file=sys.stderr)
        return 1

    try:
        prompt = load_prompt(args.prompt_file)
        model = args.model or resolve_model(args.server_url)
        image_files = sorted(iter_image_files(image_dir, args.recursive))
        if not image_files:
            print("No JPEG or PNG files were found.", file=sys.stderr)
            return 1

        hostname = parse.urlparse(args.server_url).hostname or ""
        # Connecting to 'localhost' on Windows can sometimes invoke IPv6 resolution, 
        # which heavily delays request timeouts under certain configurations. 
        # Using the direct IPv4 '127.0.0.1' address prevents this stall.
        if "localhost" in hostname.lower():
            print(
                "Note: if requests are slow on Windows, try --server-url http://127.0.0.1:8080/v1",
                file=sys.stderr,
            )

        for image_path in image_files:
            print(f"Processing {image_path}...")
            caption = generate_caption(args.server_url, model, prompt, image_path)
            output_path = save_caption(image_path, caption)
            print(f"Saved caption: {output_path}")
    except RuntimeError as exc:
        print(str(exc), file=sys.stderr)
        return 1
    except OSError as exc:
        print(f"Filesystem error: {exc}", file=sys.stderr)
        return 1

    return 0


if __name__ == "__main__":
    raise SystemExit(main())

prompt_anima_for_character.md

Describe the image in English using a combination of tags and natural language, approximately 300 words total.

Overall Format

Use the following format. Do not include any text other than the description.

[Tag description], [Natural language description]

Tag Description Rules

Describe the image using Danbooru-style tags. All tags must be lowercase. Separate tags with , followed by a space.

Tag Sections and Order

[meta/safety tags], 1girl, kagasaka sonoha, orichara, [general tags]

Tags within each section can be in any order.

Meta Tags

Tags describing the image format and attributes. The following are examples; other tags may also be used.

highres, absurdres, anime screenshot, jpeg artifacts, official art, etc.

Safety Tags

Tags indicating the safety level of the image, particularly regarding sexual content. Assign exactly one of the following:

safe, sensitive, nsfw, explicit

1girl

The image should include exactly one girl, and this tag should be included.

Character

The character depicted in the image is Kagasaka Sonoha, and this tag should be included.

Series

The character Kagasaka Sonoha is from the series Orichara, and this tag should be included.

General Tags

Tags describing the content of the image. Examples: long hair, blue eyes, smile, bare shoulders, hair ribbon, bracelet, zettai ryouiki, classroom, cityscape, star (symbol), simple background, monochrome, 2koma, etc.

The first tag in the general tags section should be "solo" to indicate that there is only one character in the image. The remaining general tags should be chosen based on the content of the image, such as the character's appearance, clothing, pose, background, and other visible elements.

Reference Information:

Kagasaka Sonoha has short black hair with blue inner layers, red eyes. So tags such as "short hair, black hair, blue hair, colored inner hair, multicolored hair, red eyes" may be appropriate depending on the image.

Natural Language Description Rules

Describe the image in English, approximately 200 words.

Describing Characters

The image features Kagasaka Sonoha, so state her name first, then describe her appearance (e.g., "Digital artwork of Kagasaka Sonoha from Orichara, with short black hair that has blue inner layers, and red eyes...").

Image Description

Describe the content of the image in as much detail as possible, including characters, background, composition, and color palette. Avoid speculating about mood or emotions; describe only what is visible in the image.

prompt_anima_v1.md

Describe the image in English using a combination of tags and natural language, approximately 300 words total.

Overall Format

Use the following format. Do not include any text other than the description.

[Tag description], [Natural language description]

Tag Description Rules

Describe the image using Danbooru-style tags. All tags must be lowercase. Separate tags with , followed by a space.

Tag Sections and Order

[meta/safety tags] [1girl/1boy/1other etc] [character] [series] [general tags]

Tags within each section can be in any order.

Meta Tags

Tags describing the image format and attributes. The following are examples; other tags may also be used.

highres, absurdres, anime screenshot, jpeg artifacts, official art, etc.

Safety Tags

Tags indicating the safety level of the image, particularly regarding sexual content. Assign exactly one of the following:

safe, sensitive, nsfw, explicit

1girl/1boy/1other etc

Tags indicating the gender and number of people in the image. Examples: 1girl, 2boys, 3girls, 1other, etc.

Character

Tag the name of the character depicted in the image. Examples: remilia scarlet, marisa kirisame, etc.

Series

Tag the name of the series the character belongs to. Examples: touhou project, vocaloid, etc.

General Tags

Tags describing the content of the image. Examples: long hair, blue eyes, smile, bare shoulders, hair ribbon, bracelet, zettai ryouiki, classroom, cityscape, star (symbol), simple background, monochrome, 2koma, etc.

Natural Language Description Rules

Describe the image in English, approximately 200 words.

When People Are Present

If the character can be identified, state the character name first, then describe their appearance (e.g., "Digital artwork of Fern from Sousou no Frieren, with long purple hair and purple eyes, wearing a black coat over a white dress with puffy sleeves..."). If the character cannot be identified, describe their appearance directly (e.g., "A digital artwork of a young girl with long brown hair and green eyes, wearing a red dress with white frills...").

Image Description

Describe the content of the image in as much detail as possible, including characters, background, composition, and color palette. Avoid speculating about mood or emotions; describe only what is visible in the image.

llama.cppなどでサーバを起動し（--mmprojを指定し画像入力を有効にしておく）、以下のコマンドで起動（URLはIPアドレスでの指定を推奨します）。

python image_caption_cli.py path/to/image_dir --server-url http://127.0.0.1:8080/v1 --prompt-file prompt_anima.md

プロンプトはAnima用に調整済みですが、quality tagとartist tagは出ないようにしているのと、キャラクタやシリーズの精度も低いので、手作業での修正を推奨します。タグ部分はtaggerの使用もご検討ください。

キャラクタが特定人物の場合はタグやキャラクタ名が固定できるので、prompt_anima_for_character.mdをキャラクタ名、シリーズ、タグ、特徴などを修正してご利用ください。