Update app.py

app.py

@@ -1,18 +1,14 @@
 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
-import zipfile
-import os
+import tempfile
 import time
+import os
 
-# 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,72 +17,101 @@ def load_model(scale):
         print(f"Weights for scale {scale} loaded successfully.")
     except Exception as e:
         print(f"Error loading weights for scale {scale}: {e}")
+        model.load_weights(weights_path, download=False)
     return model
 
 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)
-        return Image.fromarray(np.uint8(result))
+        print(f"Image converted to numpy array: shape {image_np.shape}, dtype {image_np.dtype}")
+        
+        if scale == '2x':
+            result = model2.predict(image_np)
+        elif scale == '4x':
+            result = model4.predict(image_np)
+        else:
+            result = model8.predict(image_np)
+            
+        enhanced_image = Image.fromarray(np.uint8(result))
+        print(f"Image enhanced in {time.time() - start_time:.2f} seconds")
+        return enhanced_image
     except Exception as e:
         print(f"Error enhancing image: {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 muda_dpi(input_image, dpi):
+    dpi_tuple = (dpi, dpi)
+    image = Image.fromarray(input_image.astype('uint8'), 'RGB')
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg')
+    image.save(temp_file, format='JPEG', dpi=dpi_tuple)
+    temp_file.close()
+    return Image.open(temp_file.name)
 
-def process_images(image_files, enhance, scale, generate_description):
+def resize_image(input_image, width, height):
+    image = Image.fromarray(input_image.astype('uint8'), 'RGB')
+    resized_image = image.resize((width, height))
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg')
+    resized_image.save(temp_file, format='JPEG')
+    temp_file.close()
+    return Image.open(temp_file.name)
+
+def process_images(image_files, enhance, scale, adjust_dpi, dpi, resize, width, height):
     processed_images = []
-    descriptions = []
-    zip_buffer = io.BytesIO()
+    file_paths = []
     
     for image_file in image_files:
-        image = Image.open(image_file).convert('RGB')
+        input_image = np.array(Image.open(image_file).convert('RGB'))
+        original_image = Image.fromarray(input_image.astype('uint8'), 'RGB')
         
         if enhance:
-            image = enhance_image(image, scale)
+            original_image = enhance_image(original_image, scale)
+        
+        if adjust_dpi:
+            original_image = muda_dpi(np.array(original_image), dpi)
+            
+        if resize:
+            original_image = resize_image(np.array(original_image), width, height)
+        
+        # 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(' ', '_')
         
-        if generate_description:
-            description = describe_image(image)
-            descriptions.append(description)
+        # 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')
         
-        # 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())))
-        with zipfile.ZipFile(zip_buffer, 'a') as zipf:
-            zipf.writestr(os.path.basename(image_file.name), buffer.getvalue())
+        processed_images.append(original_image)
+        file_paths.append(output_path)
     
-    zip_buffer.seek(0)
-    return processed_images, zip_buffer, descriptions
+    return processed_images, file_paths
 
 iface = gr.Interface(
     fn=process_images,
     inputs=[
-        gr.Files(label="Upload Image Files"),
+        gr.Files(label="Upload Image Files"),  # Use gr.Files for multiple file uploads
         gr.Checkbox(label="Enhance Images (ESRGAN)"),
         gr.Radio(['2x', '4x', '8x'], type="value", value='2x', label='Resolution model'),
-        gr.Checkbox(label="Generate Image Descriptions")
+        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)
     ],
     outputs=[
-        gr.Gallery(label="Enhanced Images"),
-        gr.File(label="Download Enhanced Images (ZIP)"),
-        gr.Textbox(label="Generated Descriptions", lines=5)
+        gr.Gallery(label="Final Images"),  # Use gr.Gallery to display multiple images
+        gr.Files(label="Download Final Images")
     ],
-    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."
+    title="Multi-Image Enhancer",
+    description="Upload multiple images (.jpg, .png), enhance using AI, adjust DPI, resize, and download the final results."
 )
 
-iface.launch()
+iface.launch(debug=True)