|
|
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: |
|
|
|
|
|
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() |
