Update app.py

app.py

@@ -47,9 +47,34 @@ state_dict = torch.load(filepath, map_location="cpu")
 model.load_state_dict(state_dict)
 model = model.to(DEVICE).eval()
 
+def findNormals(gray_depth, format):
+
+    d_im = cv2.cvtColor(cv2.imread(gray_depth).astype(np.uint8), cv2.COLOR_BGR2GRAY)
+    zy, zx = np.gradient(d_im)  
+    # You may also consider using Sobel to get a joint Gaussian smoothing and differentation
+    # to reduce noise
+    #zx = cv2.Sobel(d_im, cv2.CV_64F, 1, 0, ksize=5)     
+    #zy = cv2.Sobel(d_im, cv2.CV_64F, 0, 1, ksize=5)
+
+    if format == "opengl":
+        zy = -zy
+        
+    normal = np.dstack((np.ones_like(d_im), -zy, -zx))
+    n = np.linalg.norm(normal, axis=2)
+    normal[:, :, 0] /= n
+    normal[:, :, 1] /= n
+    normal[:, :, 2] /= n
+
+    # offset and rescale values to be in 0-255
+    normal += 1
+    normal /= 2
+    normal *= 255
+
+    return (normal[:, :, ::-1]).astype(np.uint8)
+
 title = "# Depth Anything V2"
-description = """Official demo for **Depth Anything V2**.
-Please refer to our [paper](https://arxiv.org/abs/2406.09414), [project page](https://depth-anything-v2.github.io), and [github](https://github.com/DepthAnything/Depth-Anything-V2) for more details."""
+description = """Unofficial demo for **Depth Anything V2**.
+Please refer to their [paper](https://arxiv.org/abs/2406.09414), [project page](https://depth-anything-v2.github.io), and [github](https://github.com/DepthAnything/Depth-Anything-V2) for more details."""
 
 @spaces.GPU
 def predict_depth(image):
@@ -61,12 +86,18 @@ with gr.Blocks(css=css) as demo:
     gr.Markdown("### Depth Prediction demo")
 
     with gr.Row():
-        input_image = gr.Image(label="Input Image", type='numpy', elem_id='img-display-input')
-        depth_image_slider = ImageSlider(label="Depth Map with Slider View", elem_id='img-display-output', position=0.5)
-    submit = gr.Button(value="Compute Depth")
-    gray_depth_file = gr.File(label="Grayscale depth map", elem_id="download",)
+        input_image = gr.ImageEditor(label="Input Image", layers=True, sources=('upload', 'clipboard'), show_download_button=True, type="numpy", interactive=True, transforms=(None,), eraser=gr.Eraser(), brush=gr.Brush(default_size=1, colors=['black', '#505050', '#a0a0a0', 'white']), elem_id="img-display-input")
+        with gr.Tab("Depth"):
+            depth_image_slider = ImageSlider(label="Depth Map with Slider View", elem_id='img-display-output', position=0.5)
+            gray_depth_file = gr.File(label="Grayscale depth map", elem_id="download",)
+            submit = gr.Button(value="Compute Depth")
+        with gr.Tab("Normals"):
+            normals_out = gr.Image(label="Normal map", interactive=False)
+            format_normals = gr.Radio(choices=["directx", "opengl"])
+            find_normals = gr.Button("Find normals")
+            find_normals.click(fn=findNormals, inputs=[gray_depth_file, format_normals], outputs=[normals_out])
+    
     raw_file = gr.File(label="16-bit raw output (can be considered as disparity)", elem_id="download",)
-
     cmap = matplotlib.colormaps.get_cmap('Spectral_r')
 
     def on_submit(image):
@@ -92,9 +123,9 @@ with gr.Blocks(css=css) as demo:
 
     submit.click(on_submit, inputs=[input_image], outputs=[depth_image_slider, gray_depth_file, raw_file])
 
-    example_files = os.listdir('assets/examples')
+    example_files = os.listdir('assets/drawn_examples')
     example_files.sort()
-    example_files = [os.path.join('assets/examples', filename) for filename in example_files]
+    example_files = [os.path.join('assets/drawn_examples', filename) for filename in example_files]
     examples = gr.Examples(examples=example_files, inputs=[input_image], outputs=[depth_image_slider, gray_depth_file, raw_file], fn=on_submit)