AI Image Tag Generator

 """

AI Image Tag Generator (transformers + PIL)

- Uses Hugging Face transformers image-classification pipeline (ViT)

- Returns top-k labels as tags with confidence scores

- Default input path is the uploaded file: /mnt/data/image.png

"""

from transformers import pipeline

from PIL import Image

import argparse

import os

# Default path (file uploaded in this session)

DEFAULT_IMAGE_PATH = "/mnt/data/image.png"

def generate_tags(image_path, model_name="google/vit-base-patch16-224", top_k=5):

    """

    Generate tags for an image using a HF transformers image-classification pipeline.

    Returns a list of (label, score) tuples.

    """

    if not os.path.exists(image_path):

        raise FileNotFoundError(f"Image not found: {image_path}")

    classifier = pipeline("image-classification", model=model_name)

    img = Image.open(image_path).convert("RGB")

    preds = classifier(img, top_k=top_k)

    # preds is a list of dicts: [{"label": "...", "score": 0.XX}, ...]

    tags = [(p["label"], float(p["score"])) for p in preds]

    return tags

def pretty_print_tags(tags):

    print("Generated tags:")

    for label, score in tags:

        print(f" - {label}  ({score:.2f})")

def parse_args():

    p = argparse.ArgumentParser(description="AI Image Tag Generator")

    p.add_argument("--image", "-i", default=DEFAULT_IMAGE_PATH, help="Path to input image")

    p.add_argument("--model", "-m", default="google/vit-base-patch16-224", help="HuggingFace model name")

    p.add_argument("--topk", type=int, default=5, help="Number of top tags to return")

    return p.parse_args()

def main():

    args = parse_args()

    try:

        tags = generate_tags(args.image, model_name=args.model, top_k=args.topk)

        pretty_print_tags(tags)

    except Exception as e:

        print("Error:", e)

if __name__ == "__main__":

    main()