import streamlit as st
from rembg import remove, new_session
from PIL import Image
import io
import fitz  # PyMuPDF for handling PDFs

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

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

# 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)

    # Create a high-quality model session
    session = new_session("u2net_hq")  # Load the HQ model

    # Remove background using the HQ model
    with st.spinner("Removing background with high quality..."):
        output_image = remove(image, session=session)

    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 (U2Net HQ)")