rename-images-in-a-directory-with-ai-locally.py

import requests
import json
import base64
import os
from typing import List


def get_keywords(image: str) -> List[str]:
    body = {
        "model": "llava-phi",
        "format": "json",
        "prompt": "Describe the image as 6 keywords. Output in JSON format. Use the following schema: { filename: str, keywords: List[str] }",
        "images": [image],
        "stream": False,
    }

    response = requests.post("http://localhost:11434/api/generate", json=body)

    if response.status_code != 200:
        print(f"Error: {response.status_code}: {response.reason}")
        return []
    else:
        json_response = response.json()
        keywords = json.loads(json_response["response"])
        return keywords.get("keywords", [])


def create_file_name(keywords: List[str], file_ext: str) -> str:
    new_file_name = ""
    if keywords:
        file_parts = [k.replace(" ", "_") for k in keywords]
        filtered_words = [w for w in file_parts if len(new_file_name) + len(w) <= 30]
        new_file_name = "-".join(filtered_words) + "." + file_ext
    return new_file_name


def main():
    current_path = os.getcwd()
    for file in os.listdir("."):
        if (
            file.endswith(".JPG")
            or file.endswith(".webp")
            or file.endswith(".WEBP")
            or file.endswith(".PNG")
        ):
            with open(os.path.join(current_path, file), "rb") as image_file:
                b64 = base64.b64encode(image_file.read()).decode()
            keywords = get_keywords(b64)
            new_file_name = create_file_name(keywords, file.split(".")[-1])
            if new_file_name:
                os.rename(
                    os.path.join(current_path, file),
                    os.path.join(current_path, new_file_name),
                )
                print(f"Copied {file} to {new_file_name}")


if __name__ == "__main__":
    main()