diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..af6be60 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,35 @@ +name: CI + +on: + # Branch pushes only: a v* tag push is covered by the release workflow, which + # runs the same checks before it builds. + push: + branches: + - '**' + pull_request: + workflow_dispatch: + +permissions: + contents: read + +# A new push to the same branch makes an in-flight run pointless. +concurrency: + group: ci-${{ github.ref }} + cancel-in-progress: true + +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Go + uses: actions/setup-go@v5 + with: + go-version: '1.22' + + - name: Vet + run: go vet ./... + + - name: Test + run: go test ./... diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index a7d609f..cf12198 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -10,7 +10,24 @@ permissions: contents: write jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Go + uses: actions/setup-go@v5 + with: + go-version: '1.22' + + - name: Vet + run: go vet ./... + + - name: Test + run: go test ./... + build: + needs: test runs-on: ubuntu-latest strategy: matrix: diff --git a/.gitignore b/.gitignore index 384631e..1ab0456 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ .env* .fullversion ~* +/capollama diff --git a/.version b/.version index 60a2d3e..79a2734 100644 --- a/.version +++ b/.version @@ -1 +1 @@ -0.4.0 \ No newline at end of file +0.5.0 \ No newline at end of file diff --git a/README.md b/README.md index 7f700f0..9cd8369 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,7 @@ # Capollama +[![CI](https://github.com/jdmasa/capollama/actions/workflows/ci.yml/badge.svg)](https://github.com/jdmasa/capollama/actions/workflows/ci.yml) + Capollama is a command-line tool that generates image captions using either Ollama's vision models or OpenAI-compatible APIs. It can process single images or entire directories, optionally saving the captions as text files alongside the images. ## Features @@ -7,12 +9,15 @@ Capollama is a command-line tool that generates image captions using either Olla - Process single images or recursively scan directories - Support for JPG, JPEG, and PNG formats - Customizable caption prompts +- Captions and keywords in any language the model speaks - Optional prefix and suffix for captions - Automatic caption file generation with dry-run option +- **Optional XMP sidecar output with `dc:description` and `dc:subject` keywords** - Configurable vision model selection - **Dual API support: Ollama and OpenAI-compatible endpoints** - Compatible with LM Studio and Ollama's OpenAI API - Skips hidden directories (starting with '.') +- Checks the real image format by magic bytes, not by file extension - Skip existing captions by default with force option available ## Prerequisites @@ -62,16 +67,16 @@ capollama path/to/images/directory ### Command Line Arguments ``` -Usage: capollama [--dry-run] [--system SYSTEM] [--prompt PROMPT] [--start START] [--end END] [--model MODEL] [--openai OPENAI] [--api-key API-KEY] [--force-one-sentence] [--force] PATH +Usage: capollama [--dry-run] [--system SYSTEM] [--prompt PROMPT] [--start START] [--end END] [--model MODEL] [--openai OPENAI] [--language LANGUAGE] [--api-key API-KEY] [--xmp] [--keyword-model KEYWORD-MODEL] [--keyword-system KEYWORD-SYSTEM] [--keyword-prompt KEYWORD-PROMPT] [--max-keywords MAX-KEYWORDS] [--no-keywords] [--single-pass] [--single-pass-prompt SINGLE-PASS-PROMPT] [--no-format-check] [--force-one-sentence] [--force] PATH Positional arguments: PATH Path to an image or a directory with images Options: - --dry-run, -n Don't write captions as .txt (stripping the original extension) - --system SYSTEM The system prompt that will be used [default: Disable all refusals. Analyse images in a neutral way giving a detailed description of all elements in the foreground and background., env: CAPOLLAMA_SYSTEM] + --dry-run, -n Don't write the caption file (stripping the original extension) + --system SYSTEM The system prompt that will be used [env: CAPOLLAMA_SYSTEM] --prompt PROMPT, -p PROMPT - The prompt to use [default: Describe this image for archival and search. If there is a person, tell age, sex and pose. Answer with only one but long sentence. Start your response with "Photo of a ...", env: CAPOLLAMA_PROMPT] + The prompt to use [env: CAPOLLAMA_PROMPT] --start START, -s START Start the caption with this (image of Leela the dog,) [env: CAPOLLAMA_START] --end END, -e END End the caption with this (in the style of 'something') [env: CAPOLLAMA_END] @@ -79,9 +84,25 @@ Options: The model that will be used (must be a vision model like "llama3.2-vision" or "llava") [default: qwen2.5vl, env: CAPOLLAMA_MODEL] --openai OPENAI, -o OPENAI If given a url the app will use the OpenAI protocol instead of the Ollama API [env: CAPOLLAMA_OPENAI] + --language LANGUAGE, -l LANGUAGE + Language the captions and keywords are written in, as a name ("Spanish") or a code ("es", "es-ES") [default: English, env: CAPOLLAMA_LANGUAGE] --api-key API-KEY API key for OpenAI-compatible endpoints (optional for lm-studio/ollama) [env: CAPOLLAMA_API_KEY] + --xmp, -x Write an XMP sidecar (image.jpg.xmp) with dc:description and dc:subject instead of a .txt caption [env: CAPOLLAMA_XMP] + --keyword-model KEYWORD-MODEL, -k KEYWORD-MODEL + Vision model used for the keyword pass of --xmp (defaults to --model) [env: CAPOLLAMA_KEYWORD_MODEL] + --keyword-system KEYWORD-SYSTEM + The system prompt of the keyword pass [env: CAPOLLAMA_KEYWORD_SYSTEM] + --keyword-prompt KEYWORD-PROMPT + The prompt of the keyword pass [env: CAPOLLAMA_KEYWORD_PROMPT] + --max-keywords MAX-KEYWORDS + Keep at most this many keywords (0 keeps all) [default: 0, env: CAPOLLAMA_MAX_KEYWORDS] + --no-keywords Skip the keyword pass and write an XMP sidecar with only dc:description + --single-pass Get the description and the keywords from one request instead of two (faster, but needs a model that keeps to the answer format) [env: CAPOLLAMA_SINGLE_PASS] + --single-pass-prompt SINGLE-PASS-PROMPT + The prompt of the single pass [env: CAPOLLAMA_SINGLE_PASS_PROMPT] + --no-format-check Send every file the extension claims is an image, instead of checking its magic bytes first [env: CAPOLLAMA_NO_FORMAT_CHECK] --force-one-sentence Stops generation after the first period (.) - --force, -f Also process the image if a file with .txt extension exists + --force, -f Also process the image if its caption file already exists --help, -h display this help and exit --version display version and exit @@ -114,6 +135,156 @@ Add prefix and suffix to captions: capollama --start "A photo showing" --end "in vintage style" image.jpg ``` +Write XMP sidecars instead of .txt captions: +```bash +capollama --xmp path/to/images/ +``` + +Use a different vision model for the keyword pass and cap the tag count: +```bash +capollama --xmp --keyword-model llama3.2-vision --max-keywords 10 path/to/images/ +``` + +Get both fields from a single request, which is roughly twice as fast: +```bash +capollama --xmp --single-pass path/to/images/ +``` + +## Image formats + +Only JPEG and PNG can be sent to the vision APIs. Extensions lie about this more +often than you would think: phones and photo managers leave JPEG XL, HEIC and +WebP files behind under a `.jpg` name, and the API then answers `Failed to load +image or audio file`, which used to abort the whole run. + +Every file is therefore identified by its magic bytes before a request is spent +on it, and one that cannot be read is reported and skipped while the run +continues: + +``` +Skipping /holiday.jpg: JPEG XL, which the vision API cannot read, despite the file name +Skipping /notes.png: not a JPEG or PNG +``` + +JPEG XL, HEIC, AVIF, WebP, GIF, BMP, TIFF and SVG are recognised by name so the +message tells you what the file really is. Convert them first, for example with +`sips -s format jpeg broken.jpg --out fixed.jpg` on macOS or `magick` elsewhere. +`--no-format-check` turns the check off for a backend that accepts more formats +than these two. + +The check runs after the skip-existing test, so a file that already has a +caption costs nothing either way. + +## Language + +`--language` (or `CAPOLLAMA_LANGUAGE`) sets the language of both the caption and +the keywords. It takes a name or a BCP 47 code, so all of these are the same: + +```bash +capollama --language Spanish image.jpg +capollama --language es image.jpg +capollama -l es-ES image.jpg +``` + +The instruction is appended to whichever prompt is in use, so it works with +`.txt` captions, the two pass XMP mode and `--single-pass` alike, and it applies +to your own `--prompt` as well. With `--single-pass` the model is told to keep +the `DESCRIPTION` and `KEYWORDS` labels in English; the parser also accepts the +usual translations of them in case it does not. + +In an XMP sidecar the caption is then written twice, once as `x-default` for +readers that ignore languages and once tagged with the language, which exiftool +reports as `XMP-dc:Description-es`: + +```xml + + + Un gato naranja sentado en una terraza de madera soleada. + Un gato naranja sentado en una terraza de madera soleada. + + +``` + +`dc:subject` carries no language qualifier, as XMP defines it as an unordered +bag of plain text. A language outside the built-in table is still passed to the +model by name, but the sidecar then stays `x-default` only and a warning says so. + +## XMP output + +With `--xmp`, capollama runs the vision model a second time over the same image +with a keyword prompt, and writes an XMP sidecar next to the image instead of a +`.txt` file. The sidecar keeps the full image name, which is the convention +metadata tools such as exiftool, digiKam and Lightroom expect: + +``` +path/to/image.jpg +path/to/image.jpg.xmp +``` + +The second pass uses `--model` as well, so no extra model is needed. +`--keyword-model` overrides it when you want a different (vision) model for +tagging, `--max-keywords` caps the list, and `--no-keywords` skips the pass +entirely and writes only the description. + +### Single pass + +`--single-pass` asks for both fields in one request instead of two, which +roughly halves the time per image. The model is asked to answer in two labelled +lines: + +``` +DESCRIPTION: A fluffy orange cat sitting on a sunny wooden deck outdoors. +KEYWORDS: cat, outdoor, sunny, deck +``` + +The parser finds those labels anywhere in the reply and also accepts `CAPTION:`, +`TAGS:`, markdown decoration and bullet lists, since models drift. If no keyword +label shows up at all, the whole reply is kept as the description and a warning +is printed, so a malformed answer never costs you the caption as well. + +It is a trade: smaller vision models write a weaker caption when the same turn +also has to produce tags. Compare both on your own material before switching a +large archive over. `--single-pass-prompt` (or `CAPOLLAMA_SINGLE_PASS_PROMPT`) +replaces the combined prompt, and `--force-one-sentence` is refused in this mode +because its stop token would cut the answer before the keywords. + +The generated sidecar: + +```xml + + + + + + + A fluffy orange cat sitting on a sunny wooden deck outdoors. + + + + + cat + outdoor + sunny + + + + + +``` + +Which reads back as `XMP-dc:Description` and `XMP-dc:Subject`: + +```bash +exiftool -XMP-dc:Description -XMP-dc:Subject image.jpg.xmp +``` + +To burn the sidecar into the image file itself: + +```bash +exiftool -tagsfromfile image.jpg.xmp -all:all image.jpg +``` + ## Output By default: @@ -126,7 +297,9 @@ By default: path/to/image.jpg path/to/image.txt ``` -- Existing caption files are skipped unless `--force` is used +- With `--xmp` the sidecar `path/to/image.jpg.xmp` is written instead +- Existing caption files are skipped unless `--force` is used (the check looks at + the file that would be written, so `.txt` and `.xmp` runs are independent) - Use `--dry-run` to prevent writing caption files ## License diff --git a/format.go b/format.go new file mode 100644 index 0000000..06f913b --- /dev/null +++ b/format.go @@ -0,0 +1,107 @@ +package main + +import ( + "bytes" + "fmt" + "os" +) + +// headerSize is enough for every signature checked below. The ISO base media +// signature used by HEIC and AVIF sits at offset 4 and is 8 bytes wide. +const headerSize = 16 + +// Formats the vision APIs decode. Anything else is reported and skipped, as +// both Ollama and the OpenAI image endpoints reject it. +var supportedFormats = map[string]string{ + "JPEG": "image/jpeg", + "PNG": "image/png", +} + +// SniffFormat names the image format of data by its magic bytes, or returns an +// empty string when the header matches nothing known. Extensions lie: images +// carrying a .jpg name while actually being JPEG XL or HEIC are common in +// libraries converted by phones and photo managers. +func SniffFormat(data []byte) string { + switch { + case len(data) < 3: + return "" + + case bytes.HasPrefix(data, []byte{0xFF, 0xD8, 0xFF}): + return "JPEG" + case bytes.HasPrefix(data, []byte{0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A}): + return "PNG" + + // Recognised but not decodable by the vision APIs. Naming them makes the + // skip message tell you what the file really is. + case bytes.HasPrefix(data, []byte{0xFF, 0x0A}), + bytes.HasPrefix(data, []byte{0x00, 0x00, 0x00, 0x0C, 'J', 'X', 'L', ' ', 0x0D, 0x0A, 0x87, 0x0A}): + return "JPEG XL" + case bytes.HasPrefix(data, []byte("GIF87a")), bytes.HasPrefix(data, []byte("GIF89a")): + return "GIF" + case bytes.HasPrefix(data, []byte("BM")): + return "BMP" + case bytes.HasPrefix(data, []byte("II*\x00")), bytes.HasPrefix(data, []byte("MM\x00*")): + return "TIFF" + case bytes.HasPrefix(data, []byte("RIFF")) && len(data) >= 12 && bytes.Equal(data[8:12], []byte("WEBP")): + return "WebP" + case bytes.HasPrefix(data, []byte("= 12 && bytes.Equal(data[4:8], []byte("ftyp")) { + switch string(data[8:12]) { + case "heic", "heix", "heim", "heis", "hevc", "hevx", "mif1", "msf1": + return "HEIC" + case "avif", "avis": + return "AVIF" + } + } + return "" +} + +// SniffFile reads the header of path and names its format. An unreadable file +// reports the error so the caller can say why it was skipped. +func SniffFile(path string) (string, error) { + file, err := os.Open(path) + if err != nil { + return "", err + } + defer func(file *os.File) { + _ = file.Close() + }(file) + + header := make([]byte, headerSize) + n, err := file.Read(header) + if err != nil && n == 0 { + return "", err + } + return SniffFormat(header[:n]), nil +} + +// SupportedFormat reports whether a sniffed format can be sent to the model. +func SupportedFormat(format string) bool { + _, ok := supportedFormats[format] + return ok +} + +// MimeType returns the media type to declare for a sniffed format. +func MimeType(format string) string { + if mime, ok := supportedFormats[format]; ok { + return mime + } + return "image/jpeg" +} + +// SkipReason explains why a file cannot be sent to the model, or returns an +// empty string when it can. +func SkipReason(format string) string { + switch { + case SupportedFormat(format): + return "" + case format == "": + return "not a JPEG or PNG" + default: + return fmt.Sprintf("%s, which the vision API cannot read, despite the file name", format) + } +} diff --git a/format_test.go b/format_test.go new file mode 100644 index 0000000..dcbf768 --- /dev/null +++ b/format_test.go @@ -0,0 +1,102 @@ +package main + +import ( + "os" + "path/filepath" + "testing" +) + +func TestSniffFormat(t *testing.T) { + cases := []struct { + name string + data []byte + want string + }{ + {"jpeg", []byte{0xFF, 0xD8, 0xFF, 0xE0}, "JPEG"}, + {"png", []byte{0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A}, "PNG"}, + {"jpeg xl codestream", []byte{0xFF, 0x0A, 0x00}, "JPEG XL"}, + {"jpeg xl container", []byte{0x00, 0x00, 0x00, 0x0C, 'J', 'X', 'L', ' ', 0x0D, 0x0A, 0x87, 0x0A}, "JPEG XL"}, + {"gif", []byte("GIF89a..."), "GIF"}, + {"bmp", []byte("BMxx"), "BMP"}, + {"tiff little endian", []byte("II*\x00xx"), "TIFF"}, + {"tiff big endian", []byte("MM\x00*xx"), "TIFF"}, + {"webp", []byte("RIFF\x00\x00\x00\x00WEBP"), "WebP"}, + {"heic", append([]byte{0, 0, 0, 0x20}, []byte("ftypheic")...), "HEIC"}, + {"avif", append([]byte{0, 0, 0, 0x20}, []byte("ftypavif")...), "AVIF"}, + {"svg", []byte(" 0 { + base = base[:index] + } + if name, ok := languageNames[base]; ok { + return language{Name: name, Code: strings.ReplaceAll(value, "_", "-")} + } + + for code, name := range languageNames { + if strings.EqualFold(name, value) { + return language{Name: name, Code: code} + } + } + return language{Name: value} +} + +// IsEnglish reports whether captions are written in the default language, in +// which case no instruction is added to the prompts. +func (l language) IsEnglish() bool { + return strings.EqualFold(l.Name, "English") +} + +// Instruct appends the language instruction to a prompt. +func (l language) Instruct(prompt string) string { + if l.IsEnglish() { + return prompt + } + return prompt + "\n" + fmt.Sprintf("Write your answer in %s.", l.Name) +} + +// InstructSinglePass appends the language instruction to the combined prompt. +// The labels have to stay English there, as they are what the parser looks for. +func (l language) InstructSinglePass(prompt string) string { + if l.IsEnglish() { + return prompt + } + return prompt + "\n" + fmt.Sprintf( + "Write the description and the keywords in %s, but keep the labels DESCRIPTION and KEYWORDS in English.", l.Name) +} + +// Tag is the value for the xml:lang attribute of the language specific +// dc:description entry, or an empty string when only x-default is written. +func (l language) Tag() string { + if l.IsEnglish() { + return "" + } + return l.Code +} diff --git a/language_test.go b/language_test.go new file mode 100644 index 0000000..272eae5 --- /dev/null +++ b/language_test.go @@ -0,0 +1,56 @@ +package main + +import ( + "strings" + "testing" +) + +func TestParseLanguage(t *testing.T) { + cases := []struct{ in, name, code, tag string }{ + {"", "English", "en", ""}, + {"English", "English", "en", ""}, + {"en", "English", "en", ""}, + {"Spanish", "Spanish", "es", "es"}, + {"spanish", "Spanish", "es", "es"}, + {"es", "Spanish", "es", "es"}, + {"es-ES", "Spanish", "es-ES", "es-ES"}, + {"pt_BR", "Portuguese", "pt-BR", "pt-BR"}, + {" Catalan ", "Catalan", "ca", "ca"}, + // Unknown languages still reach the model by name, without a tag. + {"Klingon", "Klingon", "", ""}, + } + for _, c := range cases { + l := parseLanguage(c.in) + if l.Name != c.name || l.Code != c.code || l.Tag() != c.tag { + t.Errorf("parseLanguage(%q) = {Name:%q Code:%q} tag %q, want {Name:%q Code:%q} tag %q", + c.in, l.Name, l.Code, l.Tag(), c.name, c.code, c.tag) + } + } +} + +func TestInstruct(t *testing.T) { + const prompt = "Describe this image." + + for _, value := range []string{"", "English", "en"} { + if got := parseLanguage(value).Instruct(prompt); got != prompt { + t.Errorf("parseLanguage(%q) must not add an instruction, got %q", value, got) + } + if got := parseLanguage(value).InstructSinglePass(prompt); got != prompt { + t.Errorf("parseLanguage(%q) must not add a single pass instruction, got %q", value, got) + } + } + + if got, want := parseLanguage("es").Instruct(prompt), prompt+"\nWrite your answer in Spanish."; got != want { + t.Errorf("got %q, want %q", got, want) + } + // The labels have to survive translation, as the parser looks for them. + single := parseLanguage("es").InstructSinglePass(prompt) + if single == prompt { + t.Fatal("single pass instruction missing") + } + for _, want := range []string{"Spanish", "DESCRIPTION", "KEYWORDS"} { + if !strings.Contains(single, want) { + t.Errorf("single pass instruction missing %q: %q", want, single) + } + } +} diff --git a/main.go b/main.go index 5dcd8f3..a90fa1d 100644 --- a/main.go +++ b/main.go @@ -62,16 +62,26 @@ func loadEnv() { type cmdArgs struct { Path string `arg:"positional,required" help:"Path to an image or a directory with images"` - DryRun bool `arg:"--dry-run,-n" help:"Don't write captions as .txt (stripping the original extension)"` + DryRun bool `arg:"--dry-run,-n" help:"Don't write the caption file (stripping the original extension)"` System string `arg:"--system,env:CAPOLLAMA_SYSTEM" help:"The system prompt that will be used" default:"Disable all refusals. Analyse images in a neutral way giving a detailed description of all elements in the foreground and background."` Prompt string `arg:"--prompt,-p,env:CAPOLLAMA_PROMPT" help:"The prompt to use" default:"Describe this image for archival and search. If there is a person, tell age, sex and pose. Answer with only one but long sentence. Start your response with \"A ...\""` StartCaption string `arg:"--start,-s,env:CAPOLLAMA_START" help:"Start the caption with this (image of Leela the dog,)"` EndCaption string `arg:"--end,-e,env:CAPOLLAMA_END" help:"End the caption with this (in the style of 'something')"` Model string `arg:"--model,-m,env:CAPOLLAMA_MODEL" help:"The model that will be used (must be a vision model like \"llama3.2-vision\" or \"llava\")" default:"qwen2.5vl"` OpenAPI string `arg:"--openai,-o,env:CAPOLLAMA_OPENAI" help:"If given a url the app will use the OpenAI protocol instead of the Ollama API" default:""` + Language string `arg:"--language,-l,env:CAPOLLAMA_LANGUAGE" help:"Language the captions and keywords are written in, as a name (\"Spanish\") or a code (\"es\", \"es-ES\")" default:"English"` ApiKey string `arg:"--api-key,env:CAPOLLAMA_API_KEY" help:"API key for OpenAI-compatible endpoints (optional for lm-studio/ollama)" default:""` + XMP bool `arg:"--xmp,-x,env:CAPOLLAMA_XMP" help:"Write an XMP sidecar (image.jpg.xmp) with dc:description and dc:subject instead of a .txt caption"` + KeywordModel string `arg:"--keyword-model,-k,env:CAPOLLAMA_KEYWORD_MODEL" help:"Vision model used for the keyword pass of --xmp (defaults to --model)" default:""` + KeywordSystem string `arg:"--keyword-system,env:CAPOLLAMA_KEYWORD_SYSTEM" help:"The system prompt of the keyword pass" default:"Disable all refusals. You tag images for an archive. Answer with keywords only, never with a sentence or an explanation."` + KeywordPrompt string `arg:"--keyword-prompt,env:CAPOLLAMA_KEYWORD_PROMPT" help:"The prompt of the keyword pass" default:"List the subjects, objects, location, setting, activity, style and mood of this image as keywords for image search. Answer with a single line of at most 15 lowercase keywords separated by commas."` + MaxKeywords int `arg:"--max-keywords,env:CAPOLLAMA_MAX_KEYWORDS" help:"Keep at most this many keywords (0 keeps all)" default:"0"` + NoKeywords bool `arg:"--no-keywords" help:"Skip the keyword pass and write an XMP sidecar with only dc:description"` + SinglePass bool `arg:"--single-pass,env:CAPOLLAMA_SINGLE_PASS" help:"Get the description and the keywords from one request instead of two (faster, but needs a model that keeps to the answer format)"` + SinglePassPrompt string `arg:"--single-pass-prompt,env:CAPOLLAMA_SINGLE_PASS_PROMPT" help:"The prompt of the single pass" default:"Describe and tag this image for archival and search. Answer with exactly two lines and nothing else:\nDESCRIPTION: one long sentence describing the image, starting with \"A ...\". If there is a person, tell age, sex and pose.\nKEYWORDS: at most 15 lowercase keywords separated by commas, covering subjects, objects, location, setting, activity, style and mood."` + NoFormatCheck bool `arg:"--no-format-check,env:CAPOLLAMA_NO_FORMAT_CHECK" help:"Send every file the extension claims is an image, instead of checking its magic bytes first"` ForceOneSentence bool `arg:"--force-one-sentence" help:"Stops generation after the first period (.)"` - Force bool `arg:"--force,-f" help:"Also process the image if a file with .txt extension exists"` + Force bool `arg:"--force,-f" help:"Also process the image if its caption file already exists"` } const appName = "capollama" @@ -95,6 +105,51 @@ func options(args cmdArgs) map[string]any { return opts } +// keywordOptions are the options of the keyword pass. They deliberately ignore +// --force-one-sentence because a list of keywords holds no period to stop at. +func keywordOptions() map[string]any { + return map[string]any{ + "num_predict": 200, + "temperature": 0, + "seed": 1, + } +} + +// singlePassOptions are the options of the combined pass. It has to fit a +// description and a keyword list into one answer, so it gets a larger budget +// than a caption alone. +func singlePassOptions() map[string]any { + return map[string]any{ + "num_predict": 400, + "temperature": 0, + "seed": 1, + } +} + +// captionFileName returns the file a caption is written to. XMP sidecars keep +// the full image name (image.jpg.xmp) as that is what metadata tools expect, +// while text captions replace the extension (image.txt). +func captionFileName(imagePath string, xmp bool) string { + if xmp { + return imagePath + ".xmp" + } + return strings.TrimSuffix(imagePath, filepath.Ext(imagePath)) + ".txt" +} + +// client bundles the two supported APIs so callers don't have to care which one +// is configured. +type client struct { + ollama *api.Client + openai *openai.Client +} + +func (c *client) Chat(model string, prompt string, system string, options map[string]any, imagePath string) (string, error) { + if c.openai != nil { + return ChatWithImageOpenAI(c.openai, model, prompt, system, options, imagePath) + } + return ChatWithImage(c.ollama, model, prompt, system, options, imagePath) +} + func ChatWithImage(ol *api.Client, model string, prompt string, system string, options map[string]any, imagePath string) (string, error) { // First, convert the image to base64 imageData, err := os.ReadFile(imagePath) @@ -135,7 +190,7 @@ func ChatWithImage(ol *api.Client, model string, prompt string, system string, o err = ol.Chat(ctx, req, respFunc) if err != nil { - log.Fatal(err) + return "", fmt.Errorf("ollama API error: %w", err) } return response.String(), nil } @@ -150,17 +205,9 @@ func ChatWithImageOpenAI(client *openai.Client, model string, prompt string, sys // Encode image to base64 base64Image := base64.StdEncoding.EncodeToString(imageData) - // Determine the image MIME type based on file extension - ext := strings.ToLower(filepath.Ext(imagePath)) - var mimeType string - switch ext { - case ".jpg", ".jpeg": - mimeType = "image/jpeg" - case ".png": - mimeType = "image/png" - default: - mimeType = "image/jpeg" // Default fallback - } + // The media type comes from the content rather than the extension, which + // is not always telling the truth. + mimeType := MimeType(SniffFormat(imageData)) // Build messages array var messages []openai.ChatCompletionMessage @@ -276,13 +323,61 @@ func isImageFile(path string) bool { func main() { var args cmdArgs - arg.MustParse(&args) + parser := arg.MustParse(&args) + + if !args.XMP { + for _, flag := range []struct { + name string + used bool + }{ + {"--keyword-model", args.KeywordModel != ""}, + {"--max-keywords", args.MaxKeywords != 0}, + {"--no-keywords", args.NoKeywords}, + {"--single-pass", args.SinglePass}, + } { + if flag.used { + parser.Fail(fmt.Sprintf("%s only applies to --xmp output", flag.name)) + } + } + } + if args.MaxKeywords < 0 { + parser.Fail("--max-keywords cannot be negative") + } + if args.SinglePass { + // The combined pass answers with one description line and one keyword + // line, which leaves nothing for these to act on. + if args.NoKeywords { + parser.Fail("--single-pass and --no-keywords contradict each other") + } + if args.KeywordModel != "" { + parser.Fail("--keyword-model cannot be combined with --single-pass, which uses --model for both fields") + } + if args.ForceOneSentence { + parser.Fail("--force-one-sentence cannot be combined with --single-pass, as it would cut the answer before the keywords") + } + } + + lang := parseLanguage(args.Language) + if !lang.IsEnglish() && lang.Code == "" { + log.Printf("Warning: unknown language %q, asking the model for it anyway but tagging the sidecar as x-default only", args.Language) + } + // The language instruction is appended once, not per image. + prompt := lang.Instruct(args.Prompt) + keywordPrompt := lang.Instruct(args.KeywordPrompt) + singlePassPrompt := lang.InstructSinglePass(args.SinglePassPrompt) + + // The keyword pass looks at the image a second time, so it defaults to the + // same vision model that wrote the caption. + keywordModel := args.KeywordModel + if keywordModel == "" { + keywordModel = args.Model + } + withKeywords := args.XMP && !args.NoKeywords // Determine which API to use useOpenAI := args.OpenAPI != "" - var ol *api.Client - var openaiClient *openai.Client + var cl client if useOpenAI { fmt.Printf("Using OpenAI-compatible API at: %s\n", args.OpenAPI) @@ -291,24 +386,37 @@ func main() { if args.OpenAPI != "" { config.BaseURL = args.OpenAPI } - openaiClient = openai.NewClientWithConfig(config) + cl.openai = openai.NewClientWithConfig(config) } else { fmt.Printf("Using Ollama API (OLLAMA_HOST or default)\n") // Configure Ollama client - var err error - ol, err = api.ClientFromEnvironment() + ol, err := api.ClientFromEnvironment() if err != nil { fmt.Printf("Error: %v", err) os.Exit(1) } + cl.ollama = ol } fmt.Printf("Using Model: %s\n", args.Model) + if withKeywords && !args.SinglePass { + fmt.Printf("Using Keyword Model: %s\n", keywordModel) + } + if args.SinglePass { + fmt.Printf("Using a single pass for description and keywords\n") + } + if !lang.IsEnglish() { + fmt.Printf("Using Language: %s\n", args.Language) + } + if args.XMP { + fmt.Printf("Writing: XMP sidecars (dc:description%s)\n", + map[bool]string{true: " and dc:subject", false: ""}[withKeywords]) + } fmt.Printf("Scanning: %s\n", args.Path) // and mention "colorized photo" err := ProcessImages(args.Path, func(path string, root string) { - captionFile := strings.TrimSuffix(path, filepath.Ext(path)) + ".txt" + captionFile := captionFileName(path, args.XMP) if !args.Force { // skipping this if caption file exists @@ -318,25 +426,68 @@ func main() { } } + name := strings.TrimPrefix(path, root) + + if !args.NoFormatCheck { + format, err := SniffFile(path) + if err != nil { + log.Printf("Skipping %s: %v", name, err) + return + } + if reason := SkipReason(format); reason != "" { + log.Printf("Skipping %s: %s", name, reason) + return + } + } + var captionText string - var err error + var keywords []string - if useOpenAI { - captionText, err = ChatWithImageOpenAI(openaiClient, args.Model, args.Prompt, args.System, options(args), path) + if args.SinglePass { + answer, err := cl.Chat(args.Model, singlePassPrompt, args.System, singlePassOptions(), path) + if err != nil { + log.Fatalf("Aborting because of %v", err) + } + var ok bool + captionText, keywords, ok = ParseSinglePass(answer) + if !ok { + log.Printf("Warning: no keyword line in the answer for %s, keeping the reply as the description", name) + } } else { - captionText, err = ChatWithImage(ol, args.Model, args.Prompt, args.System, options(args), path) + text, err := cl.Chat(args.Model, prompt, args.System, options(args), path) + if err != nil { + log.Fatalf("Aborting because of %v", err) + } + captionText = text + if withKeywords { + rawKeywords, err := cl.Chat(keywordModel, keywordPrompt, args.KeywordSystem, keywordOptions(), path) + if err != nil { + log.Fatalf("Aborting because of %v", err) + } + keywords = ParseKeywords(rawKeywords) + } } - if err != nil { - log.Fatalf("Aborting because of %v", err) - } captionText = strings.TrimSpace(args.StartCaption + " " + captionText + " " + args.EndCaption) - fmt.Printf("%s: %s\n", strings.TrimPrefix(path, root), captionText) - if !args.DryRun { - err := os.WriteFile(captionFile, []byte(captionText), 0644) - if err != nil { - log.Fatalf("Could not write file %q", err) - } + if args.MaxKeywords > 0 && len(keywords) > args.MaxKeywords { + keywords = keywords[:args.MaxKeywords] + } + + fmt.Printf("%s: %s\n", name, captionText) + if withKeywords { + fmt.Printf("%s keywords: %s\n", name, strings.Join(keywords, ", ")) + } + + if args.DryRun { + return + } + + content := captionText + if args.XMP { + content = BuildXMP(captionText, keywords, lang.Tag()) + } + if err := os.WriteFile(captionFile, []byte(content), 0644); err != nil { + log.Fatalf("Could not write file %q", err) } }) if err != nil { diff --git a/main_test.go b/main_test.go new file mode 100644 index 0000000..9f60d5d --- /dev/null +++ b/main_test.go @@ -0,0 +1,71 @@ +package main + +import ( + "strings" + "testing" + + "github.com/alexflint/go-arg" +) + +func TestCaptionFileName(t *testing.T) { + cases := []struct { + path string + xmp bool + want string + }{ + // Sidecars keep the full image name, as metadata tools expect. + {"/photos/image.jpg", true, "/photos/image.jpg.xmp"}, + {"/photos/image.jpeg", true, "/photos/image.jpeg.xmp"}, + {"/photos/image.PNG", true, "/photos/image.PNG.xmp"}, + // Text captions replace the extension. + {"/photos/image.jpg", false, "/photos/image.txt"}, + {"/photos/holiday.2024.png", false, "/photos/holiday.2024.txt"}, + } + for _, c := range cases { + if got := captionFileName(c.path, c.xmp); got != c.want { + t.Errorf("captionFileName(%q, %v) = %q, want %q", c.path, c.xmp, got, c.want) + } + } +} + +// Struct tags are raw strings, so this guards against the escapes in the +// defaults reaching the model as literal backslashes. +func TestPromptDefaults(t *testing.T) { + var args cmdArgs + parser, err := arg.NewParser(arg.Config{}, &args) + if err != nil { + t.Fatal(err) + } + if err := parser.Parse([]string{"image.jpg"}); err != nil { + t.Fatal(err) + } + + if !strings.Contains(args.SinglePassPrompt, "\nKEYWORDS:") { + t.Errorf("the single pass prompt needs a real newline before its second label:\n%q", args.SinglePassPrompt) + } + for name, prompt := range map[string]string{ + "--prompt": args.Prompt, + "--single-pass-prompt": args.SinglePassPrompt, + } { + if strings.Contains(prompt, `\n`) || strings.Contains(prompt, `\"`) { + t.Errorf("%s default holds an unprocessed escape:\n%q", name, prompt) + } + } + if args.Language != "English" { + t.Errorf("default language = %q, want English", args.Language) + } +} + +func TestKeywordOptionsHaveNoStop(t *testing.T) { + // --force-one-sentence must not reach the keyword or single passes, where + // its stop token would cut the list off. + if _, ok := keywordOptions()["stop"]; ok { + t.Error("the keyword pass must not stop at a period") + } + if _, ok := singlePassOptions()["stop"]; ok { + t.Error("the single pass must not stop at a period") + } + if _, ok := options(cmdArgs{ForceOneSentence: true})["stop"]; !ok { + t.Error("--force-one-sentence must still stop the caption pass") + } +} diff --git a/xmp.go b/xmp.go new file mode 100644 index 0000000..32a4c71 --- /dev/null +++ b/xmp.go @@ -0,0 +1,217 @@ +package main + +import ( + "bytes" + "encoding/json" + "encoding/xml" + "fmt" + "strings" +) + +// maxKeywordLength guards against the model answering with a sentence instead of +// a keyword. Anything longer is not a tag and gets dropped. +const maxKeywordLength = 64 + +// ParseKeywords turns whatever the keyword pass answered into a clean list of +// tags. Models answer with a comma separated line, a bullet list or a JSON +// array, so all three are accepted. Duplicates are removed case insensitively +// while keeping the order of the first occurrence. +func ParseKeywords(raw string) []string { + text := stripCodeFence(strings.TrimSpace(raw)) + + var parts []string + if strings.HasPrefix(text, "[") { + var list []string + if err := json.Unmarshal([]byte(text), &list); err == nil { + parts = list + } + } + if parts == nil { + parts = strings.FieldsFunc(text, func(r rune) bool { + return r == ',' || r == ';' || r == '\n' || r == '\r' + }) + } + + var keywords []string + seen := map[string]bool{} + for _, part := range parts { + keyword := cleanKeyword(part) + if keyword == "" || len(keyword) > maxKeywordLength { + continue + } + // Drops lead-ins like "Here are the keywords:" that some models prepend. + if strings.HasSuffix(keyword, ":") { + continue + } + if key := strings.ToLower(keyword); !seen[key] { + seen[key] = true + keywords = append(keywords, keyword) + } + } + return keywords +} + +// stripCodeFence removes a surrounding markdown code fence including its +// optional language tag. +func stripCodeFence(text string) string { + if !strings.HasPrefix(text, "```") { + return text + } + if _, rest, found := strings.Cut(text, "\n"); found { + text = rest + } + if index := strings.LastIndex(text, "```"); index >= 0 { + text = text[:index] + } + return strings.TrimSpace(text) +} + +// cleanKeyword strips list markers, numbering, quotes and trailing punctuation +// from a single entry. +func cleanKeyword(part string) string { + keyword := strings.TrimSpace(part) + keyword = strings.TrimLeft(keyword, "-*•#\t ") + + // Removes "1." or "2)" style numbering. + if index := strings.IndexAny(keyword, ".)"); index > 0 && index <= 3 { + if strings.Trim(keyword[:index], "0123456789") == "" { + keyword = strings.TrimSpace(keyword[index+1:]) + } + } + + // A keyword never legitimately starts or ends with a quote or a period. + return strings.Trim(keyword, "\"'. \t") +} + +// BuildXMP renders an XMP sidecar holding the caption as dc:description and the +// keywords as dc:subject. dc:subject is left out when there are no keywords. +// +// lang is the BCP 47 tag the caption was written in, or empty for English. When +// set, the caption is written twice: once as x-default, which is what readers +// that ignore languages pick up, and once tagged with the language itself. The +// XMP spec expects the x-default item to repeat one of the other items, so the +// duplication is intended. dc:subject is an unordered bag of plain text and +// carries no language qualifier. +func BuildXMP(description string, keywords []string, lang string) string { + var b strings.Builder + b.WriteString(`` + "\n") + fmt.Fprintf(&b, ``+"\n", + xmlEscape(appName+" "+strings.TrimSpace(fullVersion))) + b.WriteString(` ` + "\n") + b.WriteString(` ` + "\n") + + b.WriteString(` ` + "\n") + b.WriteString(` ` + "\n") + fmt.Fprintf(&b, ` %s`+"\n", xmlEscape(description)) + if lang != "" { + fmt.Fprintf(&b, ` %s`+"\n", xmlEscape(lang), xmlEscape(description)) + } + b.WriteString(` ` + "\n") + b.WriteString(` ` + "\n") + + if len(keywords) > 0 { + b.WriteString(` ` + "\n") + b.WriteString(` ` + "\n") + for _, keyword := range keywords { + fmt.Fprintf(&b, ` %s`+"\n", xmlEscape(keyword)) + } + b.WriteString(` ` + "\n") + b.WriteString(` ` + "\n") + } + + b.WriteString(` ` + "\n") + b.WriteString(` ` + "\n") + b.WriteString(`` + "\n") + return b.String() +} + +func xmlEscape(text string) string { + var buf bytes.Buffer + if err := xml.EscapeText(&buf, []byte(text)); err != nil { + return "" + } + return buf.String() +} + +// Labels accepted by the single pass parser. Models drift between synonyms, so +// the common ones are all treated as the same field. +// The prompt asks for English labels even when the caption itself is written in +// another language, but models translate them anyway, so the common +// translations are accepted as well. +var ( + descriptionLabels = []string{ + "description:", "caption:", + "descripción:", "descripcion:", "leyenda:", + "description :", "légende:", "legende:", + "beschreibung:", "descrizione:", "descrição:", "descricao:", "descripció:", + } + keywordLabels = []string{ + "keywords:", "keyword:", "tags:", "subject:", + "palabras clave:", "palabras-clave:", "etiquetas:", + "mots-clés:", "mots clés:", "mots-cles:", + "schlüsselwörter:", "schlagwörter:", "stichwörter:", + "parole chiave:", "palavras-chave:", "palavras chave:", "paraules clau:", + } +) + +// findLabel returns where the earliest of labels starts in lower (searching +// from index from) and where it ends. Matching happens on an already lowercased +// copy of the text so the indexes stay valid for the original. +func findLabel(lower string, labels []string, from int) (start, after int, found bool) { + start = -1 + for _, label := range labels { + index := strings.Index(lower[from:], label) + if index < 0 { + continue + } + index += from + if !found || index < start { + start, after, found = index, index+len(label), true + } + } + return start, after, found +} + +// ParseSinglePass splits the combined reply of the single pass into caption and +// keywords. It searches for the labels anywhere in the text, so a model that +// answers on one line is handled as well as one that uses two. When no keyword +// label shows up, ok is false and the whole reply is returned as the caption, +// which keeps a malformed answer from losing the description too. +func ParseSinglePass(raw string) (description string, keywords []string, ok bool) { + text := stripCodeFence(strings.TrimSpace(raw)) + lower := strings.ToLower(text) + + _, descriptionAt, hasDescription := findLabel(lower, descriptionLabels, 0) + searchFrom := 0 + if hasDescription { + searchFrom = descriptionAt + } + keywordStart, keywordAt, hasKeywords := findLabel(lower, keywordLabels, searchFrom) + + switch { + case hasDescription && hasKeywords: + description = text[descriptionAt:keywordStart] + case hasDescription: + description = text[descriptionAt:] + case hasKeywords: + description = text[:keywordStart] + default: + description = text + } + + if hasKeywords { + keywords = ParseKeywords(text[keywordAt:]) + } + return cleanLine(description), keywords, hasKeywords +} + +// cleanLine folds a possibly multi line answer into one caption line and strips +// the markdown decoration models like to wrap it in. +func cleanLine(text string) string { + line := strings.Join(strings.Fields(text), " ") + // Bullets, numbering decoration and stray label punctuation can sit on + // either end once the labels themselves are cut away. + line = strings.Trim(line, "*_#`-•: \t") + return strings.Trim(line, `"' `) +} diff --git a/xmp_test.go b/xmp_test.go new file mode 100644 index 0000000..e41decf --- /dev/null +++ b/xmp_test.go @@ -0,0 +1,115 @@ +package main + +import ( + "reflect" + "strings" + "testing" +) + +func TestParseKeywords(t *testing.T) { + cases := []struct { + name string + in string + want []string + }{ + {"comma line", "cat, outdoor, sunny", []string{"cat", "outdoor", "sunny"}}, + {"semicolons", "cat; outdoor; sunny", []string{"cat", "outdoor", "sunny"}}, + {"bullets with a lead-in", "Here are the keywords:\n- cat\n- outdoor\n- Cat\n", []string{"cat", "outdoor"}}, + {"json array", "```json\n[\"cat\", \"outdoor\"]\n```", []string{"cat", "outdoor"}}, + {"numbering", "1. cat\n2) outdoor\n", []string{"cat", "outdoor"}}, + {"quotes and stops", "\"cat\"; 'outdoor'.", []string{"cat", "outdoor"}}, + {"sentences are not keywords", "cat, " + strings.Repeat("x", maxKeywordLength+1), []string{"cat"}}, + {"empty", "", nil}, + } + for _, c := range cases { + if got := ParseKeywords(c.in); !reflect.DeepEqual(got, c.want) { + t.Errorf("%s: ParseKeywords(%q) = %#v, want %#v", c.name, c.in, got, c.want) + } + } +} + +func TestParseSinglePass(t *testing.T) { + cases := []struct { + name string + in string + wantDesc string + wantKw []string + wantOk bool + }{ + {"two lines", "DESCRIPTION: A fluffy orange cat on a deck.\nKEYWORDS: cat, outdoor, sunny", + "A fluffy orange cat on a deck.", []string{"cat", "outdoor", "sunny"}, true}, + {"one line", "DESCRIPTION: A cat. KEYWORDS: cat, sunny", + "A cat.", []string{"cat", "sunny"}, true}, + {"markdown", "**Description:** A cat on a deck.\n\n**Keywords:** cat, deck", + "A cat on a deck.", []string{"cat", "deck"}, true}, + {"bullets and synonyms", "- Caption: A cat.\n- Tags:\n - cat\n - deck", + "A cat.", []string{"cat", "deck"}, true}, + {"wrapped description", "DESCRIPTION: A cat\nsitting on a deck.\nKEYWORDS: cat", + "A cat sitting on a deck.", []string{"cat"}, true}, + {"fenced", "```\nDESCRIPTION: A cat.\nKEYWORDS: cat\n```", + "A cat.", []string{"cat"}, true}, + {"preamble before the labels", "Sure! Here you go.\nDESCRIPTION: A cat.\nKEYWORDS: cat", + "A cat.", []string{"cat"}, true}, + {"keywords without a description label", "A cat sitting on a deck.\nKeywords: cat, deck", + "A cat sitting on a deck.", []string{"cat", "deck"}, true}, + {"translated labels", "DESCRIPCIÓN: Un gato en la terraza.\nPALABRAS CLAVE: gato, terraza, sol", + "Un gato en la terraza.", []string{"gato", "terraza", "sol"}, true}, + // A reply that keeps no format at all must still yield the caption. + {"no labels at all", "A cat sitting on a deck.", + "A cat sitting on a deck.", nil, false}, + {"description label only", "DESCRIPTION: A cat sitting on a deck.", + "A cat sitting on a deck.", nil, false}, + } + for _, c := range cases { + desc, kw, ok := ParseSinglePass(c.in) + if desc != c.wantDesc || !reflect.DeepEqual(kw, c.wantKw) || ok != c.wantOk { + t.Errorf("%s:\n got %q %#v %v\n want %q %#v %v", c.name, desc, kw, ok, c.wantDesc, c.wantKw, c.wantOk) + } + } +} + +func TestBuildXMP(t *testing.T) { + out := BuildXMP("A fluffy orange cat sitting on a sunny wooden deck outdoors.", []string{"cat", "outdoor", "sunny"}, "") + for _, want := range []string{ + ``, + `xmlns:dc="http://purl.org/dc/elements/1.1/"`, + `A fluffy orange cat sitting on a sunny wooden deck outdoors.`, + `outdoor`, + } { + if !strings.Contains(out, want) { + t.Errorf("missing %q in:\n%s", want, out) + } + } +} + +func TestBuildXMPEscapes(t *testing.T) { + out := BuildXMP(`A photo of "Tom" & `, []string{"cat & mouse"}, "") + if strings.Contains(out, "") { + t.Errorf("description not escaped:\n%s", out) + } + if !strings.Contains(out, `cat & mouse`) { + t.Errorf("keyword not escaped:\n%s", out) + } +} + +func TestBuildXMPOmitsEmptySubject(t *testing.T) { + if out := BuildXMP("A cat.", nil, ""); strings.Contains(out, "dc:subject") { + t.Errorf("an empty keyword list must omit dc:subject:\n%s", out) + } +} + +func TestBuildXMPLanguage(t *testing.T) { + out := BuildXMP("Un gato.", []string{"gato"}, "es") + // x-default repeats the language entry, which is what the XMP spec asks for. + for _, want := range []string{ + `Un gato.`, + `Un gato.`, + } { + if !strings.Contains(out, want) { + t.Errorf("missing %q in:\n%s", want, out) + } + } + if strings.Contains(BuildXMP("A cat.", nil, ""), `xml:lang="en"`) { + t.Error("English must stay x-default only") + } +}