Update app.py

app.py

@@ -1,32 +1,104 @@
+from transformers import BlipProcessor, BlipForConditionalGeneration
+from PIL import Image
 import pandas as pd
-from transformers import pipeline
+import numpy as np
+import gradio as gr
+import tempfile
+import os
+import csv
 
-# Function to enhance title using Hugging Face model
-def generate_title_with_huggingface(description):
-    result = generator(description, max_length=100, num_return_sequences=1)
-    return result[0]['generated_text']
+# Initialize the processor and model
+processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 
-# Load the CSV file directly (replace 'input.csv' with your actual file path)
-file_name = 'input.csv'  # Update this with the correct path to your CSV file
-df = pd.read_csv(file_name)
+def generate_image_caption(image):
+    """
+    Generate a caption for the given image.
 
-# Initialize the text generation model (e.g., GPT-2)
-generator = pipeline('text-generation', model='gpt-2')
+    Args:
+        image (PIL.Image): The image object.
 
-# Apply the function to generate SEO-friendly titles
-df['Title'] = df['Description'].apply(generate_title_with_huggingface)
+    Returns:
+        str: The generated caption.
+    """
+    # Convert image to RGB format
+    image = image.convert("RGB")
+    
+    # Preprocess the image and generate a caption
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model.generate(**inputs)
+    caption = processor.decode(outputs[0], skip_special_tokens=True)
+    
+    return caption
 
-# Function to generate keywords (basic example)
-def generate_keywords(description):
-    words = set(description.replace(',', '').replace('.', '').split())
-    return ",".join(list(words)[:50])
+def generate_keywords(caption):
+    """
+    Generate a list of keywords from the caption.
 
-# Generate basic keywords from the description (you can enhance this further)
-df['Keywords'] = df['Description'].apply(generate_keywords)
+    Args:
+        caption (str): The image caption.
 
-# Save the SEO-optimized DataFrame to a new CSV file
-seo_output_file_path = 'seo_huggingface_filename_title_keywords.csv'
-df[['Filename', 'Title', 'Keywords']].to_csv(seo_output_file_path, index=False)
+    Returns:
+        list: A list of single-word keywords.
+    """
+    # Example simple keyword extraction (use a more sophisticated method if needed)
+    words = caption.split()
+    keywords = list(set(words))[:50]  # Take unique words and limit to 50
+    return keywords
 
-# The output file will be saved in the current working directory
-print(f"File saved as {seo_output_file_path}")
+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)
+            # Ensure the title is within the 70 to 100 characters range
+            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(img_file.name) for img_file in image_files], csv_file_path
+
+# Define Gradio interface
+iface = gr.Interface(
+    fn=process_images,
+    inputs=[
+        gr.Files(label="Upload Image Files", type="file", multiple=True)
+    ],
+    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)