Spaces:
No application file
No application file
| 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() | |