import streamlit as st
from rembg import remove
from PIL import Image
import io
import fitz  # PyMuPDF for PDF handling

st.set_page_config(page_title="Background Remover", page_icon="🎨", layout="centered")

st.title("🎨 Background Remover")
st.markdown("Upload an image (PNG, JPEG, JPG, WEBP, or PDF) and remove the background instantly!")

# File uploader
uploaded_file = st.file_uploader("Upload an image or PDF", type=["png", "jpeg", "jpg", "webp", "pdf"])

if uploaded_file is not None:
    file_type = uploaded_file.type
    
    if file_type == "application/pdf":
        # Convert PDF to image
        doc = fitz.open(stream=uploaded_file.read(), filetype="pdf")
        page = doc.load_page(0)  # Only process the first page
        image = Image.frombytes("RGB", [int(page.rect.width), int(page.rect.height)], page.get_pixmap().samples)
    else:
        image = Image.open(uploaded_file)
    
    st.image(image, caption="Original Image", use_column_width=True)
    
    # Remove background
    with st.spinner("Removing background..."):
        output_image = remove(image)
    
    st.image(output_image, caption="Background Removed", use_column_width=True)
    
    # Download button
    img_byte_arr = io.BytesIO()
    output_image.save(img_byte_arr, format="PNG")
    st.download_button("Download Image", img_byte_arr.getvalue(), file_name="no_bg_image.png", mime="image/png")

st.markdown("---")
st.markdown("💡 Developed with ❤️ using Streamlit and rembg")