import streamlit as st
from carvekit.api.interface import Interface
from carvekit.ml.wrap.fba_matting import FBAMatting
from carvekit.ml.wrap.tracer_b7 import TracerUniversalB7
from carvekit.pipelines.postprocessing import MattingMethod
from carvekit.pipelines.preprocessing import PreprocessingStub
from carvekit.trimap.generator import TrimapGenerator
from PIL import Image

# Create Streamlit app title
st.title("Image Background Remover")

# Create a file uploader
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png"])

if uploaded_file is not None:
    # Load the image
    image = Image.open(uploaded_file)
    
    # Set up ML pipeline
    seg_net = TracerUniversalB7(device='cpu', batch_size=1)
    fba = FBAMatting(device='cpu', input_tensor_size=2048, batch_size=1)
    trimap = TrimapGenerator()
    preprocessing = PreprocessingStub()
    postprocessing = MattingMethod(matting_module=fba, trimap_generator=trimap, device='cpu')
    interface = Interface(pre_pipe=preprocessing, post_pipe=postprocessing, seg_pipe=seg_net)
    
    # Process the image
    processed_bg = interface([image])[0]
    
    # Display original and processed images
    col1, col2 = st.columns(2)
    with col1:
        st.image(image, caption='Original Image', use_column_width=True)
    with col2:
        st.image(processed_bg, caption='Background Removed', use_column_width=True)