Update app.py

app.py

@@ -574,6 +574,24 @@ def draw_mask(l, t, v, d, evt: gr.EventData):
     return gr.ImageEditor(value=d)
 
 
+def findNormals():
+    global depths
+    depth = cv2.imread(depths[frame_selected]).astype(np.float32)
+    normals = np.zeros((depth.shape[0], depth.shape[1], 3), dtype=np.float32)
+
+    for x in range(1, depth.shape[1] - 1):
+        for y in range(1, depth.shape[0] - 1):
+            t = np.array([x, y - 1, depth[y - 1, x]], dtype=np.float32)
+            l = np.array([x - 1, y, depth[y, x - 1]], dtype=np.float32)
+            c = np.array([x, y, depth[y, x]], dtype=np.float32)
+
+            d = np.cross(l - c, t - c)
+            n = d / np.linalg.norm(d)
+            normals[y, x] = n
+
+    return (normals * 255 + 127).astype(np.uint8)
+
+
 load_model="""
 async(c, o, b, p, d, n, m)=>{
   var intv = setInterval(function(){
@@ -807,6 +825,10 @@ with gr.Blocks(css=css, js=js) as demo:
               apply.click(fn=apply_mask, inputs=[output_mask, bsize], outputs=[output_mask, output_depth, output_frame])
               reset.click(fn=reset_mask, inputs=None, outputs=[output_mask, output_depth])
 
+            normals_out = gr.Image(label="Normal map", interactive=False)
+            find_normals = gr.Button("Find normals")
+            find_normals.click(fn=findNormals, inputs=None, outputs=[normals_out])
+
         with gr.Column():
             model_type = gr.Dropdown([("small", "vits"), ("base", "vitb"), ("large", "vitl")], type="value", value="vits", label='Model Type')
             processed_video = gr.Video(label="Output Video", format="mp4", interactive=False)