import gradio as gr
from PIL import Image
import spaces
from transformers import AutoModelForImageSegmentation
import torch
from torchvision import transforms
import requests
from io import BytesIO
import os

# --- Model and Processor Setup ---
# Use a higher precision for matrix multiplication for better performance
torch.set_float32_matmul_precision("high")

# Load the BiRefNet model for image segmentation
# trust_remote_code=True is required for this model
birefnet = AutoModelForImageSegmentation.from_pretrained(
    "ZhengPeng7/BiRefNet", trust_remote_code=True
)
# Move the model to the available device (GPU if available, otherwise CPU)
device = "cuda" if torch.cuda.is_available() else "cpu"
birefnet.to(device)

# Define the image transformation pipeline
transform_image = transforms.Compose(
    [
        transforms.Resize((1024, 1024)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)

# --- Helper Function to Load Images ---
def load_image(image_source, output_type="pil"):
    """
    Loads an image from a file path, URL, or numpy array.
    """
    if image_source is None:
        return None
        
    if isinstance(image_source, str):
        if image_source.startswith("http"):
            try:
                response = requests.get(image_source)
                response.raise_for_status()
                image = Image.open(BytesIO(response.content))
            except requests.exceptions.RequestException as e:
                raise gr.Error(f"Could not fetch image from URL: {e}")
        else:
            image = Image.open(image_source)
    elif hasattr(image_source, 'shape'): # Check if it's a numpy-like array
        image = Image.fromarray(image_source)
    else:
        image = image_source # Assume it's already a PIL image

    if output_type == "pil":
        return image.convert("RGB")
    return image

# --- Core Processing Function ---
# Use @spaces.GPU decorator if you plan to run this on a GPU-enabled Hugging Face Space
# @spaces.GPU
def process_image_to_transparent(image: Image.Image) -> Image.Image:
    """
    Takes a PIL image, removes the background, and returns a PIL image with an alpha channel.
    """
    if image is None:
        return None
    image_size = image.size
    # Unsqueeze adds a batch dimension, which the model expects
    input_tensor = transform_image(image).unsqueeze(0).to(device)

    # Prediction without tracking gradients for efficiency
    with torch.no_grad():
        # The model returns multiple outputs; the last one is the primary segmentation map
        preds = birefnet(input_tensor)[-1].sigmoid().cpu()

    # Process the prediction tensor to create a mask
    pred_tensor = preds[0].squeeze()
    mask_pil = transforms.ToPILImage()(pred_tensor)
    mask_resized = mask_pil.resize(image_size)

    # Apply the mask as an alpha channel to the original image
    image.putalpha(mask_resized)
    return image

# --- Gradio Interface Functions ---

def fn(image_source):
    """
    Handles image uploads and URLs, returning the processed image.
    """
    if image_source is None:
        return None
    
    pil_image = load_image(image_source, output_type="pil")
    processed_image = process_image_to_transparent(pil_image)
    return processed_image

def process_file(image_filepath):
    """
    Handles a single file upload and returns a downloadable processed file.
    """
    if image_filepath is None:
        return None

    # Define the output path for the new PNG file
    base_name = os.path.basename(image_filepath.name) # Use .name for Gradio file objects
    name, _ = os.path.splitext(base_name)
    output_path = f"{name}_transparent.png"
    
    # Load the image from the provided file path
    pil_image = load_image(image_filepath.name, output_type="pil")
    
    # Process the image
    transparent_image = process_image_to_transparent(pil_image)
    
    # Save the processed image to the new path
    transparent_image.save(output_path)
    
    # Return the path to the newly created file for download
    return output_path

# --- Gradio UI Definition ---

# Define example images for the interface
example_image_path = "butterfly.jpeg"
# You should have a 'butterfly.jpeg' in the same directory or provide a full path
# For demonstration, let's create a dummy example image if it doesn't exist.
if not os.path.exists(example_image_path):
    print(f"'{example_image_path}' not found. Creating a dummy image for example.")
    try:
        dummy_img = Image.new('RGB', (200, 200), color = 'red')
        dummy_img.save(example_image_path)
    except Exception as e:
        print(f"Could not create dummy image: {e}")

example_url = "https://i.ibb.co/67B6Knk9/students-1807505-1280.jpg"

# Define the individual interfaces for each tab
tab1 = gr.Interface(
    fn, 
    inputs=gr.Image(label="Upload an Image", type="pil"), 
    outputs=gr.Image(label="Processed Image", format="png"), 
    examples=[[example_image_path]], 
    api_name="image"
)

tab2 = gr.Interface(
    fn, 
    inputs=gr.Textbox(label="Paste an Image URL"), 
    outputs=gr.Image(label="Processed Image", format="png"), 
    examples=[[example_url]], 
    api_name="text"
)

tab3 = gr.Interface(
    process_file, 
    inputs=gr.File(label="Upload an Image File"), 
    outputs=gr.File(label="Download Processed PNG"), 
    examples=[[example_image_path]], 
    api_name="png"
)

# Combine the interfaces into a tabbed layout
demo = gr.TabbedInterface(
    [tab1, tab2, tab3], 
    ["Image Upload", "URL Input", "File Output"], 
    title="Background Removal Tool | CodeTechDevX"
)

if __name__ == "__main__":
    demo.launch(show_error=True)