akhaliq HF staff commited on
Commit
b19928f
1 Parent(s): e0f06f0

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +50 -0
app.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from PIL import Image
3
+ import depth_pro
4
+ import numpy as np
5
+ import matplotlib.pyplot as plt
6
+
7
+ # Load model and preprocessing transform
8
+ model, transform = depth_pro.create_model_and_transforms()
9
+ model.eval()
10
+
11
+ def predict_depth(input_image):
12
+ # Preprocess the image
13
+ result = depth_pro.load_rgb(input_image.name)
14
+ image = result[0]
15
+ f_px = result[-1] # Assuming f_px is the last item in the returned tuple
16
+ image = transform(image)
17
+
18
+ # Run inference
19
+ prediction = model.infer(image, f_px=f_px)
20
+ depth = prediction["depth"] # Depth in [m]
21
+ focallength_px = prediction["focallength_px"] # Focal length in pixels
22
+
23
+ # Normalize depth for visualization
24
+ depth_normalized = (depth - np.min(depth)) / (np.max(depth) - np.min(depth))
25
+
26
+ # Create a color map
27
+ plt.figure(figsize=(10, 10))
28
+ plt.imshow(depth_normalized, cmap='viridis')
29
+ plt.colorbar(label='Depth')
30
+ plt.title('Predicted Depth Map')
31
+ plt.axis('off')
32
+
33
+ # Save the plot to a file
34
+ output_path = "depth_map.png"
35
+ plt.savefig(output_path)
36
+ plt.close()
37
+
38
+ return output_path, f"Focal length: {focallength_px:.2f} pixels"
39
+
40
+ # Create Gradio interface
41
+ iface = gr.Interface(
42
+ fn=predict_depth,
43
+ inputs=gr.Image(type="filepath"),
44
+ outputs=[gr.Image(type="filepath", label="Depth Map"), gr.Textbox(label="Focal Length")],
45
+ title="Depth Prediction Demo",
46
+ description="Upload an image to predict its depth map and focal length."
47
+ )
48
+
49
+ # Launch the interface
50
+ iface.launch()