Spaces:
Running
Running
File size: 3,065 Bytes
2722a13 d959f41 331c43f |
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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
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!") |