hardware acceleration bottlenecks fixes

Author: OlafenwaMoses

pyproject.toml

@@ -37,6 +37,7 @@ dependencies = [
     "transformers>=5.6.2",
     "pillow>=12.2.0",
     "open3d>=0.18.0",
+    "orjson>=3.9.0",
 ]
 
 # ---------------------------------------------------------------------------

tests/unit/test_coordinate_space.py

@@ -39,16 +39,16 @@
 
 W, H = 20, 16
 FX = FY = 200.0
-CX, CY = 9.5, 7.5   # exact centre of 20 × 16
+CX, CY = 9.5, 7.5  # exact centre of 20 × 16
 
 _STEREO_CFG = StereoDepthAdvancedConfig(
     focal_length=FX,
     cx=CX,
     cy=CY,
-    baseline=100.0,   # 100 mm
+    baseline=100.0,  # 100 mm
     doffs=0.0,
     z_far=1000.0,
-    conf_threshold=0.0,   # accept all pixels
+    conf_threshold=0.0,  # accept all pixels
     occ_threshold=0.0,
 )
 
@@ -95,9 +95,7 @@ def _run_depth(depth: np.ndarray, cfg: DepthEstimationAdvanceConfig = _DEPTH_CFG
     """Run DepthEstimation with a mocked depth array and return the result."""
     h, w = depth.shape
     img = _image_bytes(w, h)
-    with patch.object(
-        DepthEstimationHandler, "_run_depth_anything_checkpoint", return_value=depth
-    ):
+    with patch.object(DepthEstimationHandler, "_run_depth_anything_checkpoint", return_value=depth):
         return DepthEstimation().run(
             DepthEstimationCommand(
                 image_input=img,
@@ -302,8 +300,8 @@ class TestDepthOrdering:
     """Closer objects (larger disparity) must appear at smaller Z."""
 
     def test_stereo_larger_disparity_means_smaller_z(self):
-        disp_near = np.full((H, W), 50.0, dtype=np.float32)   # closer
-        disp_far = np.full((H, W), 10.0, dtype=np.float32)    # further
+        disp_near = np.full((H, W), 50.0, dtype=np.float32)  # closer
+        disp_far = np.full((H, W), 10.0, dtype=np.float32)  # further
         pts_near = np.asarray(_run_stereo(disp_near).point_cloud.points)
         pts_far = np.asarray(_run_stereo(disp_far).point_cloud.points)
         assert pts_near[:, 2].mean() < pts_far[:, 2].mean(), (

tests/unit/test_stereo_depth_handler.py

@@ -213,8 +213,8 @@ def test_depth_image_closer_is_brighter(self, dummy_image_bytes, fake_disp):
         # Use a disparity field where the top-left quadrant has much higher
         # disparity (closer) so we can verify the inversion.
         near_far = fake_disp.copy()
-        near_far[:24, :32] = 200.0   # very close (high disparity)
-        near_far[24:, 32:] = 1.0     # very far   (low disparity)
+        near_far[:24, :32] = 200.0  # very close (high disparity)
+        near_far[24:, 32:] = 1.0  # very far   (low disparity)
         with patch.object(StereoDepthHandler, "_run_s2m2", return_value=near_far):
             result = StereoDepthHandler().handle(
                 StereoDepthCommand(
@@ -272,7 +272,7 @@ def test_raw_depth_values_are_metric_metres(self, dummy_image_bytes):
                     advanced_config=cfg,
                 )
             )
-        expected = (100.0 * 1000.0) / (10.0 * 1000.0)   # = 10.0 m
+        expected = (100.0 * 1000.0) / (10.0 * 1000.0)  # = 10.0 m
         np.testing.assert_allclose(result.raw_depth, expected, rtol=1e-5)
 
     def test_return_point_cloud_requires_open3d(self, dummy_image_bytes, fake_disp):

uv.lock

@@ -17,6 +17,7 @@
 _DEFAULT_DEPTH_SCALE = 1000.0
 _DEFAULT_DEPTH_TRUNC = 10.0
 
+
 class DepthEstimationHandler(CommandHandler[DepthEstimationCommand, DepthEstimationResult]):
     _depth_anything_models = {}
     _model_lock = threading.Lock()
@@ -156,6 +157,12 @@ def _load_depth_anything_checkpoint(self, model_path: str):
             model.load_state_dict(state_dict)
             device = self._torch_device(torch)
             model = model.to(device).eval()
+
+            if device == "cuda":
+                torch.backends.cudnn.benchmark = True
+                torch.backends.cuda.matmul.allow_tf32 = True
+                torch.backends.cudnn.allow_tf32 = True
+
             processor = DPTImageProcessor(
                 size={"height": 518, "width": 518},
                 do_resize=True,

vizion3d/lifting/handlers.py

@@ -6,7 +6,9 @@
 ``app.include_router(lifting_router)``.
 """
 
+import orjson
 from fastapi import APIRouter, File, Form, UploadFile
+from fastapi.responses import Response
 
 from vizion3d.lifting import DepthEstimation, DepthEstimationCommand
 from vizion3d.lifting.defaults import DEFAULT_DEPTH_MODEL_URL
@@ -85,7 +87,7 @@ async def depth_estimation(
         advanced_config=advanced_config,
     )
     result = DepthEstimation().run(cmd)
-    return {
+    payload = {
         "depth_map": result.depth_map,
         "min_depth": result.min_depth,
         "max_depth": result.max_depth,
@@ -101,3 +103,4 @@ async def depth_estimation(
             else None
         ),
     }
+    return Response(content=orjson.dumps(payload), media_type="application/json")

vizion3d/server/rest/depth_estimation.py

@@ -18,7 +18,7 @@ def o3d_depth_image_to_png_bytes(o3d_image) -> bytes:
     """Encode an Open3D uint16 depth image as a PNG byte string."""
     arr = np.asarray(o3d_image)
     buf = io.BytesIO()
-    Image.fromarray(arr).save(buf, format="PNG")
+    Image.fromarray(arr).save(buf, format="PNG", compress_level=1)
     return buf.getvalue()
 
 

vizion3d/server/rest/serialisation.py

@@ -6,7 +6,9 @@
 ``app.include_router(router)``.
 """
 
+import orjson
 from fastapi import APIRouter, File, Form, UploadFile
+from fastapi.responses import Response
 
 from vizion3d.stereo import StereoDepth, StereoDepthCommand
 from vizion3d.stereo.defaults import DEFAULT_STEREO_MODEL_URL
@@ -101,7 +103,7 @@ async def stereo_depth(
         advanced_config=advanced_config,
     )
     result = StereoDepth().run(cmd)
-    return {
+    payload = {
         "depth_map": result.depth_map,
         "disparity_map": result.disparity_map,
         "min_depth": result.min_depth,
@@ -118,3 +120,4 @@ async def stereo_depth(
             else None
         ),
     }
+    return Response(content=orjson.dumps(payload), media_type="application/json")

vizion3d/server/rest/stereo_depth.py

@@ -167,6 +167,11 @@ def _load_s2m2(self, model_path: str):
             model.my_load_state_dict(state_dict)
             model = model.to(device).eval()
 
+            if device == "cuda":
+                torch.backends.cudnn.benchmark = True
+                torch.backends.cuda.matmul.allow_tf32 = True
+                torch.backends.cudnn.allow_tf32 = True
+
             self._stereo_models[model_path] = (model, torch, device)
             return self._stereo_models[model_path]
 
@@ -222,8 +227,8 @@ def _run_s2m2(
             left_t_inp = left_t
             right_t_inp = right_t
 
-        left_pad = image_pad(left_t_inp, 32).to(device)
-        right_pad = image_pad(right_t_inp, 32).to(device)
+        left_pad = image_pad(left_t_inp, 32).to(device, non_blocking=True)
+        right_pad = image_pad(right_t_inp, 32).to(device, non_blocking=True)
 
         device_type = device if isinstance(device, str) else device.type
         if device_type == "cuda":