Update app.py

app.py

@@ -240,7 +240,7 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', blur_data
             count += 1
 
         final_vid = create_video(orig_frames, frame_rate, "orig")
-        #final_vid = create_video(depth_frames, frame_rate, "depth")
+        depth_vid = create_video(depth_frames, frame_rate, "depth")
             
         final_zip = zip_files(orig_frames, depth_frames)
         raw_video.release()
@@ -263,7 +263,7 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', blur_data
         else:
             gradient = cv2.imread('./gradient_small.png').astype(np.uint8)
         
-        return final_vid, final_zip, frames, masks[frame_selected], depths #output_path
+        return final_vid, final_zip, frames, masks[frame_selected], depths, depth_vid #output_path
 
 def depth_edges_mask(depth):
     """Returns a mask of edges in the depth map.
@@ -674,6 +674,9 @@ async(c, o, p, d, n, m)=>{
         window.y = 0;
         window.xold = 0;
         window.yold = 0;
+        window.buffer = null;
+        window.ctx = null;
+        window.video = document.getElementById("depth_video").getElementsByTagName("video")[0];
 
         document.getElementById("model3D").getElementsByTagName("canvas")[0].addEventListener('pointermove', function(evt) {
         if (md === true) {
@@ -719,6 +722,85 @@ async(c, o, p, d, n, m)=>{
         document.getElementById("model3D").getElementsByTagName("canvas")[0].addEventListener('pointerout', function() {
           md = false;
         });
+
+        video.addEventListener("playing", function () {
+          const canvas = document.createElement("canvas");
+          canvas.width = video.videoWidth;
+          canvas.height = video.videoHeight;
+          ctx = canvas.getContext("2d", { willReadFrequently: true });
+
+          updateMap();
+        });
+
+        function requestMap() {
+          ctx.drawImage(video, 0, 0, video.videoWidth, video.videoHeight);
+          if (buffer) {
+            applyDisplacementMapFromBuffer(sceneToRender.meshes[sceneToRender.meshes.length-1], buffer, video.videoWidth, video.videoHeight, 0, -1, null, null, true);
+          }
+          buffer = ctx.getImageData(0, 0, video.videoWidth, video.videoHeight).data;
+          applyDisplacementMapFromBuffer(BABYLON.Engine.LastCreatedScene.meshes[BABYLON.Engine.LastCreatedScene.meshes.length-1], buffer, video.videoWidth, video.videoHeight, 0, 1, null, null, true);
+          setTimeout(updateMap, 40);
+        }
+
+        function updateMap() {
+          requestAnimationFrame(requestMap);
+        }
+
+    function applyDisplacementMapFromBuffer(
+        mesh,
+        buffer,
+        heightMapWidth,
+        heightMapHeight,
+        minHeight,
+        maxHeight,
+        uvOffset,
+        uvScale,
+        forceUpdate
+    ) {
+        if (!mesh.isVerticesDataPresent(BABYLON.VertexBuffer.PositionKind) || !mesh.isVerticesDataPresent(BABYLON.VertexBuffer.NormalKind) || !mesh.isVerticesDataPresent(BABYLON.VertexBuffer.UVKind)) {
+            alert("Cannot call applyDisplacementMap: Given mesh is not complete. Position, Normal or UV are missing");
+            return mesh;
+        }
+
+        const positions = mesh.getVerticesData(BABYLON.VertexBuffer.PositionKind, true, true);
+        const normals = mesh.getVerticesData(BABYLON.VertexBuffer.NormalKind);
+        const uvs = mesh.getVerticesData(BABYLON.VertexBuffer.UVKind);
+
+        let position = BABYLON.Vector3.Zero();
+        const normal = BABYLON.Vector3.Zero();
+        const uv = BABYLON.Vector2.Zero();
+
+        uvOffset = uvOffset || BABYLON.Vector2.Zero();
+        uvScale = uvScale || new BABYLON.Vector2(1, 1);
+
+        for (let index = 0; index < positions.length; index += 3) {
+            BABYLON.Vector3.FromArrayToRef(positions, index, position);
+            BABYLON.Vector3.FromArrayToRef(normals, index, normal);
+            BABYLON.Vector2.FromArrayToRef(uvs, (index / 3) * 2, uv);
+
+            // Compute height
+            const u = (Math.abs(uv.x * uvScale.x + (uvOffset.x % 1)) * (heightMapWidth - 1)) % heightMapWidth | 0;
+            const v = (Math.abs(uv.y * uvScale.y + (uvOffset.y % 1)) * (heightMapHeight - 1)) % heightMapHeight | 0;
+
+            const pos = (u + v * heightMapWidth) * 4;
+            const r = buffer[pos] / 255.0;
+            const g = buffer[pos + 1] / 255.0;
+            const b = buffer[pos + 2] / 255.0;
+            const a = buffer[pos + 3] / 255.0;
+
+            const gradient = r * 0.33 + g * 0.33 + b * 0.33;
+            //const gradient = a;
+
+            normal.normalize();
+            normal.scaleInPlace(minHeight + (maxHeight - minHeight) * gradient);
+            position = position.add(normal);
+
+            position.toArray(positions, index);
+        }
+
+        mesh.setVerticesData(BABYLON.VertexBuffer.PositionKind, positions);
+        return mesh;
+    }
         
         if (document.getElementById("model")) {
           document.getElementById("model").appendChild(document.getElementById("model3D"));
@@ -966,6 +1048,7 @@ with gr.Blocks(css=css, js=js, head=head) as demo:
           model_type = gr.Dropdown([("small", "vits"), ("base", "vitb"), ("large", "vitl"), ("giant", "vitg")], type="value", value="vits", label='Model Type')
           processed_video = gr.Video(label="Output Video", format="mp4", interactive=False)
           processed_zip = gr.File(label="Output Archive", interactive=False)
+          depth_video = gr.Video(label="Depth Video", format="mp4", elem_id="depth_video", interactive=False, visible=False)
           result = gr.Model3D(label="3D Mesh", clear_color=[0.5, 0.5, 0.5, 0.0], camera_position=[0, 90, 0], zoom_speed=2.0, pan_speed=2.0, interactive=True, elem_id="model3D")
           with gr.Tab("Blur"):
             chart_c = gr.HTML(elem_id="chart_c", value="""<div id='chart' onpointermove='window.drawLine(event.clientX, event.clientY);' onpointerdown='window.pointerDown(event.clientX, event.clientY);' onpointerup='window.pointerUp();' onpointerleave='window.pointerUp();' onpointercancel='window.pointerUp();' onclick='window.resetLine();'></div>
@@ -1184,12 +1267,12 @@ with gr.Blocks(css=css, js=js, head=head) as demo:
 
         return output_video_path + (json.dumps(locations),)
 
-    submit.click(on_submit, inputs=[input_video, model_type, blur_in, boffset, bsize, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, coords])
+    submit.click(on_submit, inputs=[input_video, model_type, blur_in, boffset, bsize, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, depth_video, coords])
     render.click(None, inputs=[coords, mesh_order, output_frame, output_mask, selected, output_depth], outputs=None, js=load_model)
     render.click(partial(get_mesh), inputs=[output_frame, output_mask, blur_in, load_all], outputs=[result, mesh_order])
 
     example_files = [["./examples/streetview.mp4", "vits", blurin, 1, 32, example_coords]]
-    examples = gr.Examples(examples=example_files, fn=on_submit, cache_examples=True, inputs=[input_video, model_type, blur_in, boffset, bsize, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, coords])
+    examples = gr.Examples(examples=example_files, fn=on_submit, cache_examples=True, inputs=[input_video, model_type, blur_in, boffset, bsize, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, depth_video, coords])
     
 
 if __name__ == '__main__':