Spaces:
Running
Running
File size: 2,317 Bytes
2722a13 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
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.
|