remove fp16

app.py

@@ -1,10 +1,15 @@
+"""
+Author: Luigi Piccinelli
+Licensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)
+"""
+
 import gc
 import os
 import shutil
+import sys
 import time
 from datetime import datetime
 from math import pi
-import sys
 
 import gradio as gr
 import numpy as np
@@ -12,12 +17,10 @@ import torch
 import trimesh
 from PIL import Image
 
-
-sys.path.append("unik3d/")
+sys.path.append("./unik3d/")
 
 from unik3d.models import UniK3D
 from unik3d.utils.camera import OPENCV, Fisheye624, Pinhole, Spherical
-from unik3d.utils.visualization import colorize
 
 
 def predictions_to_glb(
@@ -86,7 +89,7 @@ def instantiate_camera(camera_name, params, device):
     return eval(camera_name)(params=torch.tensor(params).float()).to(device)
 
 
-def run_model(target_dir, model_name, camera_name, params):
+def run_model(target_dir, model_name, camera_name, params, efficiency):
 
     print("Instantiating model and camera...")
     model = instantiate_model(model_name)
@@ -102,6 +105,7 @@ def run_model(target_dir, model_name, camera_name, params):
 
     # Perform inference with the model.
     print("Running inference...")
+    model.resolution_level = min(efficiency, 9.0)
     outputs = model.infer(image_tensor, camera=camera, normalize=True)
     outputs["image"] = image_tensor
 
@@ -127,8 +131,8 @@ def gradio_demo(
     hfov,
     mask_black_bg,
     mask_far_points,
+    efficiency
 ):
-    print(target_dir)
     if not os.path.isdir(target_dir) or target_dir == "None":
         return None, "No valid target directory found. Please upload first.", None
 
@@ -138,7 +142,7 @@ def gradio_demo(
     print("Running run_model...")
     params = [fx, fy, cx, cy, k1, k2, k3, k4, k5, k6, t1, t2, hfov]
     with torch.no_grad():
-        outputs = run_model(target_dir, model_name, camera_name, params)
+        outputs = run_model(target_dir, model_name, camera_name, params, efficiency)
 
     # Save predictions
     points = outputs["points"].squeeze().permute(1, 2, 0).cpu().numpy()
@@ -399,8 +403,9 @@ if __name__ == "__main__":
             <li><strong>Upload Your Image:</strong> Use the "Upload Images" panel to provide your input.</li>
             <li><strong>Run:</strong> Click the "Run UniK3D" button to start the 3D estimation process.</li>
             <li><strong>Visualize:</strong> The 3D reconstruction will appear in the viewer on the right. You can rotate, pan, and zoom to explore the model, and download the GLB file.</li>
+            <li><strong>Downstream:</strong> The 3D output can be used as reconstruction or for monocular camera calibration.</li>
         </ol>
-        <p><strong style="color: #ff7e26;">Please note:</strong> <span style="color: #ff7e26; font-weight: bold;">Our model runs on CPU on HuggingFace Space. Actual inference is less than 100ms second per image on consumer-level GPUs. Web-based 3D pointcloud visualization may be slow due to Gradio's rendering. For faster visualization, use a local machine to run our demo from our <a href="https://github.com/lpiccinelli-eth/UniK3D">GitHub repository</a>. </span></p>
+        <p><strong style="color: #ff7e26;">Please note:</strong> <span style="color: #ff7e26; font-weight: bold;">Our model runs on CPU on HuggingFace Space. Actual inference is less than 100ms second per image on consumer-level GPUs, on Spaces will take between 20s and 90s, depending on the "Speed-Resoltion Tradeoff" chosen. Web-based 3D pointcloud visualization may be slow due to Gradio's rendering. For faster visualization, use a local machine to run our demo from our <a href="https://github.com/lpiccinelli-eth/UniK3D">GitHub repository</a>. </span></p>
         </div>
         """
         )
@@ -409,7 +414,7 @@ if __name__ == "__main__":
 
         with gr.Row():
             with gr.Column():
-                camera_dropdown = gr.Dropdown(
+                camera_model = gr.Dropdown(
                     choices=[
                         "Predicted",
                         "Pinhole",
@@ -419,13 +424,14 @@ if __name__ == "__main__":
                     ],
                     label="Input Camera",
                 )
-                model_dropdown = gr.Dropdown(
+                model_size = gr.Dropdown(
                     choices=["Large", "Base", "Small"], label="Utilized Model"
                 )
                 mask_black_bg = gr.Checkbox(
                     label="Filter Black Background", value=False
                 )
                 mask_far_points = gr.Checkbox(label="Filter Far Points", value=False)
+                efficiency = gr.Slider(0, 10, step=1, value=10, label="Speed-Resolution Tradeoff", info="Lower is faster and Higher is more detailed")
 
             with gr.Column():
                 fx = gr.Number(label="Focal length x", value=500.0, visible=False)
@@ -498,6 +504,7 @@ if __name__ == "__main__":
                 0.0,
                 True,
                 False,
+                10.0,
             ],
             [
                 "assets/demo/naruto.jpg",
@@ -518,9 +525,10 @@ if __name__ == "__main__":
                 0.0,
                 False,
                 False,
+                10.0,
             ],
             [
-                "assets/demo/bears.jpg",
+                "assets/demo/bears.png",
                 "Large",
                 "Predicted",
                 0.0,
@@ -538,6 +546,7 @@ if __name__ == "__main__":
                 0.0,
                 True,
                 False,
+                10.0,
             ],
             [
                 "assets/demo/berzirk.jpg",
@@ -558,6 +567,7 @@ if __name__ == "__main__":
                 0.0,
                 True,
                 False,
+                10.0,
             ],
             [
                 "assets/demo/luke.webp",
@@ -578,6 +588,7 @@ if __name__ == "__main__":
                 0.0,
                 False,
                 False,
+                10.0,
             ],
             [
                 "assets/demo/equirectangular.jpg",
@@ -598,6 +609,7 @@ if __name__ == "__main__":
                 360.0,
                 False,
                 False,
+                10.0,
             ],
             [
                 "assets/demo/venice.jpg",
@@ -618,6 +630,7 @@ if __name__ == "__main__":
                 360.0,
                 False,
                 True,
+                10.0,
             ],
             [
                 "assets/demo/dl3dv.png",
@@ -638,9 +651,10 @@ if __name__ == "__main__":
                 0.0,
                 False,
                 False,
+                10.0,
             ],
             [
-                "assets/demo/scannet.jpg",
+                "assets/demo/scannet.png",
                 "Large",
                 "Fisheye624",
                 791.90869140625,
@@ -658,6 +672,7 @@ if __name__ == "__main__":
                 0.0,
                 False,
                 False,
+                10.0,
             ],
         ]
 
@@ -680,6 +695,7 @@ if __name__ == "__main__":
             hfov,
             mask_black_bg,
             mask_far_points,
+            efficiency
         ):
             target_dir, image_path = handle_uploads(input_image)
             glbfile, log_msg, prediction_save_path = gradio_demo(
@@ -701,6 +717,7 @@ if __name__ == "__main__":
                 hfov,
                 mask_black_bg,
                 mask_far_points,
+                efficiency
             )
             return (
                 glbfile,
@@ -716,8 +733,8 @@ if __name__ == "__main__":
             examples=examples,
             inputs=[
                 input_image,
-                model_dropdown,
-                camera_dropdown,
+                model_size,
+                camera_model,
                 fx,
                 fy,
                 cx,
@@ -733,6 +750,7 @@ if __name__ == "__main__":
                 hfov,
                 mask_black_bg,
                 mask_far_points,
+                efficiency
             ],
             outputs=[reconstruction_output, log_output, reconstruction_npy],
             fn=example_pipeline,
@@ -746,8 +764,8 @@ if __name__ == "__main__":
             fn=gradio_demo,
             inputs=[
                 target_dir_output,
-                model_dropdown,
-                camera_dropdown,
+                model_size,
+                camera_model,
                 fx,
                 fy,
                 cx,
@@ -763,6 +781,7 @@ if __name__ == "__main__":
                 hfov,
                 mask_black_bg,
                 mask_far_points,
+                efficiency
             ],
             outputs=[reconstruction_output, log_output, reconstruction_npy],
         ).then(
@@ -788,13 +807,10 @@ if __name__ == "__main__":
         )
 
         # Dynamically update intrinsic parameter visibility when camera selection changes.
-        camera_dropdown.change(
+        camera_model.change(
             fn=update_parameters,
-            inputs=camera_dropdown,
+            inputs=camera_model,
             outputs=[fx, fy, cx, cy, k1, k2, k3, k4, k5, k6, t1, t2, hfov],
         )
 
-        # demo.queue(max_size=20).launch(show_error=True, share=False, ssr_mode=False)
-        demo.launch(
-            show_error=True,
-        )
+        demo.queue(max_size=20).launch(show_error=True, share=True, ssr_mode=False)

gradio_demo.py

@@ -1,3 +1,8 @@
+"""
+Author: Luigi Piccinelli
+Licensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)
+"""
+
 import gc
 import os
 import shutil
@@ -13,7 +18,6 @@ from PIL import Image
 
 from unik3d.models import UniK3D
 from unik3d.utils.camera import OPENCV, Fisheye624, Pinhole, Spherical
-from unik3d.utils.visualization import colorize
 
 
 def predictions_to_glb(
@@ -82,7 +86,7 @@ def instantiate_camera(camera_name, params, device):
     return eval(camera_name)(params=torch.tensor(params).float()).to(device)
 
 
-def run_model(target_dir, model_name, camera_name, params):
+def run_model(target_dir, model_name, camera_name, params, efficiency):
 
     print("Instantiating model and camera...")
     model = instantiate_model(model_name)
@@ -98,6 +102,7 @@ def run_model(target_dir, model_name, camera_name, params):
 
     # Perform inference with the model.
     print("Running inference...")
+    model.resolution_level = min(efficiency, 9.0)
     outputs = model.infer(image_tensor, camera=camera, normalize=True)
     outputs["image"] = image_tensor
 
@@ -123,8 +128,8 @@ def gradio_demo(
     hfov,
     mask_black_bg,
     mask_far_points,
+    efficiency
 ):
-    print(target_dir)
     if not os.path.isdir(target_dir) or target_dir == "None":
         return None, "No valid target directory found. Please upload first.", None
 
@@ -134,7 +139,7 @@ def gradio_demo(
     print("Running run_model...")
     params = [fx, fy, cx, cy, k1, k2, k3, k4, k5, k6, t1, t2, hfov]
     with torch.no_grad():
-        outputs = run_model(target_dir, model_name, camera_name, params)
+        outputs = run_model(target_dir, model_name, camera_name, params, efficiency)
 
     # Save predictions
     points = outputs["points"].squeeze().permute(1, 2, 0).cpu().numpy()
@@ -395,8 +400,9 @@ if __name__ == "__main__":
             <li><strong>Upload Your Image:</strong> Use the "Upload Images" panel to provide your input.</li>
             <li><strong>Run:</strong> Click the "Run UniK3D" button to start the 3D estimation process.</li>
             <li><strong>Visualize:</strong> The 3D reconstruction will appear in the viewer on the right. You can rotate, pan, and zoom to explore the model, and download the GLB file.</li>
+            <li><strong>Downstream:</strong> The 3D output can be used as reconstruction or for monocular camera calibration.</li>
         </ol>
-        <p><strong style="color: #ff7e26;">Please note:</strong> <span style="color: #ff7e26; font-weight: bold;">Our model runs on CPU on HuggingFace Space. Actual inference is less than 100ms second per image on consumer-level GPUs. Web-based 3D pointcloud visualization may be slow due to Gradio's rendering. For faster visualization, use a local machine to run our demo from our <a href="https://github.com/lpiccinelli-eth/UniK3D">GitHub repository</a>. </span></p>
+        <p><strong style="color: #ff7e26;">Please note:</strong> <span style="color: #ff7e26; font-weight: bold;">Our model runs on CPU on HuggingFace Space. Actual inference is less than 100ms second per image on consumer-level GPUs, on Spaces will take between 20s and 90s, depending on the "Speed-Resoltion Tradeoff" chosen. Web-based 3D pointcloud visualization may be slow due to Gradio's rendering. For faster visualization, use a local machine to run our demo from our <a href="https://github.com/lpiccinelli-eth/UniK3D">GitHub repository</a>. </span></p>
         </div>
         """
         )
