import streamlit as st
import requests
from PIL import Image
import torch
from transformers import DepthProImageProcessorFast, DepthProForDepthEstimation
import numpy as np
import io

# Check if CUDA is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load model and processor
image_processor = DepthProImageProcessorFast.from_pretrained("apple/DepthPro-hf")
model = DepthProForDepthEstimation.from_pretrained("apple/DepthPro-hf").to(device)

# Streamlit App UI
st.title("Interactive Depth-based AR Painting App")

# Upload image through Streamlit UI
uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Add Generate button
    if st.button("Generate"):
        # Process image with DepthPro for depth estimation
        inputs = image_processor(images=image, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs = model(**inputs)

        # Post-process depth output
        post_processed_output = image_processor.post_process_depth_estimation(
            outputs, target_sizes=[(image.height, image.width)],
        )

        depth = post_processed_output[0]["predicted_depth"]
        depth = (depth - depth.min()) / (depth.max() - depth.min())
        depth = depth * 255.
        depth = depth.detach().cpu().numpy()
        depth_image = Image.fromarray(depth.astype("uint8"))

        st.subheader("Depth Map")
        st.image(depth_image, caption="Estimated Depth Map", use_column_width=True)

        # Colorize the depth map to make it more visible
        colormap = depth_image.convert("RGB")
        st.subheader("Colorized Depth Map")
        st.image(colormap, caption="Colorized Depth Map", use_column_width=True)

        # Option to save depth image
        if st.button('Save Depth Image'):
            depth_image.save('depth_image.png')
            st.success("Depth image saved successfully!")

        # Interactive Painting Feature
        st.subheader("Interactive Depth-based Painting")

        # Prepare for canvas
        canvas = st.canvas(
            width=colormap.width,
            height=colormap.height,
            drawing_mode="freedraw", 
            initial_drawing=colormap,
            key="painting_canvas"
        )

        if canvas.image_data is not None:
            # Convert canvas drawing to an image
            painted_image = Image.fromarray(canvas.image_data.astype(np.uint8))

            # You can combine the depth and painting here
            st.subheader("Canvas with Painting")
            st.image(painted_image, caption="Painting on Depth Map", use_column_width=True)

            # Option to save painted image
            if st.button('Save Painted Image'):
                painted_image.save('painted_image.png')
                st.success("Painted image saved successfully!")
        else:
            st.write("Draw on the canvas to interact with depth!")