tag images with ollama

tag_images_with_ai.sh

#!/bin/bash

# dry run mode
dry_run=0
if [[ "$1" == "-n" || "$1" == "--dry-run" ]]; then
    dry_run=1
    echo -e "\tℹ️ Dry run mode enabled. No changes will be made."
    shift  # remove the dry run argument
fi

if [ -z "$1" ]; then
    # invalid number of aguments
    echo "Usage: $0 <glob-pattern>"
    exit 1
fi

# if you run the API on another address, you
# can customize by using `TGF_API_URL` environment variable
url="${TGF_API_URL:-http://localhost:11434/api/generate}"

# you can define a custom prompt
# with `TGF_PROMPT` environment variable
prompt="${TGF_PROMPT:-What is in this picture? Answer with 10 keywords as a JSON array with format [\"keyword1\",\"keyword2\",\"keyword3\",\"keyword4\",\"keyword5\",\"keyword6\",\"keyword7\",\"keyword8\",\"keyword9\",\"keyword10\"].}"

# if you prefer another model, you
# can customize by using `TGF_LLM_MODEL` environment variable
model="${TGF_LLM_MODEL:-llava}"

# if you prefer another temperature value, you
# can customize by using `TGF_LLM_TEMPERATURE` environment variable
temperature="${TGF_LLM_TEMPERATURE:-0}"

# temporary activation of special options
shopt -s nullglob

for pattern in "$@"; do  # iterate of each pattern from command line
    for file in $pattern; do  # iterate over each existing file represented by the pattern
        full_file_path=$(realpath "$file")

        echo "Processing $full_file_path ..."
        
        # read the current file and save as Base64
        base64string=$(base64 -i "$full_file_path")

        # create a JSON for the Ollama API
        json_data=$(jq -n \
                        --arg model "$model" \
                        --arg prompt "$prompt" \
                        --arg img "$base64string" \
                        --argjson temperature "$temperature" \
                        '{model: $model, prompt: $prompt, options: {temperature: $temperature}, stream: false, images: [$img]}')

        # execute Ollama API
        echo -e "\tℹ️ Sending POST request..."
        response=$(wget --header "Content-Type: application/json" \
                        --post-data "$json_data" \
                        "$url" \
                        -q \
                        -O - 2>&1) # Captures both the output and server response headers

        # if the execution was successful, there should
        # be a `response` property in the output
        extracted_response=$(echo "$response" | jq -r '.response')

        if [ -n "$extracted_response" ]; then
            # we have a `response` property, but it can happen
            # that specific LLMs return Markdown code
            array_with_keywords=$(echo "$extracted_response" | sed -n '/```json/,/```/p' | sed -e '1d' -e '$d' | sed -e 's/^[[:space:]]*//g' -e 's/[[:space:]]*$//g' -e 's/[[:space:]]*```+.*$//')

            # check if we have a valid JSON array in `$extracted_response`
            if ! echo "$array_with_keywords" | jq -e . >/dev/null 2>&1; then
                echo -e "\t⚠️ Could not find JSON array in response!"
                continue
            fi

            # convert JSON string to sorted list of keywords
            sorted_keywords=$(echo "$array_with_keywords" | jq -r '.[]' | sort | uniq | paste -sd ',' -)

            # do not continue if we have no keywords from `jq`
            if [ -z "$sorted_keywords" ]; then
                echo -e "\t⚠️ No keywords found for '$full_file_path'!"
                continue
            fi

            # the next steps will write to file, so
            # check if dry mode is enabled
            if [[ $dry_run -eq 1 ]]; then
                echo -e "\tℹ️ Found following keywords for '$full_file_path': $sorted_keywords"
                continue
            fi

            # first try remove old keywords
            if ! exiftool -keywords= "$full_file_path" -overwrite_original >/dev/null 2>&1; then
                echo -e "\t⚠️ Failed to remove old keywords from '$full_file_path'."
                continue
            fi

            # now write the final keywords
            if exiftool -keywords="$keywords" "$full_file_path" -overwrite_original >/dev/null 2>&1; then
                echo -e "\t✅ Wrote following keywords to '$full_file_path': $sorted_keywords"
            else
                echo -e "\t⚠️ Failed to update keywords for '$full_file_path'."
            fi
        else
            echo -e "\t⚠️ No valid JSON found!"
        fi
    done
done

# now deactivate temporary options
shopt -u nullglob