Spaces:
Running
Running
| 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!") |