Update app.py

app.py

@@ -1,16 +1,16 @@
-# app.py
-import gradio as gr
+# app.py (Streamlit version)
+import streamlit as st
 from PIL import Image
 import torch
 import torchvision.transforms as transforms
 import numpy as np
 import sys
+import subprocess
 
 # Install and import (for Hugging Face environment)
 try:
     import pytorch_hub_examples
 except ModuleNotFoundError:  # Handle if not installed (Hugging Face)
-    import subprocess
     subprocess.run(["git", "clone", "https://github.com/facebookresearch/pytorch_hub_examples.git"])
     subprocess.run(["cd", "pytorch_hub_examples", "&&", "pip", "install", "-e", "."])
     import pytorch_hub_examples
@@ -49,13 +49,20 @@ def remove_background(image):
 
         return new_image
     except Exception as e:
-        return f"Error: {e}"
+        st.error(f"Error: {e}")  # Display error in Streamlit
+        return None  # Return None in case of error
 
 
-with gr.Blocks() as demo:
-    image_input = gr.Image(type="pil")
-    image_output = gr.Image()
-    image_button = gr.Button("Remove Background")
-    image_button.click(remove_background, inputs=image_input, outputs=image_output)
+st.title("Background Remover")
 
-demo.launch()
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
+if uploaded_file is not None:
+    image = Image.open(uploaded_file).convert("RGB") # Ensure RGB for transform
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    if st.button("Remove Background"):
+        with st.spinner("Removing background..."):  # Show a spinner while processing
+            result_image = remove_background(image)
+            if result_image:
+                st.image(result_image, caption="Background Removed", use_column_width=True)