@@ -405,7 +411,7 @@ if __name__ == "__main__":
 
         with gr.Row():
             with gr.Column():
-                camera_dropdown = gr.Dropdown(
+                camera_model = gr.Dropdown(
                     choices=[
                         "Predicted",
                         "Pinhole",
@@ -415,13 +421,14 @@ if __name__ == "__main__":
                     ],
                     label="Input Camera",
                 )
-                model_dropdown = gr.Dropdown(
+                model_size = gr.Dropdown(
                     choices=["Large", "Base", "Small"], label="Utilized Model"
                 )
                 mask_black_bg = gr.Checkbox(
                     label="Filter Black Background", value=False
                 )
                 mask_far_points = gr.Checkbox(label="Filter Far Points", value=False)
+                efficiency = gr.Slider(0, 10, step=1, value=10, label="Speed-Resolution Tradeoff", info="Lower is faster and Higher is more detailed")
 
             with gr.Column():
                 fx = gr.Number(label="Focal length x", value=500.0, visible=False)
@@ -494,6 +501,7 @@ if __name__ == "__main__":
                 0.0,
                 True,
                 False,
+                10.0,
             ],
             [
                 "assets/demo/naruto.jpg",
@@ -514,6 +522,7 @@ if __name__ == "__main__":
                 0.0,
                 False,
                 False,
+                10.0,
             ],
             [
                 "assets/demo/bears.png",
@@ -534,6 +543,7 @@ if __name__ == "__main__":
                 0.0,
                 True,
                 False,
+                10.0,
             ],
             [
                 "assets/demo/berzirk.jpg",
@@ -554,6 +564,7 @@ if __name__ == "__main__":
                 0.0,
                 True,
                 False,
+                10.0,
             ],
             [
                 "assets/demo/luke.webp",
@@ -574,6 +585,7 @@ if __name__ == "__main__":
                 0.0,
                 False,
                 False,
+                10.0,
             ],
             [
                 "assets/demo/equirectangular.jpg",
@@ -594,6 +606,7 @@ if __name__ == "__main__":
                 360.0,
                 False,
                 False,
+                10.0,
             ],
             [
                 "assets/demo/venice.jpg",
@@ -614,6 +627,7 @@ if __name__ == "__main__":
                 360.0,
                 False,
                 True,
+                10.0,
             ],
             [
                 "assets/demo/dl3dv.png",
@@ -634,6 +648,7 @@ if __name__ == "__main__":
                 0.0,
                 False,
                 False,
+                10.0,
             ],
             [
                 "assets/demo/scannet.png",
@@ -654,6 +669,7 @@ if __name__ == "__main__":
                 0.0,
                 False,
                 False,
+                10.0,
             ],
         ]
 
