Update app.py

app.py

@@ -1,18 +1,17 @@
-# app.py (Streamlit version)
+# app.py (Streamlit version - Corrected)
 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)
+# Install pytorch_hub_examples package (Corrected)
 try:
     import pytorch_hub_examples
-except ModuleNotFoundError:  # Handle if not installed (Hugging Face)
+except ModuleNotFoundError:
     subprocess.run(["git", "clone", "https://github.com/facebookresearch/pytorch_hub_examples.git"])
-    subprocess.run(["cd", "pytorch_hub_examples", "&&", "pip", "install", "-e", "."])
+    subprocess.run(["cd", "pytorch_hub_examples", "&&", "python", "setup.py", "install"])  # Correct installation
     import pytorch_hub_examples
 
 # Load the U-2-Net model
@@ -21,19 +20,18 @@ model.eval()
 
 # Define the transform
 transform = transforms.Compose([
-    transforms.Resize((320, 320)),  # Resize for U-2-Net
+    transforms.Resize((320, 320)),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # ImageNet stats
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-
 def remove_background(image):
     try:
         img = transform(image).unsqueeze(0)
 
         with torch.no_grad():
             out = model(img)
-            mask = torch.sigmoid(out[0]) # Get the mask
+            mask = torch.sigmoid(out[0])
 
         mask = mask.squeeze().cpu().numpy()
         mask = (mask * 255).astype(np.uint8)
@@ -49,20 +47,19 @@ def remove_background(image):
 
         return new_image
     except Exception as e:
-        st.error(f"Error: {e}")  # Display error in Streamlit
-        return None  # Return None in case of error
-
+        st.error(f"Error: {e}")
+        return None
 
 st.title("Background Remover")
 
 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
+    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
+        with st.spinner("Removing background..."):
             result_image = remove_background(image)
             if result_image:
                 st.image(result_image, caption="Background Removed", use_column_width=True)