import streamlit as st
from PIL import Image, ImageEnhance, ImageFilter
import numpy as np
import io
import cv2
from transformers import BlipProcessor, BlipForConditionalGeneration

# Load models for AI enhancements
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

# Function to enhance image quality using PIL
def enhance_image(image, enhance_type, factor):
    enhancer = None
    if enhance_type == 'Brightness':
        enhancer = ImageEnhance.Brightness(image)
    elif enhance_type == 'Contrast':
        enhancer = ImageEnhance.Contrast(image)
    elif enhance_type == 'Sharpness':
        enhancer = ImageEnhance.Sharpness(image)
    elif enhance_type == 'Color':
        enhancer = ImageEnhance.Color(image)
    
    if enhancer:
        return enhancer.enhance(factor)
    return image

# Function for applying filters (Grayscale, etc.)
def apply_filter(image, filter_type):
    if filter_type == "Grayscale":
        return image.convert("L")
    elif filter_type == "Blur":
        return image.filter(ImageFilter.GaussianBlur(radius=5))
    return image

# Function to remove background (using a simple model or API)
def remove_background(image):
    # Here you can integrate an API for background removal, e.g., remove.bg or Hugging Face models.
    # For now, it's a placeholder.
    return image  # Placeholder for background removal

# Main function
def main():
    st.title("Advanced Image Editor")

    st.sidebar.title("Image Editor Tools")
    uploaded_file = st.sidebar.file_uploader("Upload Image", type=["jpg", "png", "jpeg"])

    if uploaded_file is not None:
        image = Image.open(uploaded_file)

        st.image(image, caption="Uploaded Image", use_column_width=True)

        st.sidebar.subheader("Enhance Image")
        enhancement_type = st.sidebar.selectbox("Enhancement Type", ["Brightness", "Contrast", "Sharpness", "Color"])
        enhancement_factor = st.sidebar.slider("Enhancement Factor", 0.5, 2.0, 1.0, 0.1)

        enhanced_image = enhance_image(image, enhancement_type, enhancement_factor)
        st.sidebar.subheader("Apply Filter")
        filter_type = st.sidebar.selectbox("Select Filter", ["None", "Grayscale", "Blur"])
        filtered_image = apply_filter(enhanced_image, filter_type)

        st.image(filtered_image, caption="Processed Image", use_column_width=True)

        # Background Removal
        remove_bg = st.sidebar.checkbox("Remove Background")
        if remove_bg:
            filtered_image = remove_background(filtered_image)
            st.image(filtered_image, caption="Background Removed Image", use_column_width=True)

        # Download Button
        buf = io.BytesIO()
        filtered_image.save(buf, format="PNG")
        byte_im = buf.getvalue()
        st.sidebar.download_button("Download Edited Image", byte_im, "edited_image.png", "image/png")

if __name__ == "__main__":
    main()