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)

    # 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!")

    # Option for interactive painting (Placeholder)
    st.subheader("Interactive Depth-based Painting (Demo Placeholder)")
    st.write("This feature will allow users to paint on surfaces based on depth. For now, we can show the depth and its effects.")

# Placeholder for future interactive painting functionality.
# This could be extended with AR-based libraries or Unity integration in a full-scale app.