import gradio as gr
import spaces
import torch
from image_loader import load_image_from_url, load_image_from_file
from image_processor import process_image
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

torch.set_float32_matmul_precision(["high", "highest"][0])

try:
    birefnet = AutoModelForImageSegmentation.from_pretrained(
        "ZhengPeng7/BiRefNet", trust_remote_code=True
    )
    birefnet.to("cuda")
    logging.info("BiRefNet model loaded successfully.")
except Exception as e:
    logging.error(f"Error loading BiRefNet model: {e}")
    raise Exception(f"Error loading BiRefNet model: {e}")

def fn(image_input):
    try:
        if isinstance(image_input, str):  # URL input
            img = load_image_from_url(image_input)
        else:  # File upload
            img = load_image_from_file(image_input)

        img = img.convert("RGB")
        origin = img.copy()
        processed_image = process(img)
        return (processed_image, origin)
    except Exception as e:
        logging.error(f"Error in fn function: {e}")
        return None, None  # Return None or a placeholder image

@spaces.GPU
def process(image):
    try:
        processed_image = process_image(image, birefnet)
        return processed_image
    except Exception as e:
        logging.error(f"Error in process function: {e}")
        raise gr.Error(f"Error processing image: {e}")


def process_file(file_path):
    try:
        name_path = file_path.rsplit(".", 1)[0] + ".png"
        img = load_image_from_file(file_path)
        img = img.convert("RGB")
        transparent = process(img)
        transparent.save(name_path)
        logging.info(f"Processed image saved to: {name_path}")
        return name_path
    except Exception as e:
        logging.error(f"Error in process_file function: {e}")
        raise gr.Error(f"Error processing file: {e}")

slider1 = gr.ImageSlider(label="Processed Image", type="pil", format="png")
slider2 = gr.ImageSlider(label="Processed Image from URL", type="pil", format="png")
image_upload = gr.Image(label="Upload an image")
image_file_upload = gr.Image(label="Upload an image", type="filepath")
url_input = gr.Textbox(label="Paste an image URL")
output_file = gr.File(label="Output PNG File")

# Example images
try:
    chameleon = load_image_from_file("butterfly.jpg")
except Exception as e:
    logging.error(f"Error loading example image: {e}")
    chameleon = None  # Or a placeholder image

url_example = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"

tab1 = gr.Interface(fn, inputs=image_upload, outputs=slider1, examples=[chameleon], api_name="image")
tab2 = gr.Interface(fn, inputs=url_input, outputs=slider2, examples=[url_example], api_name="text")
tab3 = gr.Interface(process_file, inputs=image_file_upload, outputs=output_file, examples=["butterfly.jpg"], api_name="png")

demo = gr.TabbedInterface(
    [tab1, tab2, tab3], ["Image Upload", "URL Input", "File Output"], title="Background Removal Tool"
)

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