Create app.py

app.py

@@ -0,0 +1,46 @@
+# streamlit_app.py
+
+import streamlit as st
+from PIL import Image
+from transformers import AutoModelForImageSegmentation, AutoProcessor
+import numpy as np
+import matplotlib.pyplot as plt
+
+# Title of the app
+st.title("Image Segmentation App with Hugging Face and Streamlit")
+
+# Description
+st.write("Upload an image, and the Hugging Face model will segment it.")
+
+# Load the Hugging Face model and processor
+@st.cache_resource  # Cache the model to avoid reloading every time
+def load_model():
+    model_name = "ZhengPeng7/BiRefNet"
+    model = AutoModelForImageSegmentation.from_pretrained(model_name, trust_remote_code=True)
+    processor = AutoProcessor.from_pretrained(model_name)
+    return model, processor
+
+model, processor = load_model()
+
+# Upload an image
+uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+
+if uploaded_image:
+    # Display the uploaded image
+    image = Image.open(uploaded_image)
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    # Perform segmentation
+    st.write("Performing segmentation... Please wait!")
+    inputs = processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+
+    # Generate segmentation mask
+    segmentation = outputs.logits.argmax(dim=1)[0].detach().cpu().numpy()
+
+    # Display the segmentation mask
+    st.write("Segmentation mask:")
+    plt.figure(figsize=(10, 10))
+    plt.imshow(segmentation, cmap="viridis")
+    plt.axis("off")
+    st.pyplot(plt)