import gradio as gr
from PIL import Image
import numpy as np
from transformers import AutoModelForImageClassification, AutoFeatureExtractor

# Load the Marigold model
feature_extractor = AutoFeatureExtractor.from_pretrained("prs-eth/marigold-depth-lcm-v1-0")
model = AutoModelForImageClassification.from_pretrained("prs-eth/marigold-depth-lcm-v1-0")

def load_image(image):
    """ Convert uploaded image to grayscale. """
    return image.convert('L')

def compute_gradients(image):
    """ Compute horizontal and vertical gradients of the image. """
    image_array = np.asarray(image, dtype=float)
    x_gradient = np.gradient(image_array, axis=1)
    y_gradient = np.gradient(image_array, axis=0)
    return x_gradient, y_gradient

def create_normal_map(image):
    """ Generate a normal map from an image. """
    image = load_image(image)
    x_grad, y_grad = compute_gradients(image)
    
    # Normalize gradients
    max_grad = max(np.max(np.abs(x_grad)), np.max(np.abs(y_grad)))
    x_grad /= max_grad
    y_grad /= max_grad

    # Calculate z component of the normal (assumed perpendicular to the surface)
    z = np.sqrt(1 - (x_grad ** 2) - (y_grad ** 2))

    # Normalize to 0-255 and format as uint8
    normal_map = np.dstack(((x_grad * 0.5 + 0.5) * 255,
                            (y_grad * 0.5 + 0.5) * 255,
                            (z * 1.0) * 255)).astype(np.uint8)

    return Image.fromarray(normal_map, 'RGB')

def estimate_depth(image):
    """ Estimate depth using the Marigold model. """
    inputs = feature_extractor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    depth_map = outputs.logits  # Adjust if the output tensor is named differently
    return depth_map

def interface(image):
    normal_map = create_normal_map(image)
    grayscale_image = load_image(normal_map)
    depth_map = estimate_depth(image)
    return normal_map, grayscale_image, depth_map

# Set up the Gradio interface
iface = gr.Interface(
    fn=interface,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Image(type="pil", label="Normal Map"),
        gr.Image(type="pil", label="Grayscale Image"),
        gr.Image(type="pil", label="Depth Map")  # Adjust the output type if needed
    ],
    title="Normal Map, Grayscale, and Depth Map Generator",
    description="Upload an image to generate its normal map, a grayscale version, and estimate its depth."
)

iface.launch()