shukdevdatta123 commited on
Commit
2722a13
·
verified ·
1 Parent(s): 448efee

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +60 -0
app.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import requests
3
+ from PIL import Image
4
+ import torch
5
+ from transformers import DepthProImageProcessorFast, DepthProForDepthEstimation
6
+ import numpy as np
7
+ import io
8
+
9
+ # Check if CUDA is available
10
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11
+
12
+ # Load model and processor
13
+ image_processor = DepthProImageProcessorFast.from_pretrained("apple/DepthPro-hf")
14
+ model = DepthProForDepthEstimation.from_pretrained("apple/DepthPro-hf").to(device)
15
+
16
+ # Streamlit App UI
17
+ st.title("Interactive Depth-based AR Painting App")
18
+
19
+ # Upload image through Streamlit UI
20
+ uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
21
+
22
+ if uploaded_file is not None:
23
+ image = Image.open(uploaded_file)
24
+ st.image(image, caption="Uploaded Image", use_column_width=True)
25
+
26
+ # Process image with DepthPro for depth estimation
27
+ inputs = image_processor(images=image, return_tensors="pt").to(device)
28
+ with torch.no_grad():
29
+ outputs = model(**inputs)
30
+
31
+ # Post-process depth output
32
+ post_processed_output = image_processor.post_process_depth_estimation(
33
+ outputs, target_sizes=[(image.height, image.width)],
34
+ )
35
+
36
+ depth = post_processed_output[0]["predicted_depth"]
37
+ depth = (depth - depth.min()) / (depth.max() - depth.min())
38
+ depth = depth * 255.
39
+ depth = depth.detach().cpu().numpy()
40
+ depth_image = Image.fromarray(depth.astype("uint8"))
41
+
42
+ st.subheader("Depth Map")
43
+ st.image(depth_image, caption="Estimated Depth Map", use_column_width=True)
44
+
45
+ # Colorize the depth map to make it more visible
46
+ colormap = depth_image.convert("RGB")
47
+ st.subheader("Colorized Depth Map")
48
+ st.image(colormap, caption="Colorized Depth Map", use_column_width=True)
49
+
50
+ # Option to save depth image
51
+ if st.button('Save Depth Image'):
52
+ depth_image.save('depth_image.png')
53
+ st.success("Depth image saved successfully!")
54
+
55
+ # Option for interactive painting (Placeholder)
56
+ st.subheader("Interactive Depth-based Painting (Demo Placeholder)")
57
+ st.write("This feature will allow users to paint on surfaces based on depth. For now, we can show the depth and its effects.")
58
+
59
+ # Placeholder for future interactive painting functionality.
60
+ # This could be extended with AR-based libraries or Unity integration in a full-scale app.