Update app.py

app.py

@@ -41,67 +41,86 @@ css = """
 
 set_seed(666)
 
-DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-default_steps = 10
-model = PixelPerfectDepth(sampling_steps=default_steps)
-ckpt_path = hf_hub_download(
-    repo_id="gangweix/Pixel-Perfect-Depth",
-    filename="ppd.pth",
-    repo_type="model"
-)
-state_dict = torch.load(ckpt_path, map_location="cpu")
-model.load_state_dict(state_dict, strict=False)
-model = model.to(DEVICE).eval()
-
-moge_model = MoGeModel.from_pretrained("Ruicheng/moge-2-vitl-normal").to(DEVICE).eval()
-
-
-
-title = "# Pixel-Perfect Depth"
-description = """Official demo for **Pixel-Perfect Depth**.
-Please refer to our [paper](), [project page](https://pixel-perfect-depth.github.io), and [github](https://github.com/gangweix/pixel-perfect-depth) for more details."""
-
-@spaces.GPU
-def predict_depth(image, denoise_steps):
-    depth, resize_image = model.infer_image(image, sampling_steps=denoise_steps)
-    return depth, resize_image
-
-@spaces.GPU
-def predict_moge_depth(image):
-    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    image = torch.tensor(image / 255, dtype=torch.float32, device=DEVICE).permute(2, 0, 1)
-    metric_depth, mask, intrinsics = moge_model.infer(image)
-    metric_depth[~mask] = metric_depth[mask].max()
-    return metric_depth, mask, intrinsics
-
-with gr.Blocks(css=css) as demo:
-    gr.Markdown(title)
-    gr.Markdown(description)
-    gr.Markdown("### Depth Prediction demo")
-
-    with gr.Row():
-        # Left: input image + settings
-        with gr.Column():
-            input_image = gr.Image(label="Input Image", image_mode="RGB", type='numpy', elem_id='img-display-input')
-            with gr.Accordion(label="Settings", open=False):
-                denoise_steps = gr.Slider(label="Denoising Steps", minimum=1, maximum=50, value=10, step=1)
-                apply_filter = gr.Checkbox(label="Apply filter points", value=True)
-            submit_btn = gr.Button(value="Predict Depth")
-
-        # Right: 3D point cloud + depth
-        with gr.Column():
-            with gr.Tabs():
-                with gr.Tab("3D View"):
-                    model_3d = gr.Model3D(display_mode="solid", label="3D Point Map", clear_color=[1,1,1,1], height="60vh")
-                with gr.Tab("Depth"):
-                    depth_map = ImageSlider(label="Depth Map with Slider View", elem_id='img-display-output', position=0.5)
+def main(share=True):
+    print("Initializing Pixel-Perfect Depth Demo...")
     
-    concat_file = gr.File(label="Concatenated visualization (image+depth)", elem_id="image-depth-download")
-    raw_depth_file = gr.File(label="Raw depth output (saved as .npy)", elem_id="download")
-    pcd_file = gr.File(label="Point Cloud (.ply)", elem_id="download-ply")
+    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    default_steps = 10
+    model = PixelPerfectDepth(sampling_steps=default_steps)
+    ckpt_path = hf_hub_download(
+        repo_id="gangweix/Pixel-Perfect-Depth",
+        filename="ppd.pth",
+        repo_type="model"
+    )
+    state_dict = torch.load(ckpt_path, map_location="cpu")
+    model.load_state_dict(state_dict, strict=False)
+    model = model.to(DEVICE).eval()
+
+    moge_model = MoGeModel.from_pretrained("Ruicheng/moge-2-vitl-normal").to(DEVICE).eval()
 
     cmap = matplotlib.colormaps.get_cmap('Spectral')
 
+    title = "# Pixel-Perfect Depth"
+    description = """Official demo for **Pixel-Perfect Depth**.
+    Please refer to our [paper](), [project page](https://pixel-perfect-depth.github.io), and [github](https://github.com/gangweix/pixel-perfect-depth) for more details."""
+
+    @spaces.GPU
+    def predict_depth(image, denoise_steps):
+        depth, resize_image = model.infer_image(image, sampling_steps=denoise_steps)
+        return depth, resize_image
+
+    @spaces.GPU
+    def predict_moge_depth(image):
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        image = torch.tensor(image / 255, dtype=torch.float32, device=DEVICE).permute(2, 0, 1)
+        metric_depth, mask, intrinsics = moge_model.infer(image)
+        metric_depth[~mask] = metric_depth[mask].max()
+        return metric_depth, mask, intrinsics
+
+    with gr.Blocks(css=css) as demo:
+        gr.Markdown(title)
+        gr.Markdown(description)
+        gr.Markdown("### Depth Prediction demo")
+
+        with gr.Row():
+            # Left: input image + settings
+            with gr.Column():
+                input_image = gr.Image(label="Input Image", image_mode="RGB", type='numpy', elem_id='img-display-input')
+                with gr.Accordion(label="Settings", open=False):
+                    denoise_steps = gr.Slider(label="Denoising Steps", minimum=1, maximum=50, value=10, step=1)
+                    apply_filter = gr.Checkbox(label="Apply filter points", value=True)
+                submit_btn = gr.Button(value="Predict Depth")
+
+            # Right: 3D point cloud + depth
+            with gr.Column():
+                with gr.Tabs():
+                    with gr.Tab("3D View"):
+                        model_3d = gr.Model3D(display_mode="solid", label="3D Point Map", clear_color=[1,1,1,1], height="60vh")
+                    with gr.Tab("Depth"):
+                        depth_map = ImageSlider(label="Depth Map with Slider View", elem_id='img-display-output', position=0.5)
+    
+        concat_file = gr.File(label="Concatenated visualization (image+depth)", elem_id="image-depth-download")
+        raw_depth_file = gr.File(label="Raw depth output (saved as .npy)", elem_id="download")
+        pcd_file = gr.File(label="Point Cloud (.ply)", elem_id="download-ply")
+
+        submit_btn.click(
+            on_submit, 
+            inputs=[input_image, denoise_steps, apply_filter], 
+            outputs=[depth_map, model_3d, concat_file, raw_depth_file, pcd_file]
+        )
+
+        example_files = os.listdir('assets/examples')
+        example_files.sort()
+        example_files = [os.path.join('assets/examples', filename) for filename in example_files]
+        examples = gr.Examples(
+            examples=example_files, 
+            inputs=[input_image], 
+            outputs=[depth_map, model_3d, concat_file, raw_depth_file, pcd_file], 
+            fn=on_submit
+        )
+
+    demo.queue().launch(share=share) 
+
     def on_submit(image, denoise_steps, apply_filter):
 
         H, W = image.shape[:2]
@@ -148,22 +167,6 @@ with gr.Blocks(css=css) as demo:
 
         return [(image, colored_depth), tmp_ply.name, tmp_concat.name, tmp_raw_depth.name, tmp_ply.name]
 
-    submit_btn.click(
-        on_submit, 
-        inputs=[input_image, denoise_steps, apply_filter], 
-        outputs=[depth_map, model_3d, concat_file, raw_depth_file, pcd_file]
-    )
-
-    example_files = os.listdir('assets/examples')
-    example_files.sort()
-    example_files = [os.path.join('assets/examples', filename) for filename in example_files]
-    examples = gr.Examples(
-        examples=example_files, 
-        inputs=[input_image], 
-        outputs=[depth_map, model_3d, concat_file, raw_depth_file, pcd_file], 
-        fn=on_submit
-    )
-
 
 if __name__ == '__main__':
-    demo.queue().launch(share=True)
+    main(share=True)