sirishaReddy05/BackgroundImageRemoval

1import streamlit as st2from rembg import remove3from PIL import Image4import numpy as np5from io import BytesIO6import base647import os8import traceback9import time10 11st.set_page_config(layout="wide", page_title="Image Background Remover")12 13st.write("## Remove background from your image")14st.write(15    ":dog: Try uploading an image to watch the background magically removed. Full quality images can be downloaded from the sidebar. This code is open source and available [here](https://github.com/tyler-simons/BackgroundRemoval) on GitHub. Special thanks to the [rembg library](https://github.com/danielgatis/rembg) :grin:"16)17st.sidebar.write("## Upload and download :gear:")18 19# Increased file size limit20MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB21 22# Max dimensions for processing23MAX_IMAGE_SIZE = 2000  # pixels24 25# Download the fixed image26def convert_image(img):27    buf = BytesIO()28    img.save(buf, format="PNG")29    byte_im = buf.getvalue()30    return byte_im31 32# Resize image while maintaining aspect ratio33def resize_image(image, max_size):34    width, height = image.size35    if width <= max_size and height <= max_size:36        return image37    38    if width > height:39        new_width = max_size40        new_height = int(height * (max_size / width))41    else:42        new_height = max_size43        new_width = int(width * (max_size / height))44    45    return image.resize((new_width, new_height), Image.LANCZOS)46 47@st.cache_data48def process_image(image_bytes):49    """Process image with caching to avoid redundant processing"""50    try:51        image = Image.open(BytesIO(image_bytes))52        # Resize large images to prevent memory issues53        resized = resize_image(image, MAX_IMAGE_SIZE)54        # Process the image55        fixed = remove(resized)56        return image, fixed57    except Exception as e:58        st.error(f"Error processing image: {str(e)}")59        return None, None60 61def fix_image(upload):62    try:63        start_time = time.time()64        progress_bar = st.sidebar.progress(0)65        status_text = st.sidebar.empty()66        67        status_text.text("Loading image...")68        progress_bar.progress(10)69        70        # Read image bytes71        if isinstance(upload, str):72            # Default image path73            if not os.path.exists(upload):74                st.error(f"Default image not found at path: {upload}")75                return76            with open(upload, "rb") as f:77                image_bytes = f.read()78        else:79            # Uploaded file80            image_bytes = upload.getvalue()81        82        status_text.text("Processing image...")83        progress_bar.progress(30)84        85        # Process image (using cache if available)86        image, fixed = process_image(image_bytes)87        if image is None or fixed is None:88            return89        90        progress_bar.progress(80)91        status_text.text("Displaying results...")92        93        # Display images94        col1.write("Original Image :camera:")95        col1.image(image)96        97        col2.write("Fixed Image :wrench:")98        col2.image(fixed)99        100        # Prepare download button101        st.sidebar.markdown("\n")102        st.sidebar.download_button(103            "Download fixed image", 104            convert_image(fixed), 105            "fixed.png", 106            "image/png"107        )108        109        progress_bar.progress(100)110        processing_time = time.time() - start_time111        status_text.text(f"Completed in {processing_time:.2f} seconds")112        113    except Exception as e:114        st.error(f"An error occurred: {str(e)}")115        st.sidebar.error("Failed to process image")116        # Log the full error for debugging117        print(f"Error in fix_image: {traceback.format_exc()}")118 119# UI Layout120col1, col2 = st.columns(2)121my_upload = st.sidebar.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])122 123# Information about limitations124with st.sidebar.expander("ℹ️ Image Guidelines"):125    st.write("""126    - Maximum file size: 10MB127    - Large images will be automatically resized128    - Supported formats: PNG, JPG, JPEG129    - Processing time depends on image size130    """)131 132# Process the image133if my_upload is not None:134    if my_upload.size > MAX_FILE_SIZE:135        st.error(f"The uploaded file is too large. Please upload an image smaller than {MAX_FILE_SIZE/1024/1024:.1f}MB.")136    else:137        fix_image(upload=my_upload)138else:139    # Try default images in order of preference140    default_images = ["./zebra.jpg", "./wallaby.png"]141    for img_path in default_images:142        if os.path.exists(img_path):143            fix_image(img_path)144            break145    else:146        st.info("Please upload an image to get started!")