@@ -676,6 +692,7 @@ if __name__ == "__main__":
             hfov,
             mask_black_bg,
             mask_far_points,
+            efficiency
         ):
             target_dir, image_path = handle_uploads(input_image)
             glbfile, log_msg, prediction_save_path = gradio_demo(
@@ -697,6 +714,7 @@ if __name__ == "__main__":
                 hfov,
                 mask_black_bg,
                 mask_far_points,
+                efficiency
             )
             return (
                 glbfile,
@@ -712,8 +730,8 @@ if __name__ == "__main__":
             examples=examples,
             inputs=[
                 input_image,
-                model_dropdown,
-                camera_dropdown,
+                model_size,
+                camera_model,
                 fx,
                 fy,
                 cx,
@@ -729,6 +747,7 @@ if __name__ == "__main__":
                 hfov,
                 mask_black_bg,
                 mask_far_points,
+                efficiency
             ],
             outputs=[reconstruction_output, log_output, reconstruction_npy],
             fn=example_pipeline,
@@ -742,8 +761,8 @@ if __name__ == "__main__":
             fn=gradio_demo,
             inputs=[
                 target_dir_output,
-                model_dropdown,
-                camera_dropdown,
+                model_size,
+                camera_model,
                 fx,
                 fy,
                 cx,
@@ -759,6 +778,7 @@ if __name__ == "__main__":
                 hfov,
                 mask_black_bg,
                 mask_far_points,
+                efficiency
             ],
             outputs=[reconstruction_output, log_output, reconstruction_npy],
         ).then(
@@ -784,13 +804,10 @@ if __name__ == "__main__":
         )
 
         # Dynamically update intrinsic parameter visibility when camera selection changes.
