# 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)