Python Hugging Face image renamer using BeitForImageClassification (useful when you have thousands of badly named images)

image_renamer.py

from PIL import Image
from transformers import BeitImageProcessor, BeitForImageClassification
import os

# Define the directory containing images
image_dir = "e:/media"

# Load the model and processor
processor = BeitImageProcessor.from_pretrained('microsoft/beit-base-patch16-224-pt22k-ft22k')
model = BeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224-pt22k-ft22k')

# Loop through each image file in the directory
for filename in os.listdir(image_dir):
    if filename.endswith((".jpg", ".png", ".jpeg")):  # Check for common image formats
        # Construct full image path
        image_path = os.path.join(image_dir, filename)

        try:
            # Open the image
            image = Image.open(image_path)

            # Preprocess the image
            inputs = processor(images=[image], return_tensors="pt")

            # Run inference
            outputs = model(**inputs)
            logits = outputs.logits

            # Get the predicted class index and label
            predicted_class_idx = logits.argmax(-1).item()
            predicted_label = model.config.id2label[predicted_class_idx]

            # Generate a new filename with predicted class as a prefix (modify as needed)
            new_filename = f"{predicted_label}_{filename}"

            # Construct the new file path
            new_path = os.path.join(image_dir, new_filename)

            # Rename the file
            os.rename(image_path, new_path)

            print(f"Image '{filename}' renamed to '{new_filename}'")

        except Exception as e:  # Handle potential errors like invalid images or missing files
            print(f"Error processing '{filename}': {e}")

print("Renaming process completed.")

Results vary of course. A better model would probably do a better job, but the code would remain pretty much the same.