-        camera_dropdown.change(
+        camera_model.change(
             fn=update_parameters,
-            inputs=camera_dropdown,
+            inputs=camera_model,
             outputs=[fx, fy, cx, cy, k1, k2, k3, k4, k5, k6, t1, t2, hfov],
         )
 
-        # demo.queue(max_size=20).launch(show_error=True, share=False, ssr_mode=False)
-        demo.launch(
-            show_error=True,
-        )
+        demo.queue(max_size=20).launch(show_error=True, share=True, ssr_mode=False)

unik3d/models/unik3d.py

@@ -22,6 +22,7 @@ from unik3d.utils.distributed import is_main_process
 from unik3d.utils.misc import get_params, last_stack, match_gt
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+ENABLED = torch.cuda.is_available()
 
 
 def orthonormal_init(num_tokens, dims):
@@ -146,7 +147,7 @@ class UniK3D(
         )
 
         # compute loss!
-        inputs["radius"] = torch.norm(pts_gt, dim=1, keepdim=True)
+        inputs["distance"] = torch.norm(pts_gt, dim=1, keepdim=True)
         inputs["points"] = pts_gt
         inputs["depth_mask"] = mask
         losses = self.compute_losses(outputs, inputs, image_metas)
@@ -241,8 +242,8 @@ class UniK3D(
         ).reshape(B)
         loss = self.losses["depth"]
         depth_losses = loss(
-            outputs["depth"],
-            target=inputs["depth"],
+            outputs["distance"],
+            target=inputs["distance"],
             mask=inputs["depth_mask"].clone(),
             si=si,
         )
@@ -264,6 +265,7 @@ class UniK3D(
             target_pred=outputs["depth"],
             mask=inputs["depth_mask"].clone(),
         )
+        print(conf_losses, camera_losses, depth_losses)
         losses["opt"][loss.name + "_conf"] = loss.weight * conf_losses.mean()
         losses_to_be_computed.remove("confidence")
 
@@ -274,7 +276,7 @@ class UniK3D(
         return losses
 
     @torch.no_grad()
-    @torch.autocast(device_type=DEVICE, enabled=True, dtype=torch.float16)
+    @torch.autocast(device_type=DEVICE, enabled=ENABLED, dtype=torch.float16)
     def infer(
         self,
         rgb: torch.Tensor,