Upload app.py

app.py

@@ -4,7 +4,7 @@ import torch
 from RealESRGAN import RealESRGAN
 from io import BytesIO
 
-# Define the target size for the image
+# Define the target size for the image (used for initial resizing before enhancement)
 TARGET_SIZE = (240, 240)
 
 # Function to load the model based on scale and anime toggle
@@ -21,21 +21,32 @@ def load_model(scale, anime=False):
     return model
 
 def enhance_image(image, scale, anime):
-    model = load_model(scale, anime=anime)
-    
-    # Convert image to RGB if it has an alpha channel
-    if image.mode != 'RGB':
-        image = image.convert('RGB')
-    
-    # Resize image to target dimensions
-    image = image.resize(TARGET_SIZE)
-    
-    sr_image = model.predict(image)
-    
-    buffer = BytesIO()
-    sr_image.save(buffer, format="PNG")
-    buffer.seek(0)
-    return sr_image, buffer
+    try:
+        model = load_model(scale, anime=anime)
+        
+        # Convert image to RGB if it has an alpha channel
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        # Store original image size
+        original_size = image.size
+        
+        # Resize image to target dimensions for processing
+        image = image.resize(TARGET_SIZE)
+        
+        # Perform image enhancement
+        sr_image = model.predict(image)
+        
+        # Resize the enhanced image back to the original size
+        sr_image = sr_image.resize(original_size)
+        
+        buffer = BytesIO()
+        sr_image.save(buffer, format="PNG")
+        buffer.seek(0)
+        
+        return sr_image, buffer, None
+    except Exception as e:
+        return None, None, str(e)
 
 def main():
     st.title("Generative AI Image Restoration")
@@ -59,22 +70,25 @@ def main():
         
         # Enhance button
         if st.button("Restore Image"):
-            enhanced_image, buffer = enhance_image(image, scale_value, anime)
-            
-            # Show images side by side
-            col1, col2 = st.columns(2)
-            with col1:
-                st.image(image, caption="Original Image", use_column_width=True)
-            with col2:
-                st.image(enhanced_image, caption="Enhanced Image", use_column_width=True)
+            enhanced_image, buffer, error_message = enhance_image(image, scale_value, anime)
             
-            # Download button
-            st.download_button(
-                label="Download Enhanced Image",
-                data=buffer,
-                file_name="enhanced_image.png",
-                mime="image/png"
-            )
+            if error_message:
+                st.error(f"An error occurred: {error_message}")
+            else:
+                # Show images side by side
+                col1, col2 = st.columns(2)
+                with col1:
+                    st.image(image, caption="Original Image", use_column_width=True)
+                with col2:
+                    st.image(enhanced_image, caption="Enhanced Image", use_column_width=True)
+                
+                # Download button
+                st.download_button(
+                    label="Download Enhanced Image",
+                    data=buffer,
+                    file_name="enhanced_image.png",
+                    mime="image/png"
+                )
 
 if __name__ == "__main__":
-    main()
+    main()