import streamlit as st
from rembg import remove
from PIL import Image
import numpy as np
import requests
from io import BytesIO


def process_image(image):
    output = remove(image)  
    
    input_rgb = np.array(image)[:, :, 0:3]
    output_rgba = np.array(output)

    alpha = output_rgba[:, :, 3]
    alpha3 = np.dstack((alpha, alpha, alpha))

    background_rgb = input_rgb.astype(float) * (1 - alpha3.astype(float) / 255)
    background_rgb = background_rgb.astype(np.uint8)

    background = Image.fromarray(background_rgb)
    return output, background

def display_images(image):
    img1, img2 = process_image(image)

    col3, col4 = st.columns(2)
    with col3:
        st.image(img1, caption='Foreground Extraction', use_column_width=True)
    
    with col4:
        st.image(img2, caption='Background Extraction', use_column_width=True)

def main():
    st.title("Foreground & Background Extraction")
    images = {
        "Sample 1": "https://images2.alphacoders.com/931/931778.jpg",
        "Sample 2": "https://64.media.tumblr.com/862739618bd130769be6efed4d2b8841/63e31bdeb0842a99-ee/s1280x1920/49eb215c37a2235c915ded605f0ebcb81962af6c.jpg"
    }
    col1, col2 = st.columns(2)
    sample_keys = list(images.keys())
    with col1:
        st.image(f"{images.get(sample_keys[0])}", caption=f"{sample_keys[0]}", width = 300)

    with col2:
        st.image(f"{images.get(sample_keys[1])}", caption=f"{sample_keys[1]}", width = 150)
    option = st.radio(
        "Choose an option:",
        ("Upload an Image", f"{sample_keys[0]}", f"{sample_keys[1]}")
    )
    
    if option in images:
        # sample_image = Image.open(images[option])
        response = requests.get(images[option])
        sample_image = Image.open(BytesIO(response.content))
        st.image(sample_image, caption=f'{option}', use_column_width=True)
        display_images(sample_image)

    elif option == "Upload an Image":
        uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
        
        if uploaded_file is not None:
            original_image = Image.open(uploaded_file)
            st.image(original_image, caption='Uploaded Image', use_column_width=True)
            display_images(original_image)

if __name__ == "__main__":
    main()