rename.py

import os
import sys
from base64 import b64encode
from io import BytesIO

import ollama
from PIL import Image


def load_image(image_path: str) -> Image.Image:
    image = Image.open(image_path)
    if image.mode == "P" and "transparency" in image.info:
        image = image.convert("RGBA")
    return image


def encode_image(image: Image.Image) -> str:
    buffered = BytesIO()
    image.save(buffered, format=image.format)
    return b64encode(buffered.getvalue()).decode("utf-8")


def generate_filename_from_answer(answer: str) -> str:
    # Remove file extensions if present
    answer = answer.rsplit(".", 1)[0]

    filename = "".join(char.lower() if char.isalnum() else "_" for char in answer)
    words = filter(None, filename.split("_"))
    filtered_words = []
    for word in words:
        if not word.isdigit():
            filtered_words.append(word)

    if not filtered_words:
        return "unnamed_image"  # Fallback if all words are numbers

    return "_".join(filtered_words)


def generate_filename(image_path: str) -> str:
    image = load_image(image_path)
    encoded_image = encode_image(image)
    response = ollama.chat(
        model="llava-llama3",
        messages=[
            {
                "role": "user",
                "content": (
                    "Analyze the image and generate a concise, specific filename "
                    "that uniquely describes its content. Rules:\n"
                    "1. No special characters, only lowercase letters and underscores\n"
                    "2. Include only key elements\n"
                    "3. Prefer noun-verb format\n"
                    "4. Focus on main subject(s) and any notable action\n"
                    "5. Include distinctive visual elements or style if relevant\n"
                    "6. Avoid generic terms, numbers, or dates\n"
                    "The filename should clearly distinguish this image in a large collection. "
                    "Only output the filename and nothing else."
                ),
                "images": [encoded_image],
            }
        ],
    )
    answer = response["message"]["content"]
    return generate_filename_from_answer(answer)


def get_new_path(image_path: str, filename: str) -> str:
    directory = os.path.dirname(image_path)
    base_path = os.path.join(directory, filename)
    extension = os.path.splitext(image_path)[1]
    new_path = base_path + extension
    counter = 1
    while os.path.exists(new_path):
        new_path = f"{base_path}_{counter}{extension}"
        counter += 1
    return new_path


def rename_image(image_path: str, new_path: str) -> None:
    try:
        os.rename(image_path, new_path)
    except Exception as e:
        print(f"Error renaming image {image_path} to {new_path}: {e}")


def process_image(image_path: str, index: int, total: int) -> None:
    try:
        filename = generate_filename(image_path)
        new_path = get_new_path(image_path, filename)
        if new_path != image_path:
            rename_image(image_path, new_path)
            print(f"{index}/{total} Image renamed to: {new_path}")
        else:
            print(f"{index}/{total} Image already has the suggested name: {image_path}")
    except Exception as e:
        print(f"Error processing image {image_path}: {e}")


def process_folder(folder_path: str) -> None:
    supported_extensions = (".jpg", ".jpeg", ".png", ".bmp", ".gif")
    image_paths = [
        os.path.join(root, file)
        for root, _, files in os.walk(folder_path)
        for file in files
        if file.lower().endswith(supported_extensions)
    ]

    image_count = len(image_paths)
    print(f"Found {image_count} image{'s' if image_count != 1 else ''} in the folder.")

    for index, image_path in enumerate(image_paths, start=1):
        process_image(image_path, index, image_count)


def main(path: str) -> None:
    if os.path.isfile(path):
        process_image(path, 1, 1)
    elif os.path.isdir(path):
        process_folder(path)
    else:
        print(f"Invalid path: {path}")
        sys.exit(1)


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python test.py <image_path_or_folder_path>")
        sys.exit(1)

    path = sys.argv[1]
    main(path)