Spaces:
Build error
Build error
| from transformers import BlipProcessor, BlipForConditionalGeneration | |
| from PIL import Image | |
| import gradio as gr | |
| import tempfile | |
| import os | |
| import csv | |
| # Initialize the processor and model | |
| processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") | |
| model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base") | |
| def generate_image_caption(image): | |
| """ | |
| Generate a caption for the given image. | |
| Args: | |
| image (PIL.Image): The image object. | |
| Returns: | |
| str: The generated caption. | |
| """ | |
| image = image.convert("RGB") | |
| inputs = processor(images=image, return_tensors="pt") | |
| outputs = model.generate(**inputs) | |
| caption = processor.decode(outputs[0], skip_special_tokens=True) | |
| return caption | |
| def generate_keywords(caption): | |
| """ | |
| Generate a list of keywords from the caption, focusing on SEO relevance. | |
| Args: | |
| caption (str): The image caption. | |
| Returns: | |
| list: A list of 50 single-word keywords. | |
| """ | |
| words = caption.split() | |
| # Remove common stopwords and keep unique words for SEO relevance | |
| stopwords = {"a", "the", "and", "of", "in", "on", "with", "at", "for", "to", "is"} | |
| keywords = list(set(word.lower() for word in words if word.lower() not in stopwords)) | |
| # Ensure the list contains exactly 50 keywords (add repetitions if needed) | |
| if len(keywords) > 50: | |
| keywords = keywords[:50] | |
| elif len(keywords) < 50: | |
| keywords.extend(keywords[:50-len(keywords)]) | |
| return keywords | |
| def process_images(image_files): | |
| """ | |
| Process uploaded images to generate metadata and create a CSV file. | |
| Args: | |
| image_files (list of file-like objects): List of uploaded image files. | |
| Returns: | |
| tuple: A list of PIL images, path to the CSV file. | |
| """ | |
| metadata = [] | |
| temp_dir = tempfile.mkdtemp() | |
| for image_file in image_files: | |
| filename = os.path.basename(image_file.name) | |
| image = Image.open(image_file) | |
| caption = generate_image_caption(image) | |
| if caption: | |
| keywords = generate_keywords(caption) | |
| title = caption if 70 <= len(caption) <= 100 else caption[:100] | |
| metadata.append({ | |
| 'filename': filename, | |
| 'title': title, | |
| 'keywords': keywords | |
| }) | |
| # Create CSV file | |
| csv_file_path = os.path.join(temp_dir, 'images_metadata.csv') | |
| with open(csv_file_path, mode='w', newline='', encoding='utf-8') as file: | |
| writer = csv.writer(file) | |
| writer.writerow(['Filename', 'Title', 'Keywords']) | |
| for data in metadata: | |
| filename = data['filename'] | |
| title = data['title'] | |
| keywords = ','.join(data['keywords']) | |
| writer.writerow([filename, title, keywords]) | |
| return [Image.open(image_file.name) for image_file in image_files], csv_file_path | |
| # Define Gradio interface | |
| iface = gr.Interface( | |
| fn=process_images, | |
| inputs=[ | |
| gr.Files(label="Upload Image Files") # Use gr.Files for multiple file uploads | |
| ], | |
| outputs=[ | |
| gr.Gallery(label="Processed Images"), | |
| gr.File(label="Download Metadata CSV") | |
| ], | |
| title="Image Captioning and Metadata Generator", | |
| description="Upload multiple images to generate captions and metadata. Download the metadata as a CSV file." | |
| ) | |
| # Launch the interface | |
| iface.launch(debug=True) | |