Upload app.py

app.py

@@ -18,32 +18,28 @@ def load_model(scale, anime=False):
         model.load_weights(model_path)
         return model
     except Exception as e:
-        st.error(f"Error loading the model: {e}")
+        st.error(f"Failed to load the model: {e}")
         return None
 
 def enhance_image(image, scale, anime):
-    model = load_model(scale, anime=anime)
-    if model is None:
-        return None, None
-
     try:
+        model = load_model(scale, anime=anime)
+        if model is None:
+            return None, None
+        
         # Convert image to RGB if it has an alpha channel
         if image.mode != 'RGB':
             image = image.convert('RGB')
         
-        # Process the image with the model
         sr_image = model.predict(image)
         
-        # Ensure the enhanced image has the same size as the original
-        sr_image = sr_image.resize(image.size)
-        
-        # Save enhanced image to buffer
         buffer = BytesIO()
         sr_image.save(buffer, format="PNG")
         buffer.seek(0)
         return sr_image, buffer
+    
     except Exception as e:
-        st.error(f"Error enhancing the image: {e}")
+        st.error(f"An error occurred during image enhancement: {e}")
         return None, None
 
 def main():
@@ -53,38 +49,41 @@ def main():
     uploaded_image = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg"])
     
     if uploaded_image is not None:
-        image = Image.open(uploaded_image)
-        
-        # Anime toggle
-        anime = st.checkbox("Anime Image", value=False)
-        
-        # Conditional scale options
-        if anime:
-            scale = "4x"  # Set to 4x automatically when anime is selected
-        else:
-            scale = st.radio("Upscaling Factor", ["2x", "4x", "8x"], index=0)
-        
-        scale_value = int(scale.replace('x', ''))
-        
-        # Enhance button
-        if st.button("Restore Image"):
-            enhanced_image, buffer = enhance_image(image, scale_value, anime)
+        try:
+            image = Image.open(uploaded_image)
+            
+            # Anime toggle
+            anime = st.checkbox("Anime Image", value=False)
+            
+            # Conditional scale options
+            if anime:
+                scale = "4x"  # Set to 4x automatically when anime is selected
+            else:
+                scale = st.radio("Upscaling Factor", ["2x", "4x", "8x"], index=0)
+            
+            scale_value = int(scale.replace('x', ''))
             
-            if enhanced_image:
-                # 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)
+            # Enhance button
+            if st.button("Restore Image"):
+                enhanced_image, buffer = 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 enhanced_image:
+                    # 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"
+                    )
+        except Exception as e:
+            st.error(f"An error occurred while processing the image: {e}")
 
 if __name__ == "__main__":
     main()