Created
March 2, 2024 12:02
-
-
Save fancellu/8d0c48a24353ec01acfabd0e82c3e65b to your computer and use it in GitHub Desktop.
Python Hugging Face image renamer using BeitForImageClassification (useful when you have thousands of badly named images)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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.") |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Results vary of course. A better model would probably do a better job, but the code would remain pretty much the same.