Update app.py

app.py

@@ -1,6 +1,7 @@
 import torch
 from PIL import Image
 from RealESRGAN import RealESRGAN
+from transformers import BlipProcessor, BlipForConditionalGeneration  # Example for Hugging Face model
 import gradio as gr
 import numpy as np
 import io
@@ -11,6 +12,7 @@ import time
 # Set the device to CUDA if available, otherwise CPU
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+# Load the RealESRGAN models for enhancement
 def load_model(scale):
     model = RealESRGAN(device, scale=scale)
     weights_path = f'weights/RealESRGAN_x{scale}.pth'
@@ -21,44 +23,33 @@ def load_model(scale):
         print(f"Error loading weights for scale {scale}: {e}")
     return model
 
-# Load models for different scales
 model2 = load_model(2)
 model4 = load_model(4)
 model8 = load_model(8)
 
+# Load Hugging Face model and processor for image description
+processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+caption_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
+
 def enhance_image(image, scale):
     try:
-        print(f"Enhancing image with scale {scale}...")
-        start_time = time.time()
         image_np = np.array(image.convert('RGB'))
         model = model2 if scale == '2x' else model4 if scale == '4x' else model8
         result = model.predict(image_np)
-        enhanced_image = Image.fromarray(np.uint8(result))
-        print(f"Image enhanced in {time.time() - start_time:.2f} seconds")
-        return enhanced_image
+        return Image.fromarray(np.uint8(result))
     except Exception as e:
         print(f"Error enhancing image: {e}")
         return image
 
-def muda_dpi(image, dpi):
-    try:
-        with io.BytesIO() as output:
-            image.save(output, format='JPEG', dpi=(dpi, dpi))
-            return Image.open(output)
-    except Exception as e:
-        print(f"Error adjusting DPI: {e}")
-        return image
+def describe_image(image):
+    inputs = processor(image, return_tensors="pt").to(device)
+    generated_ids = caption_model.generate(**inputs)
+    description = processor.decode(generated_ids[0], skip_special_tokens=True)
+    return description
 
-def resize_image(image, width, height):
-    try:
-        resized_image = image.resize((width, height))
-        return resized_image
-    except Exception as e:
-        print(f"Error resizing image: {e}")
-        return image
-
-def process_images(image_files, enhance, scale, adjust_dpi, dpi, resize, width, height):
+def process_images(image_files, enhance, scale, generate_description):
     processed_images = []
+    descriptions = []
     zip_buffer = io.BytesIO()
     
     for image_file in image_files:
@@ -67,13 +58,11 @@ def process_images(image_files, enhance, scale, adjust_dpi, dpi, resize, width,
         if enhance:
             image = enhance_image(image, scale)
         
-        if adjust_dpi:
-            image = muda_dpi(image, dpi)
-        
-        if resize:
-            image = resize_image(image, width, height)
+        if generate_description:
+            description = describe_image(image)
+            descriptions.append(description)
         
-        # Save image to the in-memory ZIP buffer
+        # Save enhanced image to ZIP in-memory buffer
         buffer = io.BytesIO()
         image.save(buffer, format='JPEG')
         processed_images.append(Image.open(io.BytesIO(buffer.getvalue())))
@@ -81,26 +70,23 @@ def process_images(image_files, enhance, scale, adjust_dpi, dpi, resize, width,
             zipf.writestr(os.path.basename(image_file.name), buffer.getvalue())
     
     zip_buffer.seek(0)
-    return processed_images, zip_buffer
+    return processed_images, zip_buffer, descriptions
 
 iface = gr.Interface(
     fn=process_images,
     inputs=[
-        gr.Files(label="Upload Image Files"),  # Use gr.Files for multiple file uploads
+        gr.Files(label="Upload Image Files"),
         gr.Checkbox(label="Enhance Images (ESRGAN)"),
         gr.Radio(['2x', '4x', '8x'], type="value", value='2x', label='Resolution model'),
-        gr.Checkbox(label="Adjust DPI"),
-        gr.Number(label="DPI", value=300),
-        gr.Checkbox(label="Resize"),
-        gr.Number(label="Width", value=512),
-        gr.Number(label="Height", value=512)
+        gr.Checkbox(label="Generate Image Descriptions")
     ],
     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="Enhanced Images"),
+        gr.File(label="Download Enhanced Images (ZIP)"),
+        gr.Textbox(label="Generated Descriptions", lines=5)
     ],
-    title="Multi-Image Enhancer",
-    description="Upload multiple images (.jpg, .png), enhance using AI, adjust DPI, resize, and download the final results as a ZIP file."
+    title="Image Enhancer with Description Generator",
+    description="Upload multiple images, enhance using AI, generate descriptions using Hugging Face, and download results as a ZIP file."
 )
 
-iface.launch(debug=True)
+iface.launch()