Update app.py

app.py

@@ -1,16 +1,19 @@
 
+
+import torch
 from PIL import Image
 from RealESRGAN import RealESRGAN
 import gradio as gr
 import numpy as np
 import tempfile
 import time
-import zipfile
 import os
+from transformers import pipeline  # For Hugging Face image description generation
 
-# Set the device to CUDA if available, otherwise CPU
+# Check for GPU availability
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+# Load RealESRGAN model with specified scale
 def load_model(scale):
     model = RealESRGAN(device, scale=scale)
     weights_path = f'weights/RealESRGAN_x{scale}.pth'
@@ -27,6 +30,10 @@ model2 = load_model(2)
 model4 = load_model(4)
 model8 = load_model(8)
 
+# Hugging Face image description pipeline
+description_generator = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
+
+# Enhance image based on selected scale
 def enhance_image(image, scale):
     try:
         print(f"Enhancing image with scale {scale}...")
@@ -48,6 +55,18 @@ def enhance_image(image, scale):
         print(f"Error enhancing image: {e}")
         return image
 
+# Generate image description using Hugging Face Transformers
+def generate_description(image):
+    try:
+        print("Generating description for the image...")
+        description = description_generator(image)[0]['generated_text']
+        print(f"Description generated: {description}")
+        return description
+    except Exception as e:
+        print(f"Error generating description: {e}")
+        return "Description unavailable."
+
+# Adjust DPI of an image
 def muda_dpi(input_image, dpi):
     dpi_tuple = (dpi, dpi)
     image = Image.fromarray(input_image.astype('uint8'), 'RGB')
@@ -56,6 +75,7 @@ def muda_dpi(input_image, dpi):
     temp_file.close()
     return Image.open(temp_file.name)
 
+# Resize an image to specified dimensions
 def resize_image(input_image, width, height):
     image = Image.fromarray(input_image.astype('uint8'), 'RGB')
     resized_image = image.resize((width, height))
@@ -64,9 +84,11 @@ def resize_image(input_image, width, height):
     temp_file.close()
     return Image.open(temp_file.name)
 
+# Process a list of images with various options
 def process_images(image_files, enhance, scale, adjust_dpi, dpi, resize, width, height):
     processed_images = []
-    temp_dir = tempfile.mkdtemp()
+    file_paths = []
+    descriptions = []  # List to store descriptions
     
     for image_file in image_files:
         input_image = np.array(Image.open(image_file).convert('RGB'))
@@ -81,23 +103,27 @@ def process_images(image_files, enhance, scale, adjust_dpi, dpi, resize, width,
         if resize:
             original_image = resize_image(np.array(original_image), width, height)
         
-        # Save each image as JPEG, preserving the original filename
-        file_name = os.path.basename(image_file.name)
-        output_path = os.path.join(temp_dir, file_name)
+        # Generate description
+        description = generate_description(original_image)
+        descriptions.append(description)
+        
+        # Sanitize the base filename
+        base_name = os.path.basename(image_file.name)
+        file_name, _ = os.path.splitext(base_name)
+        
+        # Remove any characters that aren't alphanumeric, spaces, underscores, or hyphens
+        file_name = ''.join(e for e in file_name if e.isalnum() or e in (' ', '_', '-')).strip().replace(' ', '_')
+        
+        # Create a final file path without unnecessary suffixes
+        output_path = os.path.join(tempfile.gettempdir(), f"{file_name}.jpg")
         original_image.save(output_path, format='JPEG')
-        processed_images.append(output_path)
-    
-    # Create a ZIP file with all processed images
-    zip_path = os.path.join(temp_dir, 'processed_images.zip')
-    with zipfile.ZipFile(zip_path, 'w') as zipf:
-        for file_path in processed_images:
-            zipf.write(file_path, os.path.basename(file_path))
-    
-    # Load images for display in the gallery
-    display_images = [Image.open(img_path) for img_path in processed_images]
+        
+        processed_images.append(original_image)
+        file_paths.append(output_path)
     
-    return display_images, zip_path
+    return processed_images, file_paths, descriptions
 
+# Set up Gradio interface with share=True for public access
 iface = gr.Interface(
     fn=process_images,
     inputs=[
@@ -111,11 +137,12 @@ iface = gr.Interface(
         gr.Number(label="Height", value=512)
     ],
     outputs=[
-        gr.Gallery(label="Final Images"),  # Display the processed images
-        gr.File(label="Download Final Images (ZIP)")  # Provide a ZIP file for download
+        gr.Gallery(label="Final Images"),  # Use gr.Gallery to display multiple images
+        gr.Files(label="Download Final Images"),
+        gr.Textbox(label="Image Descriptions", lines=5)  # Display generated descriptions
     ],
-    title="bulk image upscaler",
-    description="Upload multiple images (.jpg, .png), enhance using AI, adjust DPI, resize, and download the final results as a ZIP file."
+    title="Multi-Image Enhancer with Hugging Face Descriptions",
+    description="Upload multiple images (.jpg, .png), enhance using AI, adjust DPI, resize, generate descriptions, and download the final results."
 )
 
-iface.launch(debug=True)
+iface.launch(debug=True, share=True)