viruthik's picture
Update Background-Remover-Flask-App-using-Python-main/app.py
4743e33
import streamlit as st
import os
from rembg import remove
from PIL import Image
from werkzeug.utils import secure_filename
UPLOAD_FOLDER = 'static/uploads'
ALLOWED_EXTENSIONS = set(['png', 'jpg', 'jpeg', 'webp'])
if 'static' not in os.listdir('.'):
os.mkdir('static')
if 'uploads' not in os.listdir('static/'):
os.mkdir('static/uploads')
def allowed_file(filename):
return '.' in filename and filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS
def remove_background(input_path, output_path):
input_image = Image.open(input_path)
output_image = remove(input_image)
output_image.save(output_path)
def main():
st.title("Remove Background")
uploaded_file = st.file_uploader("Choose a file", type=['png', 'jpg', 'jpeg', 'webp'])
if uploaded_file is not None and allowed_file(uploaded_file.name):
filename = secure_filename(uploaded_file.name)
file_path = os.path.join(UPLOAD_FOLDER, filename)
uploaded_file.save(file_path)
rembg_img_name = filename.split('.')[0] + "_rembg.png"
remove_background(file_path, os.path.join(UPLOAD_FOLDER, rembg_img_name))
st.image(Image.open(file_path), caption='Original Image', use_column_width=True)
st.image(Image.open(os.path.join(UPLOAD_FOLDER, rembg_img_name)),
caption='Image with Background Removed', use_column_width=True)
if __name__ == '__main__':
main()