import gradio as gr
from rembg import remove
from PIL import Image, ImageEnhance
import os
import requests
from io import BytesIO

def enhance_image(image):
    """Enhance the image by adjusting sharpness, contrast, and brightness."""
    enhancer = ImageEnhance.Sharpness(image)
    image = enhancer.enhance(2.0)  # Increase sharpness

    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(1.2)  # Increase contrast

    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(1.1)  # Increase brightness

    return image

def remove_background(image, filename, enhance):
    # Remove the background from the image
    output = remove(image)
    
    # Enhance the image after removing the background (if enabled)
    if enhance:
        output = enhance_image(output)
    
    # Save the image with the specified filename
    if not filename.endswith(".png"):
        filename += ".png"
    output_path = os.path.join("outputs", filename)
    os.makedirs("outputs", exist_ok=True)  # Create the folder if it doesn't exist
    output.save(output_path)
    
    return output, output_path

# Load the example image from the URL
def load_example_image():
    url = "https://huggingface.co/spaces/Doubleupai/RemBgV0/resolve/main/example.jpg?raw=true"  # Use the raw URL
    response = requests.get(url)
    image = Image.open(BytesIO(response.content))
    return image

# Create the Gradio interface
iface = gr.Interface(
    fn=remove_background, 
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Textbox(label="Filename (without extension)", placeholder="Enter filename"),
        gr.Checkbox(label="Enhance Image After Background Removal", value=True)  # Checkbox for enhancement
    ], 
    outputs=[
        gr.Image(type="pil", label="Background-Removed Image"),
        gr.File(label="Download Image")
    ],
    title="RemBgV0",  # Application title
    description="Upload an image, specify a filename, and choose whether to enhance the image after removing the background.",
    examples=[
        [load_example_image(), "example_output", True]  # Example image, filename, and enhancement option
    ]
)

# Launch the application
iface.launch()