Update app.py

app.py

@@ -1,10 +1,15 @@
-import gradio as gr
 import os
 from PIL import Image
 import numpy as np
 import torch
 import torchvision.transforms as T
+import streamlit as st
 from model.u2net import U2NET
+from io import BytesIO
+
+# Constants
+MAX_FILE_SIZE = 5 * 1024 * 1024  # 5 MB
+DEFAULT_IMAGE_PATH = "default_image.png"  # Path to your default image
 
 # Initialize the U2NET model
 u2net = U2NET(in_ch=3, out_ch=1)
@@ -88,7 +93,7 @@ def apply_mask(image, mask):
     return masked_image
 
 def segment_image(image):
-    """Function to be used with Gradio for segmentation."""
+    """Function to be used for segmentation."""
     # Ensure image is a PIL Image
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
@@ -98,13 +103,57 @@ def segment_image(image):
     masked_image = apply_mask(image, prediction_u2net)
     return masked_image
 
-# Define the Gradio interface
-iface = gr.Interface(
-    fn=segment_image,
-    inputs=gr.Image(type="numpy"),
-    outputs=gr.Image(type="pil",format="png"),
-    title="Remove Background From Image"
-)
+def fix_image(upload=None):
+    """Processes an uploaded image or a default image."""
+    if upload is not None:
+        image = Image.open(upload)
+    else:
+        image = Image.open(DEFAULT_IMAGE_PATH)  # Load default image
+
+    st.image(image, caption='Selected Image', use_column_width=True)
+    
+    if st.button('Segment Image'):
+        masked_image = segment_image(image)
+        st.image(masked_image, caption='Segmented Image', use_column_width=True, format="PNG")
+        # Save the image to a BytesIO object for downloading
+        buffered = BytesIO()
+        masked_image.save(buffered, format="PNG")
+        st.download_button(
+            label="Download Segmented Image",
+            data=buffered.getvalue(),
+            file_name="segmented_image.png",
+            mime="image/png"
+        )
+
+# Define the pages
+def page_one():
+    """Page for image segmentation."""
+    st.title("Image Segmentation with U2NET")
+    st.write("Upload an image to segment it using the U2NET model. The background of the segmented output will be transparent.")
+
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
+    # Determine image processing
+    if uploaded_file is not None:
+        if uploaded_file.size > MAX_FILE_SIZE:
+            st.error("The uploaded file is too large. Please upload an image smaller than 5MB.")
+        else:
+            fix_image(upload=uploaded_file)
+    else:
+        fix_image()  # Use default image if none uploaded
+
+def page_two():
+    """Page for other code."""
+    st.title("Other Feature")
+    st.write("This page is for the second feature you want to implement.")
+    # Add other code or features here
+
+# Sidebar navigation
+st.sidebar.title("Navigation")
+page = st.sidebar.radio("Go to", ("Image Segmentation", "Other Feature"))
 
-# Launch the interface
-iface.launch()
+# Page selection logic
+if page == "Image Segmentation":
+    page_one()
+elif page == "Other Feature":
+    page_two()