import streamlit as st
import cv2
import numpy as np
from PIL import Image
import io

def remove_background(img):
    # Convert to grayscale and create mask for white background
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    _, mask = cv2.threshold(gray, 245, 255, cv2.THRESH_BINARY_INV)

    # Optionally refine mask with morphology (optional)
    kernel = np.ones((3, 3), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)

    # Keep original image where mask is set
    result = cv2.bitwise_and(img, img, mask=mask)
    return result

def convert_to_jpeg(image_array):
    rgb_image = cv2.cvtColor(image_array, cv2.COLOR_BGR2RGB)
    pil_img = Image.fromarray(rgb_image)
    buffer = io.BytesIO()
    pil_img.save(buffer, format="JPEG")
    buffer.seek(0)
    return buffer

st.title("🧼 Background Remover App")
st.write("Upload an image and remove the background using OpenCV. Your image content will not be changed.")

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    image = cv2.imdecode(file_bytes, 1)

    st.image(image, caption='Original Image', use_column_width=True)

    # Remove background (only background, keep rest intact)
    no_bg_image = remove_background(image)
    st.image(no_bg_image, caption='Image with Background Removed', use_column_width=True)

    jpeg_buffer = convert_to_jpeg(no_bg_image)
    st.download_button(
        label="📥 Download Image as JPEG",
        data=jpeg_buffer,
        file_name="no_background.jpg",
        mime="image/jpeg"
    )