File size: 2,282 Bytes
1521950 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
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() |