import gradio as gr
import torch
from torchvision import models, transforms
from PIL import Image
import numpy as np
import cv2

# Load pre-trained segmentation model
model = models.segmentation.deeplabv3_resnet101(pretrained=True).eval()

# Function to perform segmentation
def segment_person(image):
    try:
        if not isinstance(image, Image.Image):
            raise ValueError("Invalid image format. Please upload a valid image.")

        preprocess = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        input_tensor = preprocess(image).unsqueeze(0)

        with torch.no_grad():
            output = model(input_tensor)['out'][0]

        output_predictions = output.argmax(0)
        person_class = 15  # Person class in DeepLabV3
        mask = output_predictions == person_class
        mask_np = mask.byte().cpu().numpy()
        mask_resized = cv2.resize(mask_np, (image.width, image.height), interpolation=cv2.INTER_NEAREST)

        person = np.array(image)
        alpha_channel = (mask_resized * 255).astype(np.uint8)
        rgba_image = np.dstack((person, alpha_channel))
        return Image.fromarray(rgba_image, 'RGBA')

    except Exception as e:
        raise RuntimeError(f"Error in segmenting person: {e}")

# Function to create stereoscopic images with inserted person
def create_stereo_images(stereo_image, person_image, depth_level, x_position, y_position):
    try:
        if not isinstance(stereo_image, np.ndarray):
            raise ValueError("Invalid stereo image format. Please upload a valid image.")

        height, width, _ = stereo_image.shape
        midpoint = width // 2
        left_image = stereo_image[:, :midpoint]
        right_image = stereo_image[:, midpoint:]

        left_image = Image.fromarray(cv2.cvtColor(left_image, cv2.COLOR_BGR2RGBA))
        right_image = Image.fromarray(cv2.cvtColor(right_image, cv2.COLOR_BGR2RGBA))

        disparity_map = {"close": 10, "medium": 1, "far": -10}
        scaling_factors = {"close": 1.2, "medium": 1.0, "far": 0.8}

        if depth_level not in disparity_map:
            raise ValueError("Invalid depth level. Choose from 'close', 'medium', or 'far'.")

        disparity = disparity_map[depth_level]
        person_scale = scaling_factors[depth_level]
        person_image = person_image.resize(
            (int(person_image.width * person_scale), int(person_image.height * person_scale)),
            Image.LANCZOS
        )

        left_x, left_y = x_position, y_position
        right_x = left_x + disparity

        left_image.paste(person_image, (left_x, left_y), person_image)
        right_image.paste(person_image, (right_x, left_y), person_image)

        return left_image, right_image

    except Exception as e:
        raise RuntimeError(f"Error in creating stereo images: {e}")

# Function to create an anaglyph image
def create_anaglyph(left_image, right_image):
    try:
        left_r, _, _ = left_image.convert("RGB").split()
        _, right_g, right_b = right_image.convert("RGB").split()
        anaglyph_image = Image.merge("RGB", (left_r, right_g, right_b))
        return anaglyph_image
    except Exception as e:
        raise RuntimeError(f"Error in creating anaglyph image: {e}")

# Gradio interface function
def process_images(person_image, stereo_image, depth_level, x_position, y_position):
    try:
        person_image = segment_person(person_image)
        stereo_image = np.array(stereo_image)

        left_image, right_image = create_stereo_images(stereo_image, person_image, depth_level, x_position, y_position)
        anaglyph_image = create_anaglyph(left_image, right_image)

        return anaglyph_image
    except Exception as e:
        return f"An error occurred: {e}"


# Gradio app setup
with gr.Blocks() as demo:
    gr.Markdown("# 3D Anaglyph Image Creator")
    gr.Markdown("Upload a person image, select depth, adjust position, and create a 3D anaglyph image.")
    
    with gr.Row():
        person_image = gr.Image(label="Upload Person Image", type="pil")
        stereo_image = gr.Image(label="Upload Stereo Background Image", type="pil")
    
    depth_level = gr.Radio(["close", "medium", "far"], label="Select Depth Level")
    x_position = gr.Slider(0, 2000, step=1, label="X Position")
    y_position = gr.Slider(0, 2000, step=1, label="Y Position")
    
    output_image = gr.Image(label="Generated Anaglyph Image")
    
    generate_button = gr.Button("Generate Anaglyph Image")
    generate_button.click(process_images, [person_image, stereo_image, depth_level, x_position, y_position], output_image)
    # Add reminder text below the button
    gr.Markdown("**Reminder:** After adjusting the depth level or position, click the 'Generate Anaglyph Image' button to update the image.")

# Launch the app
demo.launch()