diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 228ea4e..ea31b29 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -5,6 +5,9 @@ on: branches: [main] pull_request: +permissions: + contents: read + # Cancel in-progress runs for the same branch on new pushes. # Never cancel runs on main (every merge must be fully tested). concurrency: @@ -78,11 +81,25 @@ jobs: uses: actions/cache@v4 with: path: ${{ env.VIZION3D_MODEL_CACHE }} - # Key is tied to the model filename — bump if the model changes - key: depth-anything-v2-vitb-pth + # Key is tied to the model filenames — bump if either model changes + key: depth-anything-v2-vitb-stereo-depth-s2m2-l-pth + + - name: Show integration timing limits + env: + VIZION3D_TEST_DEPTH_COLD_LIMIT: ${{ vars.VIZION3D_TEST_DEPTH_COLD_LIMIT || vars.VIZION3D_TEST_COLD_LIMIT || '10' }} + VIZION3D_TEST_DEPTH_WARM_LIMIT: ${{ vars.VIZION3D_TEST_DEPTH_WARM_LIMIT || vars.VIZION3D_TEST_WARM_LIMIT || '1' }} + VIZION3D_TEST_STEREO_COLD_LIMIT: ${{ vars.VIZION3D_TEST_STEREO_COLD_LIMIT || '60' }} + VIZION3D_TEST_STEREO_WARM_LIMIT: ${{ vars.VIZION3D_TEST_STEREO_WARM_LIMIT || '5' }} + run: | + echo "VIZION3D_TEST_DEPTH_COLD_LIMIT=$VIZION3D_TEST_DEPTH_COLD_LIMIT" + echo "VIZION3D_TEST_DEPTH_WARM_LIMIT=$VIZION3D_TEST_DEPTH_WARM_LIMIT" + echo "VIZION3D_TEST_STEREO_COLD_LIMIT=$VIZION3D_TEST_STEREO_COLD_LIMIT" + echo "VIZION3D_TEST_STEREO_WARM_LIMIT=$VIZION3D_TEST_STEREO_WARM_LIMIT" - name: Run integration tests env: - VIZION3D_TEST_COLD_LIMIT: ${{ vars.VIZION3D_TEST_COLD_LIMIT }} - VIZION3D_TEST_WARM_LIMIT: ${{ vars.VIZION3D_TEST_WARM_LIMIT }} + VIZION3D_TEST_DEPTH_COLD_LIMIT: ${{ vars.VIZION3D_TEST_DEPTH_COLD_LIMIT || vars.VIZION3D_TEST_COLD_LIMIT || '10' }} + VIZION3D_TEST_DEPTH_WARM_LIMIT: ${{ vars.VIZION3D_TEST_DEPTH_WARM_LIMIT || vars.VIZION3D_TEST_WARM_LIMIT || '1' }} + VIZION3D_TEST_STEREO_COLD_LIMIT: ${{ vars.VIZION3D_TEST_STEREO_COLD_LIMIT || '60' }} + VIZION3D_TEST_STEREO_WARM_LIMIT: ${{ vars.VIZION3D_TEST_STEREO_WARM_LIMIT || '5' }} run: uv run pytest tests/integration/ -v diff --git a/.github/workflows/enforce-release-branch.yml b/.github/workflows/enforce-release-branch.yml index b60d963..7f291e5 100644 --- a/.github/workflows/enforce-release-branch.yml +++ b/.github/workflows/enforce-release-branch.yml @@ -5,6 +5,9 @@ on: branches: - release +permissions: + contents: read + jobs: check-source: name: Source branch check diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 0f3918c..9d40be9 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -6,6 +6,9 @@ on: workflow_dispatch: inputs: version: + +permissions: + contents: read description: "Release version (e.g. 1.2.0)" required: true type: string diff --git a/docs/api/lifting.md b/docs/api/lifting.md index e46aa58..521509d 100644 --- a/docs/api/lifting.md +++ b/docs/api/lifting.md @@ -20,6 +20,14 @@ Input contract for the depth estimation task. All inference parameters are decla --- +## DepthEstimationAdvanceConfig + +Camera intrinsics and depth range settings. Pass an instance of this model as `advanced_config` on `DepthEstimationCommand` to override the PrimeSense defaults used for point cloud unprojection. + +::: vizion3d.lifting.models.DepthEstimationAdvanceConfig + +--- + ## DepthEstimationResult Output contract returned by `DepthEstimation.run()`. All fields are always present; optional geometry fields are `None` when the corresponding `return_*` flag was not set. diff --git a/docs/features/depth_estimation.md b/docs/features/depth_estimation.md index 1e3da1c..8c1c00c 100644 --- a/docs/features/depth_estimation.md +++ b/docs/features/depth_estimation.md @@ -9,6 +9,15 @@ Depth estimation predicts the per-pixel distance from the camera for every pixel ## Model backends +Default checkpoint download: +[depth_anything_v2_vitb.pth](https://github.com/OlafenwaMoses/vizion3D/releases/download/essentials-v1/depth_anything_v2_vitb.pth) + +```bash +curl -L \ + https://github.com/OlafenwaMoses/vizion3D/releases/download/essentials-v1/depth_anything_v2_vitb.pth \ + -o depth_anything_v2_vitb.pth +``` + | Value | What happens | |---|---| | *(default)* | Downloads the vizion3D release checkpoint (`depth_anything_v2_vitb.pth`) to `~/.cache/vizion3d/models/` on first use, then loads it directly | @@ -32,6 +41,7 @@ Set `VIZION3D_MODEL_CACHE` in your environment to change the default cache direc | `return_depth_image` | `bool` | No | `False` | If `True`, the result includes a 16-bit grayscale Open3D Image of the depth map. | | `return_point_cloud` | `bool` | No | `False` | If `True`, the result includes an Open3D PointCloud unprojected from the RGB-D image. | | `return_mesh` | `bool` | No | `False` | If `True`, the result includes an Open3D TriangleMesh reconstructed from the point cloud via ball-pivoting. | +| `advanced_config` | `DepthEstimationAdvanceConfig` | No | PrimeSense defaults | Camera intrinsics and depth range settings. See [Advanced config](#10-advanced-config-camera-intrinsics-depth-range) below. | --- @@ -46,7 +56,7 @@ Set `VIZION3D_MODEL_CACHE` in your environment to change the default cache direc | `max_depth` | `float` | Yes | Maximum value in `depth_map`. Guaranteed `max_depth >= min_depth`. | | `backend_used` | `str` | Yes | Resolved model identifier that processed the request (local file path). | | `depth_image` | `open3d.geometry.Image \| None` | When `return_depth_image=True` | 16-bit grayscale image, dtype `uint16`, shape `(H, W)`. The full 0–65535 range maps to `[min_depth, max_depth]`. | -| `point_cloud` | `open3d.geometry.PointCloud \| None` | When `return_point_cloud=True` | Coloured 3D point cloud unprojected from the RGB-D image using PrimeSense default intrinsics. Coordinates are in metres. | +| `point_cloud` | `open3d.geometry.PointCloud \| None` | When `return_point_cloud=True` | Coloured 3D point cloud unprojected from the RGB-D image using the intrinsics in `advanced_config`. Coordinates are in metres. | | `mesh` | `open3d.geometry.TriangleMesh \| None` | When `return_mesh=True` | Triangle mesh surface reconstructed from the point cloud via ball-pivoting. Includes vertex colours. | | `point_cloud_scale` | `float` | Yes | Scale factor: multiply any distance measured between two points in the point cloud by this value to get the equivalent distance in metres. Always `1.0` — Open3D produces point cloud coordinates directly in metres. | @@ -219,7 +229,10 @@ print(f"Backend: {result.backend_used}") # Remote checkpoint URL (downloaded and cached on first use) cmd = DepthEstimationCommand( image_input="scene.png", - model_backend="https://example.com/weights/depth_anything_v2_vits.pth", + model_backend=( + "https://github.com/OlafenwaMoses/vizion3D/releases/download/" + "essentials-v1/depth_anything_v2_vitb.pth" + ), ) result = DepthEstimation().run(cmd) print(f"Backend: {result.backend_used}") @@ -229,7 +242,7 @@ print(f"Backend: {result.backend_used}") ## 8. REST API -Start the server: +Start the server with all REST features enabled: **pip / Poetry** ```bash @@ -241,6 +254,31 @@ vizion3d-serve-rest uv run vizion3d-serve-rest ``` +To preload a depth-estimation checkpoint into memory at startup, pass +`--depth_model`. This also enables the depth-estimation endpoint. If this flag +is omitted, the default vizion3D release model is downloaded on first inference +and cached under `~/.cache/vizion3d/models/`. + +```bash +uv run vizion3d-serve-rest --depth_model /models/depth_anything_v2_vitb.pth +``` + +The REST server can also expose only selected features. If none of +`--depth_estimation`, `--stereo_depth`, `--depth_model`, or `--stereo_model` is +provided, all features are enabled. If any of those flags is provided, only the +selected features are enabled. A model path flag selects and preloads its +feature: + +```bash +# Only POST /lifting/depth-estimation +uv run vizion3d-serve-rest --depth_estimation + +# Only depth estimation, with the model loaded before the first request +uv run vizion3d-serve-rest \ + --depth_estimation \ + --depth_model /models/depth_anything_v2_vitb.pth +``` + Send a request with `multipart/form-data`: ```bash @@ -294,9 +332,38 @@ print(f"Backend : {response.backend_used}") --- +## 10. Advanced config: camera intrinsics & depth range + +`DepthEstimationAdvanceConfig` lets you supply the actual camera intrinsics and depth range for your sensor, replacing the built-in PrimeSense defaults. This is required for accurate metric 3D geometry when your camera is not a 640×480 PrimeSense sensor. + +```python +from vizion3d.lifting import ( + DepthEstimation, + DepthEstimationAdvanceConfig, + DepthEstimationCommand, +) + +result = DepthEstimation().run( + DepthEstimationCommand( + image_input="scene.png", + return_point_cloud=True, + advanced_config=DepthEstimationAdvanceConfig( + fx=909.15, + fy=908.48, + cx=640.0, + cy=360.0, + depth_trunc=6.0, + ), + ) +) +``` + +The same config is available in the REST and gRPC entry points. See [Advanced Config](depth_estimation_advanced_config.md) for the full field reference, formulas, entry-point examples, and camera presets. + +--- + ## Known limitations - **Relative depth only** — the default monocular backend produces relative (inverse) depth, not metric depth. Point cloud distances are internally consistent but not calibrated to real-world scale without a known reference distance. -- **Fixed camera intrinsics** — point cloud unprojection uses `PrimeSenseDefault` (640×480) intrinsics regardless of input image resolution. For accurate metric geometry, supply real camera intrinsics. - **Ball-pivoting mesh quality** — the mesh reconstructor works best on dense, evenly sampled point clouds. Sparse or noisy clouds may produce gaps or missing faces. - **Python 3.12 required for Open3D** — `return_depth_image`, `return_point_cloud`, and `return_mesh` require Open3D, which currently only supports Python 3.12 in this project. diff --git a/docs/features/depth_estimation_advanced_config.md b/docs/features/depth_estimation_advanced_config.md new file mode 100644 index 0000000..c332c43 --- /dev/null +++ b/docs/features/depth_estimation_advanced_config.md @@ -0,0 +1,339 @@ +# Advanced Config: Camera Intrinsics & Depth Range + +`DepthEstimationAdvanceConfig` lets you override the camera intrinsics and depth range parameters that control how a raw depth map is lifted into a 3D point cloud. Without it, vizion3d uses built-in PrimeSense defaults; with it, you can match your actual camera and scene requirements precisely. + +--- + +## Background: the pinhole camera model + +Every point in a point cloud is computed by inverting the pinhole camera projection. Given a pixel at image coordinates `(u, v)` with a depth value `d` (in metres), its 3D position `(X, Y, Z)` is: + +``` +Z = d +X = (u - cx) * d / fx +Y = (v - cy) * d / fy +``` + +All four intrinsic parameters — `fx`, `fy`, `cx`, `cy` — appear in this formula. Getting them wrong produces a point cloud that is geometrically distorted: correct topology but wrong angles, skewed shapes, or objects that appear compressed or stretched. + +--- + +## Config fields + +### `fx` — horizontal focal length (pixels) + +**Default:** `525.0` + +The horizontal focal length of the camera in pixels. It is the product of the physical focal length (mm) and the horizontal pixel density (pixels/mm). A larger `fx` means the camera has a narrower horizontal field of view; the same scene width maps to fewer pixels. + +**Effect on the point cloud:** `fx` controls the horizontal spread of 3D points. If `fx` is too small, the point cloud is horizontally compressed. If too large, it is horizontally stretched. + +**How to find it:** Use your camera's calibration matrix `K[0][0]`, or compute it from the horizontal field of view `FoV_h`: + +``` +fx = (image_width / 2) / tan(FoV_h / 2) +``` + +--- + +### `fy` — vertical focal length (pixels) + +**Default:** `525.0` + +The vertical focal length in pixels. For cameras with square pixels, `fy ≈ fx`. Cameras with non-square sensors may have `fy ≠ fx`. + +**Effect on the point cloud:** Controls vertical spread analogously to `fx`. Incorrect `fy` produces vertically compressed or stretched geometry. + +**How to find it:** `K[1][1]` from the calibration matrix, or: + +``` +fy = (image_height / 2) / tan(FoV_v / 2) +``` + +--- + +### `cx` — horizontal principal point (pixels) + +**Default:** `319.5` + +The horizontal image coordinate of the optical axis — ideally the exact centre of the sensor. For a 640-wide image the ideal value is `319.5`; for a 1920-wide image it is typically near `959.5`. + +**Effect on the point cloud:** Shifts the entire point cloud left or right. A wrong `cx` makes the scene appear to be viewed from an off-centre vantage point, introducing a lateral tilt. + +--- + +### `cy` — vertical principal point (pixels) + +**Default:** `239.5` + +The vertical image coordinate of the optical axis. For a 480-tall image the ideal value is `239.5`. + +**Effect on the point cloud:** Shifts the entire point cloud up or down. Like `cx`, an incorrect value introduces a tilt — vertical in this case. + +--- + +### `depth_scale` — depth value scale factor + +**Default:** `1000.0` + +The divisor applied to the raw uint16 depth buffer before passing depth values to Open3D's `RGBDImage.create_from_color_and_depth`. Open3D divides the stored integer depth by `depth_scale` to obtain a value in metres. The default of `1000.0` means the uint16 range `[0, 65535]` maps to `[0, 65.535]` metres. + +**Effect on the point cloud:** Changing `depth_scale` rescales all Z values (and therefore X/Y values, since `X = (u - cx) * Z / fx`). Doubling `depth_scale` halves all distances. This does **not** change the relative shape of the cloud — only the metric scale. + +**When to adjust:** Only change this if you are supplying a depth buffer in a different unit (e.g. centimetres instead of millimetres). In vizion3d the depth map is internally normalised before being encoded into uint16, so the default `1000.0` is correct for the standard workflow. + +--- + +### `depth_trunc` — maximum depth clip distance (metres) + +**Default:** `10.0` + +Points with a depth value greater than `depth_trunc` metres are discarded by Open3D before building the point cloud. This controls the far clipping plane. + +**Effect on the point cloud:** Lowering `depth_trunc` removes distant background points and produces a denser, cleaner cloud for near objects. Setting it to a very small value will discard almost all points. Setting it too large can include noisy, low-confidence depth estimates at the scene boundary. + +**Practical guidance:** +- Indoor close-up scenes: `2.0–5.0` m +- Room-scale scenes: `5.0–10.0` m (default) +- Outdoor or large-scale: `10.0–30.0` m + +--- + +## Default values and PrimeSense + +The built-in defaults match the **PrimeSense / Microsoft Kinect v1** sensor at 640×480 VGA resolution: + +| Parameter | Default | PrimeSense VGA | +|---|---|---| +| `fx` | `525.0` | 525.0 px | +| `fy` | `525.0` | 525.0 px | +| `cx` | `319.5` | 319.5 px | +| `cy` | `239.5` | 239.5 px | +| `depth_scale` | `1000.0` | — | +| `depth_trunc` | `10.0` | — | + +These are reasonable placeholders for any RGB camera with a ~60° horizontal FoV. For accurate metric reconstruction, always supply intrinsics from your actual camera calibration. + +--- + +## Usage: Direct Python + +```python +from vizion3d.lifting import ( + DepthEstimation, + DepthEstimationAdvanceConfig, + DepthEstimationCommand, +) + +# Full custom intrinsics (e.g. Intel RealSense D435 at 1280×720) +config = DepthEstimationAdvanceConfig( + fx=909.15, + fy=908.48, + cx=640.0, + cy=360.0, + depth_scale=1000.0, + depth_trunc=6.0, +) + +with open("scene.png", "rb") as f: + img_bytes = f.read() + +result = DepthEstimation().run( + DepthEstimationCommand( + image_input=img_bytes, + return_point_cloud=True, + advanced_config=config, + ) +) + +import numpy as np +points = np.asarray(result.point_cloud.points) +print(f"Points: {len(points)}") +``` + +Partial overrides work too — unspecified fields keep their defaults: + +```python +# Only change depth_trunc; everything else stays at PrimeSense defaults +result = DepthEstimation().run( + DepthEstimationCommand( + image_input=img_bytes, + return_point_cloud=True, + advanced_config=DepthEstimationAdvanceConfig(depth_trunc=3.0), + ) +) +``` + +--- + +## Usage: REST API + +All six config parameters are optional form fields on the `POST /lifting/depth-estimation` endpoint. + +```bash +# Full custom intrinsics +curl -X POST "http://localhost:8000/lifting/depth-estimation" \ + -F "image=@scene.png" \ + -F "return_point_cloud=true" \ + -F "fx=909.15" \ + -F "fy=908.48" \ + -F "cx=640.0" \ + -F "cy=360.0" \ + -F "depth_scale=1000.0" \ + -F "depth_trunc=6.0" +``` + +Partial overrides — omit any field to keep its default: + +```bash +# Only override depth_trunc +curl -X POST "http://localhost:8000/lifting/depth-estimation" \ + -F "image=@scene.png" \ + -F "return_point_cloud=true" \ + -F "depth_trunc=3.0" +``` + +Python `requests` equivalent: + +```python +import requests + +with open("scene.png", "rb") as f: + img_bytes = f.read() + +response = requests.post( + "http://localhost:8000/lifting/depth-estimation", + files={"image": ("scene.png", img_bytes, "image/png")}, + data={ + "return_point_cloud": "true", + "fx": "909.15", + "fy": "908.48", + "cx": "640.0", + "cy": "360.0", + "depth_trunc": "6.0", + }, +) +data = response.json() +print(f"Depth range: {data['min_depth']:.4f} → {data['max_depth']:.4f}") +``` + +--- + +## Usage: gRPC API + +The `DepthEstimationAdvanceConfig` proto message mirrors the Python model. All fields are `optional`, so any omitted field falls back to the server-side default. + +```python +import grpc +from vizion3d.proto import lifting_pb2, lifting_pb2_grpc + +channel = grpc.insecure_channel("localhost:50051") +stub = lifting_pb2_grpc.LiftingServiceStub(channel) + +with open("scene.png", "rb") as f: + img_bytes = f.read() + +# Full custom intrinsics +request = lifting_pb2.DepthEstimationRequest( + image_bytes=img_bytes, + return_point_cloud=True, + advanced_config=lifting_pb2.DepthEstimationAdvanceConfig( + fx=909.15, + fy=908.48, + cx=640.0, + cy=360.0, + depth_scale=1000.0, + depth_trunc=6.0, + ), +) +response = stub.RunDepthEstimation(request) +print(f"Depth range: {response.min_depth:.4f} → {response.max_depth:.4f}") +``` + +Partial override — only `depth_trunc`: + +```python +request = lifting_pb2.DepthEstimationRequest( + image_bytes=img_bytes, + return_point_cloud=True, + advanced_config=lifting_pb2.DepthEstimationAdvanceConfig(depth_trunc=3.0), +) +response = stub.RunDepthEstimation(request) +``` + +--- + +## How to get your camera intrinsics + +### Option 1: camera datasheet or SDK + +Most camera SDKs expose the intrinsic matrix directly: + +```python +# Intel RealSense +import pyrealsense2 as rs +pipeline = rs.pipeline() +profile = pipeline.start() +intr = profile.get_stream(rs.stream.color).as_video_stream_profile().intrinsics +config = DepthEstimationAdvanceConfig( + fx=intr.fx, fy=intr.fy, cx=intr.ppx, cy=intr.ppy +) +``` + +### Option 2: OpenCV calibration + +Run a standard checkerboard calibration with `cv2.calibrateCamera`. The returned `camera_matrix` is: + +``` +[[fx, 0, cx], + [ 0, fy, cy], + [ 0, 0, 1]] +``` + +```python +import cv2 +import numpy as np + +# After calibrating… +_, camera_matrix, _, _, _ = cv2.calibrateCamera(obj_points, img_points, image_size, None, None) + +config = DepthEstimationAdvanceConfig( + fx=float(camera_matrix[0, 0]), + fy=float(camera_matrix[1, 1]), + cx=float(camera_matrix[0, 2]), + cy=float(camera_matrix[1, 2]), +) +``` + +### Option 3: approximate from field of view + +If you know the camera's horizontal field of view `FoV_h` (in degrees) and image dimensions: + +```python +import math + +image_width = 1920 +image_height = 1080 +fov_h_deg = 69.0 # horizontal FoV in degrees + +fx = (image_width / 2) / math.tan(math.radians(fov_h_deg / 2)) +fy = fx # assumes square pixels +cx = image_width / 2 - 0.5 +cy = image_height / 2 - 0.5 + +config = DepthEstimationAdvanceConfig(fx=fx, fy=fy, cx=cx, cy=cy) +``` + +--- + +## Common camera presets + +These are approximate values for common cameras. Always prefer calibrated values over these presets. + +| Camera | Resolution | fx | fy | cx | cy | +|---|---|---|---|---|---| +| PrimeSense / Kinect v1 | 640×480 | 525.0 | 525.0 | 319.5 | 239.5 | +| Intel RealSense D415 | 1920×1080 | 1382.0 | 1382.0 | 960.5 | 540.5 | +| Intel RealSense D435 | 1280×720 | 909.0 | 908.0 | 640.0 | 360.0 | +| iPhone 14 wide (approx.) | 4032×3024 | 5500.0 | 5500.0 | 2016.0 | 1512.0 | +| Webcam 1080p (typical) | 1920×1080 | 1400.0 | 1400.0 | 960.0 | 540.0 | diff --git a/docs/features/depth_estimation_vs_stereo_depth.md b/docs/features/depth_estimation_vs_stereo_depth.md new file mode 100644 index 0000000..99b9e85 --- /dev/null +++ b/docs/features/depth_estimation_vs_stereo_depth.md @@ -0,0 +1,165 @@ +# Depth Estimation vs Stereo Depth + +vizion3d offers two different approaches to recovering depth from images. This page explains how they differ, when to choose each one, and what the practical trade-offs are. + +--- + +## What is Depth Estimation? + +[Depth Estimation](depth_estimation.md) uses a single RGB image and a monocular neural network (Depth Anything V2) to predict a per-pixel depth map. The network infers depth from visual cues — perspective, texture gradients, occlusion — without any geometric information about the camera. + +The output is **relative** (inverse) depth: closer objects have higher values, but the actual distances in metres are unknown. Point clouds built from monocular depth are geometrically consistent in shape but not anchored to a real-world scale. + +## What is Stereo Depth? + +[Stereo Depth](stereo_depth.md) takes a **rectified pair** of images (left and right, from two cameras with a known physical separation) and finds corresponding pixels across both views. The horizontal shift between matched pixels — called **disparity** — directly encodes depth through the stereo geometry formula: + +``` +depth_m = baseline_mm × focal_length_px / disparity_px / 1000 +``` + +Provided the camera calibration is accurate, the output is **real metric depth in metres** — every point in the point cloud has a physically meaningful distance from the camera. + +--- + +## Side-by-side comparison + +| | Stereo Depth (S2M2) | Depth Estimation (Depth Anything V2) | +|---|---|---| +| **Input** | Rectified left + right image pair | Single RGB image | +| **Depth type** | Metric (real metres) | Relative (inverse depth, arbitrary scale) | +| **Coordinate system** | Camera space | Camera space | +| **Units** | Metres (real) | Metres (fictitious — mapped to `[0, depth_trunc]`) | +| **Object at 2.4 m reads as 2.4 m** | Yes — if calibration is correct | No — depends on scene content | +| **Scale factor to world** | 1.0 (accurate) | Unknown, scene-dependent | +| **`point_cloud_scale` field** | 1.0 (accurate) | 1.0 (misleading — not real metres) | +| **Shape / topology correct** | Yes | Yes, if correct intrinsics supplied via `DepthEstimationAdvanceConfig` | +| **Camera calibration needed** | Yes — `focal_length`, `baseline`, `cx`, `cy` | Optional — only affects point cloud geometry | +| **Input requirements** | Stereo rig, rectified images | Any single photo | +| **Depth completeness** | Gaps in occluded / textureless regions | Dense — every pixel has a prediction | +| **Runtime** | Moderate (transformer-based matching) | Moderate (ViT-based encoder-decoder) | + +--- + +## When to use each + +### Use Depth Estimation when: +- You only have a single camera or single image. +- You need **dense** depth predictions (no holes from occlusion). +- You want to visualise relative 3D structure without exact scale. +- You are doing scene understanding, novel view synthesis, or artistic depth-of-field effects. + +### Use Stereo Depth when: +- You have a stereo camera rig with known calibration. +- You need **real metric distances** — for robotics, measurement, AR anchoring. +- Textureless or low-contrast regions can be handled by the rig geometry. +- You need consistent scale across different scenes and camera positions. + +--- + +## Working with point clouds from each + +### Monocular point cloud (Depth Estimation) + +```python +from vizion3d.lifting import DepthEstimation, DepthEstimationAdvanceConfig, DepthEstimationCommand +import numpy as np + +result = DepthEstimation().run( + DepthEstimationCommand( + image_input="scene.png", + return_point_cloud=True, + advanced_config=DepthEstimationAdvanceConfig( + fx=909.15, fy=908.48, cx=640.0, cy=360.0, + ), + ) +) + +points = np.asarray(result.point_cloud.points) # shape (N, 3) +# point_cloud_scale == 1.0, but distances are NOT real metres — +# the depth model output is relative and mapped to depth_trunc. +print(f"point_cloud_scale: {result.point_cloud_scale}") # 1.0 (misleading) +``` + +### Stereo point cloud (Stereo Depth) + +```python +from vizion3d.stereo import StereoDepth, StereoDepthAdvancedConfig, StereoDepthCommand +import numpy as np + +result = StereoDepth().run( + StereoDepthCommand( + left_image="left.png", + right_image="right.png", + return_point_cloud=True, + advanced_config=StereoDepthAdvancedConfig( + focal_length=1733.74, + cx=792.27, + cy=541.89, + baseline=536.62, + ), + ) +) + +points = np.asarray(result.point_cloud.points) # shape (N, 3), real metres +dist = np.linalg.norm(points[0] - points[1]) * result.point_cloud_scale +print(f"Real distance between p0 and p1: {dist:.4f} m") # actual metres +print(f"point_cloud_scale: {result.point_cloud_scale}") # 1.0 (accurate) +``` + +--- + +## Output differences at a glance + +| Output field | Depth Estimation | Stereo Depth | +|---|---|---| +| `depth_map` | Relative depth (fictitious metres) | Metric depth (real metres) | +| `disparity_map` | Not present | Pixel disparity (always returned) | +| `min_depth` / `max_depth` | Relative range | Real range in metres | +| `point_cloud_scale` | 1.0 (misleading) | 1.0 (accurate) | +| `backend_used` | Local path to Depth Anything V2 `.pth` | Local path to S2M2 `.pth` | + +--- + +## Advanced config comparison + +Both tasks expose a camera configuration object, but the fields are different because the underlying geometry differs: + +### DepthEstimationAdvanceConfig (monocular) + +Controls how the relative depth map is converted into a point cloud: + +```python +from vizion3d.lifting import DepthEstimationAdvanceConfig + +cfg = DepthEstimationAdvanceConfig( + fx=525.0, # horizontal focal length (pixels) + fy=525.0, # vertical focal length (pixels) + cx=319.5, # principal point x + cy=239.5, # principal point y + depth_scale=1000.0, # uint16 → metres divisor + depth_trunc=10.0, # max depth in metres +) +``` + +### StereoDepthAdvancedConfig (stereo) + +Controls the stereo geometry and point cloud quality filters: + +```python +from vizion3d.stereo import StereoDepthAdvancedConfig + +cfg = StereoDepthAdvancedConfig( + focal_length=1000.0, # focal length in pixels (fx = fy assumed) + cx=640.0, # principal point x + cy=360.0, # principal point y + baseline=100.0, # stereo baseline in millimetres + doffs=0.0, # disparity offset (Middlebury-style calibration) + z_far=10.0, # max depth in metres + conf_threshold=0.1, # min confidence score for point inclusion + occ_threshold=0.5, # min occlusion score for point inclusion + scale_factor=1.0, # input downscale for speed/quality tradeoff +) +``` + +The key stereo-only parameters are `baseline` (physical rig geometry) and `doffs` (calibration offset), which have no equivalent in monocular depth — they are meaningless without a second camera. diff --git a/docs/features/stereo_depth.md b/docs/features/stereo_depth.md new file mode 100644 index 0000000..488f9f1 --- /dev/null +++ b/docs/features/stereo_depth.md @@ -0,0 +1,426 @@ +# Stereo Depth + +**Category:** Lifting (2D → 3D) +**Experimental:** No + +Stereo depth estimation recovers per-pixel **metric depth** (in metres) from a pair of rectified left/right RGB images by matching corresponding pixels across the two views and applying the stereo geometry formula: + +``` +depth_m = baseline_mm × focal_length_px / disparity_px / 1000 +``` + +vizion3d uses [S2M2](https://github.com/Dongyeop-Yoo/S2M2) (Stereo Matching Model with Multi-scale transformer) as its stereo backend. Unlike [Depth Estimation](depth_estimation.md), stereo depth produces **real-world metric distances** — provided the camera calibration parameters are correct. + +--- + +## Model backends + +Default checkpoint download: +[stereo-depth-s2m2-L.pth](https://github.com/OlafenwaMoses/vizion3D/releases/download/essentials-v1/stereo-depth-s2m2-L.pth) + +```bash +curl -L \ + https://github.com/OlafenwaMoses/vizion3D/releases/download/essentials-v1/stereo-depth-s2m2-L.pth \ + -o stereo-depth-s2m2-L.pth +``` + +| Value | What happens | +|---|---| +| *(default)* | Downloads the vizion3D release checkpoint (`stereo-depth-s2m2-L.pth`, the L variant) to `~/.cache/vizion3d/models/` on first use, then loads it | +| An HTTPS URL ending in `.pth` or `.pt` | Downloaded to the cache directory on first use, then loaded as an S2M2 checkpoint | +| A local `.pth` or `.pt` file path | Loaded directly — no download | + +Models are kept in memory after the first inference. Set `VIZION3D_MODEL_CACHE` to override the cache directory. + +### S2M2 variants + +The S2M2 architecture comes in four size variants. The correct one is detected automatically from the checkpoint filename: + +| Variant | Channels | Transformers | Speed | Quality | +|---|---|---|---|---| +| S (`-S.pth`) | 128 | 1 | Fastest | Good | +| M (`-M.pth`) | 192 | 2 | Fast | Better | +| L (`-L.pth`) | 256 | 3 | Balanced | Best (default) | +| XL (`-XL.pth`) | 384 | 3 | Slowest | Best | + +--- + +## Command parameters + +`StereoDepthCommand` is the input contract for this task. + +| Parameter | Type | Required | Default | Description | +|---|---|---|---|---| +| `left_image` | `str \| bytes` | **Yes** | — | Left-camera image. Pass a file path string or raw image bytes. | +| `right_image` | `str \| bytes` | **Yes** | — | Right-camera image (same resolution, horizontally offset from `left_image`). | +| `model_backend` | `str` | No | vizion3D release checkpoint URL | S2M2 checkpoint. See [Model backends](#model-backends) above. | +| `return_depth_image` | `bool` | No | `False` | If `True`, the result includes a 16-bit grayscale Open3D Image of the depth map. | +| `return_point_cloud` | `bool` | No | `False` | If `True`, the result includes an Open3D PointCloud in metres. | +| `return_mesh` | `bool` | No | `False` | If `True`, the result includes an Open3D TriangleMesh reconstructed via ball-pivoting. | +| `advanced_config` | `StereoDepthAdvancedConfig` | No | 1280×720 @ 100 mm baseline defaults | Camera intrinsics and inference settings. See [Advanced config](#advanced-config) below. | + +--- + +## Result fields + +`StereoDepthResult` is the output contract. + +| Field | Type | Always present | Description | +|---|---|---|---| +| `depth_map` | `list[list[float]]` | Yes | Metric depth in **metres**, shape `[H][W]`. Real-world distances (assuming correct calibration). | +| `disparity_map` | `list[list[float]]` | Yes | Raw disparity in **pixels**, shape `[H][W]`. Horizontal pixel offset between matched features. | +| `min_depth` | `float` | Yes | Minimum value in `depth_map` (metres). | +| `max_depth` | `float` | Yes | Maximum value in `depth_map`. Guaranteed `max_depth >= min_depth`. | +| `backend_used` | `str` | Yes | Resolved local file path of the checkpoint used. | +| `depth_image` | `open3d.geometry.Image \| None` | When `return_depth_image=True` | 16-bit grayscale image, dtype `uint16`. The full 0–65535 range maps to `[min_depth, max_depth]` in metres. | +| `point_cloud` | `open3d.geometry.PointCloud \| None` | When `return_point_cloud=True` | Coloured 3D point cloud, coordinates in **metres**. | +| `mesh` | `open3d.geometry.TriangleMesh \| None` | When `return_mesh=True` | Surface mesh from ball-pivoting. Includes vertex colours. | +| `point_cloud_scale` | `float` | Yes | Always `1.0` — stereo depth produces real metric coordinates. | + +--- + +## 1. Direct Python import — image bytes + +```python +from vizion3d.stereo import StereoDepth, StereoDepthCommand + +with open("left.png", "rb") as f: + left_bytes = f.read() +with open("right.png", "rb") as f: + right_bytes = f.read() + +cmd = StereoDepthCommand(left_image=left_bytes, right_image=right_bytes) +result = StereoDepth().run(cmd) + +print(f"Depth range : {result.min_depth:.2f} → {result.max_depth:.2f} m") +print(f"Backend : {result.backend_used}") +``` + +--- + +## 2. Direct Python import — file paths + +```python +from vizion3d.stereo import StereoDepth, StereoDepthCommand + +cmd = StereoDepthCommand( + left_image="left.png", + right_image="right.png", +) +result = StereoDepth().run(cmd) + +print(f"Depth range: {result.min_depth:.2f} → {result.max_depth:.2f} m") +``` + +--- + +## 3. Disparity map + +The raw disparity map (in pixels) is always returned alongside the depth map. + +```python +import numpy as np +from vizion3d.stereo import StereoDepth, StereoDepthCommand + +cmd = StereoDepthCommand(left_image="left.png", right_image="right.png") +result = StereoDepth().run(cmd) + +disp = np.array(result.disparity_map) +print(f"Disparity range: {disp.min():.1f} → {disp.max():.1f} px") +``` + +--- + +## 4. Depth image (16-bit PNG) + +```python +import numpy as np +from PIL import Image as PILImage +from vizion3d.stereo import StereoDepth, StereoDepthCommand + +cmd = StereoDepthCommand( + left_image="left.png", + right_image="right.png", + return_depth_image=True, +) +result = StereoDepth().run(cmd) + +depth_array = np.asarray(result.depth_image) # shape (H, W), dtype uint16 +PILImage.fromarray(depth_array).save("depth.png") +``` + +--- + +## 5. Point cloud + +Point coordinates are in **real metres** — `point_cloud_scale` is always `1.0`. + +```python +import numpy as np +import open3d as o3d +from vizion3d.stereo import StereoDepth, StereoDepthAdvancedConfig, StereoDepthCommand + +cmd = StereoDepthCommand( + left_image="left.png", + right_image="right.png", + return_point_cloud=True, + advanced_config=StereoDepthAdvancedConfig( + focal_length=1733.74, + cx=792.27, + cy=541.89, + baseline=536.62, # mm + ), +) +result = StereoDepth().run(cmd) + +pcd = result.point_cloud +points = np.asarray(pcd.points) # shape (N, 3), metres +print(f"Points: {len(points):,}") +print(f"Scale : {result.point_cloud_scale} m/unit") # always 1.0 + +# Real-world distance between two points +dist = np.linalg.norm(points[0] - points[1]) * result.point_cloud_scale +print(f"p0→p1: {dist:.4f} m") + +o3d.io.write_point_cloud("scene.ply", pcd) +``` + +--- + +## 6. Surface mesh + +```python +import open3d as o3d +from vizion3d.stereo import StereoDepth, StereoDepthCommand + +cmd = StereoDepthCommand( + left_image="left.png", + right_image="right.png", + return_mesh=True, +) +result = StereoDepth().run(cmd) + +mesh = result.mesh +print(f"Vertices : {len(mesh.vertices)}") +print(f"Triangles : {len(mesh.triangles)}") +o3d.io.write_triangle_mesh("scene_mesh.ply", mesh) +``` + +--- + +## 7. All outputs at once + +```python +import numpy as np +import open3d as o3d +from vizion3d.stereo import StereoDepth, StereoDepthCommand + +cmd = StereoDepthCommand( + left_image="left.png", + right_image="right.png", + return_depth_image=True, + return_point_cloud=True, + return_mesh=True, +) +result = StereoDepth().run(cmd) + +print(f"Depth range : {result.min_depth:.2f} → {result.max_depth:.2f} m") +depth_arr = np.asarray(result.depth_image) # uint16 (H, W) +o3d.io.write_point_cloud("scene.ply", result.point_cloud) +o3d.io.write_triangle_mesh("scene_mesh.ply", result.mesh) +``` + +--- + +## 8. Speed vs quality: scale factor + +Use `scale_factor < 1.0` to downsample input before inference for faster results: + +```python +from vizion3d.stereo import StereoDepth, StereoDepthAdvancedConfig, StereoDepthCommand + +cmd = StereoDepthCommand( + left_image="left.png", + right_image="right.png", + advanced_config=StereoDepthAdvancedConfig( + scale_factor=0.5, # half-resolution → ~3–4× faster + ), +) +result = StereoDepth().run(cmd) +``` + +--- + +## 9. REST API + +Start the server with all REST features enabled: + +```bash +uv run vizion3d-serve-rest +``` + +To preload the stereo checkpoint into memory at startup, pass `--stereo_model`. +This also enables the stereo-depth endpoint. If this flag is omitted, the +default vizion3D release model is downloaded on first inference and cached under +`~/.cache/vizion3d/models/`. + +```bash +uv run vizion3d-serve-rest \ + --stereo_model /models/stereo-depth-s2m2-L.pth +``` + +The REST server can expose only selected features. If none of +`--depth_estimation`, `--stereo_depth`, `--depth_model`, or `--stereo_model` is +provided, all features are enabled. If any of those flags is provided, only the +selected features are enabled. A model path flag selects and preloads its +feature: + +```bash +# Only POST /lifting/stereo-depth +uv run vizion3d-serve-rest --stereo_depth + +# Only stereo depth, with the model loaded before the first request +uv run vizion3d-serve-rest \ + --stereo_depth \ + --stereo_model /models/stereo-depth-s2m2-L.pth + +# Enable both depth estimation and stereo depth explicitly +uv run vizion3d-serve-rest \ + --depth_estimation \ + --stereo_depth \ + --depth_model /models/depth_anything_v2_vitb.pth \ + --stereo_model /models/stereo-depth-s2m2-L.pth +``` + +Send a request with two image files: + +```bash +curl -X POST "http://localhost:8000/lifting/stereo-depth" \ + -F "left_image=@left.png" \ + -F "right_image=@right.png" \ + -F "focal_length=1733.74" \ + -F "baseline=536.62" \ + -F "cx=792.27" \ + -F "cy=541.89" \ + -F "return_point_cloud=true" +``` + +The response is a JSON-serialised `StereoDepthResult`. Binary fields (`depth_image`, `point_cloud_ply`, `mesh_ply`) are base64-encoded. + +--- + +## 10. gRPC API + +Start the server: + +```bash +uv run vizion3d-serve-grpc +``` + +Call from a gRPC client: + +```python +import grpc +from vizion3d.proto import lifting_pb2, lifting_pb2_grpc + +channel = grpc.insecure_channel("localhost:50051") +stub = lifting_pb2_grpc.LiftingServiceStub(channel) + +with open("left.png", "rb") as f: + left_bytes = f.read() +with open("right.png", "rb") as f: + right_bytes = f.read() + +request = lifting_pb2.StereoDepthRequest( + left_image_bytes=left_bytes, + right_image_bytes=right_bytes, + return_point_cloud=True, + advanced_config=lifting_pb2.StereoDepthAdvancedConfig( + focal_length=1733.74, + baseline=536.62, + cx=792.27, + cy=541.89, + ), +) +response = stub.RunStereoDepth(request) +print(f"Min depth : {response.min_depth:.2f} m") +print(f"Max depth : {response.max_depth:.2f} m") +print(f"Backend : {response.backend_used}") +``` + +--- + +## Advanced config + +`StereoDepthAdvancedConfig` supplies the camera calibration needed for accurate metric depth. + +| Field | Type | Default | Description | +|---|---|---|---| +| `focal_length` | `float` | `1000.0` | Focal length in pixels (assumes fx = fy). Override with your calibration. | +| `cx` | `float` | `640.0` | Principal point x (pixel column of optical axis). | +| `cy` | `float` | `360.0` | Principal point y (pixel row of optical axis). | +| `baseline` | `float` | `100.0` | Stereo baseline in **millimetres**. | +| `doffs` | `float` | `0.0` | Disparity offset (non-zero for Middlebury-style calibration). | +| `z_far` | `float` | `10.0` | Max depth in metres for point cloud. | +| `conf_threshold` | `float` | `0.1` | Min per-pixel confidence score for point cloud inclusion. | +| `occ_threshold` | `float` | `0.5` | Min occlusion score for point cloud inclusion. | +| `scale_factor` | `float` | `1.0` | Input downscale factor (`0.5` = half-res, ~3–4× faster). | + +### How to obtain camera intrinsics + +**From a calibration file (e.g. Middlebury):** +```python +# calib.txt format: cam0=[fx 0 cx; 0 fy cy; 0 0 1] +# baseline=B (mm), doffs=d +from vizion3d.stereo import StereoDepthAdvancedConfig + +cfg = StereoDepthAdvancedConfig( + focal_length=1733.74, # from calib.txt + cx=792.27, + cy=541.89, + baseline=536.62, # B in mm + doffs=0.0, # d from calib.txt +) +``` + +**From Intel RealSense SDK:** +```python +import pyrealsense2 as rs + +pipeline = rs.pipeline() +profile = pipeline.start() +left_stream = profile.get_stream(rs.stream.infrared, 1) +intrinsics = left_stream.as_video_stream_profile().get_intrinsics() + +cfg = StereoDepthAdvancedConfig( + focal_length=intrinsics.fx, + cx=intrinsics.ppx, + cy=intrinsics.ppy, + baseline=50.0, # RealSense D435 baseline ≈ 50 mm +) +``` + +**Approximation from field of view:** +```python +import math + +hfov_deg = 90.0 # horizontal FOV from camera spec +image_width = 1280 +focal_length = image_width / (2 * math.tan(math.radians(hfov_deg / 2))) + +cfg = StereoDepthAdvancedConfig( + focal_length=focal_length, + cx=image_width / 2 - 0.5, + cy=720 / 2 - 0.5, + baseline=100.0, +) +``` + +--- + +## Known limitations + +- **Rectified pairs required** — images must be stereo-rectified so corresponding points lie on the same horizontal scanline. Un-rectified pairs will produce incorrect results. +- **Metric scale depends on calibration** — an incorrect `baseline` or `focal_length` scales all depth values uniformly. Always use calibrated values for real applications. +- **Ball-pivoting mesh quality** — works best on dense, evenly sampled point clouds. Sparse or noisy clouds from occluded regions may produce gaps or missing faces. +- **Python 3.12 required for Open3D** — `return_depth_image`, `return_point_cloud`, and `return_mesh` require Open3D, which currently only supports Python 3.12 in this project. diff --git a/mkdocs.yml b/mkdocs.yml index 57651a7..fdaabb0 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -26,5 +26,8 @@ nav: - Home: index.md - Features: - Depth Estimation: features/depth_estimation.md + - Depth Estimation Advanced Config: features/depth_estimation_advanced_config.md + - Stereo Depth: features/stereo_depth.md + - Depth Estimation vs Stereo Depth: features/depth_estimation_vs_stereo_depth.md - API Reference: - Lifting (2D → 3D): api/lifting.md diff --git a/pyproject.toml b/pyproject.toml index 4377016..eb8ee69 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -96,9 +96,8 @@ addopts = "-v" [tool.ruff] line-length = 100 target-version = "py312" -extend-exclude = ["*_pb2*.py"] +extend-exclude = ["*_pb2*.py", "stereo_sample.py"] [tool.ruff.lint] select = ["E", "F", "I"] ignore = [] - diff --git a/tests/assets/stereo/teddy/SOURCE.md b/tests/assets/stereo/teddy/SOURCE.md new file mode 100644 index 0000000..5c09636 --- /dev/null +++ b/tests/assets/stereo/teddy/SOURCE.md @@ -0,0 +1,13 @@ +Middlebury Stereo Evaluation v3 quarter-resolution Teddy sample. + +Source: https://vision.middlebury.edu/stereo/submit3/ +Archive: https://vision.middlebury.edu/stereo/submit3/zip/MiddEval3-data-Q.zip +Scene path in archive: MiddEval3/trainingQ/Teddy/ + +Files: +- `left.png` is `im0.png` +- `right.png` is `im1.png` +- `calib.txt` is the original Middlebury calibration file + +The calibration file uses the Middlebury stereo format: +`Z = baseline * f / (disparity + doffs)`. diff --git a/tests/assets/stereo/teddy/calib.txt b/tests/assets/stereo/teddy/calib.txt new file mode 100644 index 0000000..afeb0ad --- /dev/null +++ b/tests/assets/stereo/teddy/calib.txt @@ -0,0 +1,12 @@ +cam0=[1500 0 199; 0 1500 187; 0 0 1] +cam1=[1500 0 251; 0 1500 187; 0 0 1] +doffs=52 +baseline=80 +width=450 +height=375 +ndisp=64 +isint=0 +vmin=12 +vmax=55 +dyavg=0 +dymax=0 diff --git a/tests/assets/stereo/teddy/left.png b/tests/assets/stereo/teddy/left.png new file mode 100644 index 0000000..c8e7d27 Binary files /dev/null and b/tests/assets/stereo/teddy/left.png differ diff --git a/tests/assets/stereo/teddy/right.png b/tests/assets/stereo/teddy/right.png new file mode 100644 index 0000000..dde5d9a Binary files /dev/null and b/tests/assets/stereo/teddy/right.png differ diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 7bbe771..8399fd1 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -3,14 +3,17 @@ Provides -------- -indoor_image_bytes — real 640×480 indoor-scene JPEG for inference -local_model_path — explicit .pth path in a session-scoped tmp dir - (symlinked from the default cache if available, - otherwise freshly downloaded; cleaned up by pytest) -grpc_client_stub — live LiftingService stub backed by an in-process - gRPC server running in a background thread-pool -timing_collector — session-wide store that every test appends to -pytest_terminal_summary — pretty inference-timing report printed at the end +indoor_image_bytes — real 640×480 indoor-scene JPEG for inference +stereo_image_pair — real calibrated Middlebury Teddy stereo pair +stereo_advanced_config — Teddy calibration mapped to StereoDepthAdvancedConfig +local_model_path — explicit .pth path for the depth-estimation model + (symlinked from cache if available, else downloaded) +local_stereo_model_path — explicit .pth path for the stereo-depth model + (symlinked from cache if available, else downloaded) +grpc_client_stub — live LiftingService stub backed by an in-process + gRPC server running in a background thread-pool +timing_collector — session-wide store that every test appends to +pytest_terminal_summary — pretty inference-timing report printed at the end """ from __future__ import annotations @@ -36,6 +39,7 @@ # Timing collector # ────────────────────────────────────────────────────────────────────────────── + @dataclass class TimingRecord: entry_point: str @@ -57,9 +61,7 @@ def add( duration: float, output_dir: str = "", ) -> None: - self.records.append( - TimingRecord(entry_point, scenario, run, duration, output_dir) - ) + self.records.append(TimingRecord(entry_point, scenario, run, duration, output_dir)) # Module-level singleton — pytest_terminal_summary reads from it after the session @@ -72,9 +74,10 @@ def timing_collector() -> InferenceTimingCollector: # ────────────────────────────────────────────────────────────────────────────── -# Image fixture +# Image fixtures # ────────────────────────────────────────────────────────────────────────────── + @pytest.fixture(scope="session") def indoor_image_bytes() -> bytes: path = ASSETS_DIR / "indoor_scene.jpg" @@ -86,10 +89,66 @@ def indoor_image_bytes() -> bytes: return path.read_bytes() +def _parse_middlebury_calibration(path: Path) -> dict[str, float]: + values: dict[str, float] = {} + for line in path.read_text().splitlines(): + if not line or "=" not in line: + continue + key, raw_value = line.split("=", 1) + if key == "cam0": + rows = raw_value.strip()[1:-1].split(";") + matrix = [[float(v) for v in row.split()] for row in rows] + values["focal_length"] = matrix[0][0] + values["cx"] = matrix[0][2] + values["cy"] = matrix[1][2] + elif key in {"doffs", "baseline", "width", "height", "ndisp"}: + values[key] = float(raw_value) + return values + + +@pytest.fixture(scope="session") +def stereo_image_pair() -> tuple[bytes, bytes]: + """Return a real calibrated ``(left_bytes, right_bytes)`` stereo pair. + + The fixture uses Middlebury's quarter-resolution Teddy sample. Each image is + roughly 300 KB, and the paired ``calib.txt`` provides the camera intrinsics, + disparity offset, and baseline used by ``stereo_advanced_config``. + """ + left = ASSETS_DIR / "stereo" / "teddy" / "left.png" + right = ASSETS_DIR / "stereo" / "teddy" / "right.png" + assert left.exists(), f"Stereo left image not found: {left}" + assert right.exists(), f"Stereo right image not found: {right}" + return left.read_bytes(), right.read_bytes() + + +@pytest.fixture(scope="session") +def stereo_calibration_values() -> dict[str, float]: + calib = ASSETS_DIR / "stereo" / "teddy" / "calib.txt" + assert calib.exists(), f"Stereo calibration not found: {calib}" + return _parse_middlebury_calibration(calib) + + +@pytest.fixture(scope="session") +def stereo_advanced_config(stereo_calibration_values): + from vizion3d.stereo import StereoDepthAdvancedConfig + + return StereoDepthAdvancedConfig( + focal_length=stereo_calibration_values["focal_length"], + cx=stereo_calibration_values["cx"], + cy=stereo_calibration_values["cy"], + baseline=stereo_calibration_values["baseline"], + doffs=stereo_calibration_values["doffs"], + z_far=10.0, + conf_threshold=0.0, + occ_threshold=0.0, + ) + + # ────────────────────────────────────────────────────────────────────────────── -# Local-model-path fixture +# Local-model-path fixtures # ────────────────────────────────────────────────────────────────────────────── + @pytest.fixture(scope="session") def local_model_path(tmp_path_factory) -> str: """ @@ -117,14 +176,40 @@ def local_model_path(tmp_path_factory) -> str: return str(dest) +@pytest.fixture(scope="session") +def local_stereo_model_path(tmp_path_factory) -> str: + """ + Provide a local .pth path for the stereo-depth model, cleaned up after + the session. If the model is already in the default vizion3d cache we + symlink it (free); otherwise we download it fresh. + """ + from vizion3d.lifting.defaults import default_model_cache_dir, download_model + from vizion3d.stereo.defaults import ( + DEFAULT_STEREO_MODEL_FILENAME, + DEFAULT_STEREO_MODEL_URL, + ) + + default_cache = default_model_cache_dir() / DEFAULT_STEREO_MODEL_FILENAME + tmp_dir = tmp_path_factory.mktemp("local_stereo_model") + dest = tmp_dir / DEFAULT_STEREO_MODEL_FILENAME + + if default_cache.exists(): + dest.symlink_to(default_cache.resolve()) + else: + download_model(DEFAULT_STEREO_MODEL_URL, cache_dir=tmp_dir) + + assert dest.exists() or dest.is_symlink(), f"Stereo model not found at {dest}" + return str(dest) + + # ────────────────────────────────────────────────────────────────────────────── # gRPC server + client stub fixture # ────────────────────────────────────────────────────────────────────────────── -_MAX_MSG = 500 * 1024 * 1024 # match server cap +_MAX_MSG = 500 * 1024 * 1024 # match server cap _GRPC_OPTIONS = [ - ("grpc.max_send_message_length", _MAX_MSG), + ("grpc.max_send_message_length", _MAX_MSG), ("grpc.max_receive_message_length", _MAX_MSG), ] @@ -142,10 +227,8 @@ def grpc_client_stub(): futures.ThreadPoolExecutor(max_workers=4), options=_GRPC_OPTIONS, ) - lifting_pb2_grpc.add_LiftingServiceServicer_to_server( - LiftingServiceServicer(), server - ) - port = server.add_insecure_port("[::]:0") # 0 → OS picks a free port + lifting_pb2_grpc.add_LiftingServiceServicer_to_server(LiftingServiceServicer(), server) + port = server.add_insecure_port("[::]:0") # 0 → OS picks a free port server.start() channel = grpc.insecure_channel(f"localhost:{port}", options=_GRPC_OPTIONS) @@ -161,17 +244,18 @@ def grpc_client_stub(): # Terminal report (hook) # ────────────────────────────────────────────────────────────────────────────── -def pytest_terminal_summary(terminalreporter, exitstatus, config): # noqa: ARG001 + +def pytest_terminal_summary(terminalreporter, exitstatus, config): # noqa: ARG001 records = _COLLECTOR.records if not records: return - W = 82 - EP = 10 # entry-point col width - SC = 16 # scenario col width - RN = 4 # run col width - DU = 10 # duration col width - ST = 20 # status col width + W = 82 + EP = 10 # entry-point col width + SC = 16 # scenario col width + RN = 4 # run col width + DU = 10 # duration col width + ST = 20 # status col width def _write(line: str = "") -> None: terminalreporter.write_line(line) @@ -180,14 +264,10 @@ def _thick() -> None: _write("━" * W) def _thin() -> None: - _write( - f" {'─'*EP}─┼─{'─'*SC}─┼─{'─'*RN}─┼─{'─'*(DU)}─┼─{'─'*ST}" - ) + _write(f" {'─' * EP}─┼─{'─' * SC}─┼─{'─' * RN}─┼─{'─' * (DU)}─┼─{'─' * ST}") def _row(ep="", sc="", run="", dur="", status="") -> None: - _write( - f" {ep:<{EP}} │ {sc:<{SC}} │ {run:^{RN}} │ {dur:>{DU}} │ {status}" - ) + _write(f" {ep:<{EP}} │ {sc:<{SC}} │ {run:^{RN}} │ {dur:>{DU}} │ {status}") _write() _thick() @@ -197,17 +277,17 @@ def _row(ep="", sc="", run="", dur="", status="") -> None: _row("Entry Point", "Scenario", "Run", "Duration", "Status") _thin() - def sort_key(r): return (r.entry_point, r.scenario, r.run) - def group_key(r): return (r.entry_point, r.scenario) + def sort_key(r): + return (r.entry_point, r.scenario, r.run) + + def group_key(r): + return (r.entry_point, r.scenario) first_loads: list[float] = [] - warm_times: list[float] = [] + warm_times: list[float] = [] sorted_records = sorted(records, key=sort_key) - groups = [ - (k, list(v)) - for k, v in groupby(sorted_records, key=group_key) - ] + groups = [(k, list(v)) for k, v in groupby(sorted_records, key=group_key)] for g_idx, ((ep, sc), recs) in enumerate(groups): if g_idx > 0: @@ -217,15 +297,15 @@ def group_key(r): return (r.entry_point, r.scenario) first_dur = recs[0].duration for i, rec in enumerate(recs): - ep_label = ep if i == 0 else "" - sc_label = sc if i == 0 else "" - dur_str = f"{rec.duration:7.3f}s" + ep_label = ep if i == 0 else "" + sc_label = sc if i == 0 else "" + dur_str = f"{rec.duration:7.3f}s" if rec.run == 1: status = "◉ COLD LOAD" first_loads.append(rec.duration) else: - pct = (1.0 - rec.duration / first_dur) * 100.0 + pct = (1.0 - rec.duration / first_dur) * 100.0 status = f"⚡ {pct:4.1f}% faster" warm_times.append(rec.duration) @@ -238,8 +318,8 @@ def group_key(r): return (r.entry_point, r.scenario) if first_loads and warm_times: avg_load = sum(first_loads) / len(first_loads) avg_warm = sum(warm_times) / len(warm_times) - speedup = avg_load / avg_warm if avg_warm > 0 else float("inf") - total = len(records) + speedup = avg_load / avg_warm if avg_warm > 0 else float("inf") + total = len(records) pad = 42 _write(f" {'SUMMARY'}") diff --git a/tests/integration/test_direct.py b/tests/integration/test_depth_estimation_direct.py similarity index 64% rename from tests/integration/test_direct.py rename to tests/integration/test_depth_estimation_direct.py index ad7ad78..e37bb4e 100644 --- a/tests/integration/test_direct.py +++ b/tests/integration/test_depth_estimation_direct.py @@ -22,36 +22,39 @@ pytest.importorskip("open3d", reason="open3d required — run: uv python pin 3.12 && uv sync") -from vizion3d.lifting import DepthEstimation, DepthEstimationCommand # noqa: E402 +from vizion3d.lifting import ( # noqa: E402 + DepthEstimation, + DepthEstimationAdvanceConfig, + DepthEstimationCommand, +) from vizion3d.lifting.defaults import DEFAULT_DEPTH_MODEL_URL # noqa: E402 from vizion3d.lifting.handlers import DepthEstimationHandler # noqa: E402 from vizion3d.lifting.utils import create_ply_binary # noqa: E402 N_RUNS = 5 -COLD_LIMIT = float(os.environ.get("VIZION3D_TEST_COLD_LIMIT", "10.0")) -WARM_LIMIT = float(os.environ.get("VIZION3D_TEST_WARM_LIMIT", "1.0")) +DEPTH_COLD_LIMIT = float(os.environ.get("VIZION3D_TEST_DEPTH_COLD_LIMIT", "10.0")) +DEPTH_WARM_LIMIT = float(os.environ.get("VIZION3D_TEST_DEPTH_WARM_LIMIT", "1.0")) # ────────────────────────────────────────────────────────────────────────────── # Helpers # ────────────────────────────────────────────────────────────────────────────── + def _save_outputs(result, run_dir: Path, run: int) -> None: run_dir.mkdir(parents=True, exist_ok=True) prefix = run_dir / f"run_{run:02d}" - (prefix.with_suffix(".depth_map.json")).write_text( - json.dumps(result.depth_map) - ) + (prefix.with_suffix(".depth_map.json")).write_text(json.dumps(result.depth_map)) if result.depth_image is not None: - arr = np.asarray(result.depth_image) # uint16 (H, W) + arr = np.asarray(result.depth_image) # uint16 (H, W) PILImage.fromarray(arr).save(str(prefix) + "_depth.png") if result.point_cloud is not None and result.point_cloud.has_points(): - pts = np.asarray(result.point_cloud.points).astype(np.float32) + pts = np.asarray(result.point_cloud.points).astype(np.float32) cols = (np.asarray(result.point_cloud.colors) * 255).astype(np.uint8) - ply = create_ply_binary(pts, cols) + ply = create_ply_binary(pts, cols) (str(prefix) + "_point_cloud.ply") Path(str(prefix) + "_point_cloud.ply").write_bytes(ply) @@ -70,7 +73,7 @@ def _run_group( timings: list[float] = [] for run in range(1, N_RUNS + 1): - t0 = time.perf_counter() + t0 = time.perf_counter() result = DepthEstimation().run( DepthEstimationCommand( image_input=indoor_image_bytes, @@ -88,25 +91,29 @@ def _run_group( # ── per-run assertions ──────────────────────────────────────────────── assert isinstance(result.depth_map, list) and len(result.depth_map) > 0 assert all(isinstance(row, list) for row in result.depth_map) - assert result.max_depth > result.min_depth, \ + assert result.max_depth > result.min_depth, ( "Depth variation expected for a real indoor scene" + ) assert result.depth_image is not None, "depth_image was requested but missing" - assert result.point_cloud is not None and result.point_cloud.has_points(), \ + assert result.point_cloud is not None and result.point_cloud.has_points(), ( "point_cloud was requested but missing or empty" + ) # ── caching assertion: model must be in memory after run 1 ─────────────── - assert len(DepthEstimationHandler._depth_anything_models) > 0, \ + assert len(DepthEstimationHandler._depth_anything_models) > 0, ( "Model should be in DepthEstimationHandler._depth_anything_models after first inference" + ) # ── timing assertions ───────────────────────────────────────────────────── - assert timings[0] < COLD_LIMIT, ( - f"[{entry_point} / {scenario}] " - f"Cold load took {timings[0]:.3f}s — expected < {COLD_LIMIT}s" + assert timings[0] < DEPTH_COLD_LIMIT, ( + f"[{entry_point} / {scenario}] Cold load took {timings[0]:.3f}s " + f"— expected < {DEPTH_COLD_LIMIT}s" ) for i, t in enumerate(timings[1:], start=2): - assert t < WARM_LIMIT, ( + assert t < DEPTH_WARM_LIMIT, ( f"[{entry_point} / {scenario}] " - f"Run {i} took {t:.3f}s — expected < {WARM_LIMIT}s (model should be cached)" + f"Run {i} took {t:.3f}s — expected < {DEPTH_WARM_LIMIT}s " + "(model should be cached)" ) return timings @@ -116,6 +123,7 @@ def _run_group( # Tests # ────────────────────────────────────────────────────────────────────────────── + def test_direct_default_model(indoor_image_bytes, tmp_path, timing_collector): """5 inferences with the default model backend (downloads to vizion3d cache).""" _run_group( @@ -128,9 +136,7 @@ def test_direct_default_model(indoor_image_bytes, tmp_path, timing_collector): ) -def test_direct_local_model( - indoor_image_bytes, local_model_path, tmp_path, timing_collector -): +def test_direct_local_model(indoor_image_bytes, local_model_path, tmp_path, timing_collector): """5 inferences with an explicit local .pth path in a tmp directory.""" _run_group( model_backend=local_model_path, @@ -140,3 +146,51 @@ def test_direct_local_model( scenario="Local model", timing_collector=timing_collector, ) + + +def test_direct_advanced_config_custom_intrinsics(indoor_image_bytes, local_model_path): + """Custom fx/fy/cx/cy produce a valid point cloud without errors.""" + DepthEstimationHandler._depth_anything_models.clear() + result = DepthEstimation().run( + DepthEstimationCommand( + image_input=indoor_image_bytes, + model_backend=local_model_path, + return_point_cloud=True, + advanced_config=DepthEstimationAdvanceConfig(fx=615.0, fy=615.0, cx=320.0, cy=240.0), + ) + ) + assert isinstance(result.depth_map, list) and len(result.depth_map) > 0 + assert result.point_cloud is not None + assert result.point_cloud.has_points() + + +def test_direct_advanced_config_tight_depth_trunc_yields_fewer_points( + indoor_image_bytes, local_model_path +): + """A very tight depth_trunc maps depth values to near-zero → fewer/no points.""" + DepthEstimationHandler._depth_anything_models.clear() + + default_result = DepthEstimation().run( + DepthEstimationCommand( + image_input=indoor_image_bytes, + model_backend=local_model_path, + return_point_cloud=True, + ) + ) + + tight_result = DepthEstimation().run( + DepthEstimationCommand( + image_input=indoor_image_bytes, + model_backend=local_model_path, + return_point_cloud=True, + advanced_config=DepthEstimationAdvanceConfig(depth_trunc=0.0001), + ) + ) + + import numpy as np + + default_pts = len(np.asarray(default_result.point_cloud.points)) + tight_pts = len(np.asarray(tight_result.point_cloud.points)) + assert tight_pts < default_pts, ( + f"Expected tight depth_trunc to yield fewer points ({tight_pts} vs default {default_pts})" + ) diff --git a/tests/integration/test_grpc.py b/tests/integration/test_depth_estimation_grpc.py similarity index 64% rename from tests/integration/test_grpc.py rename to tests/integration/test_depth_estimation_grpc.py index acb18b1..91573e5 100644 --- a/tests/integration/test_grpc.py +++ b/tests/integration/test_depth_estimation_grpc.py @@ -25,14 +25,15 @@ from vizion3d.proto import lifting_pb2 # noqa: E402 N_RUNS = 5 -COLD_LIMIT = float(os.environ.get("VIZION3D_TEST_COLD_LIMIT", "10.0")) -WARM_LIMIT = float(os.environ.get("VIZION3D_TEST_WARM_LIMIT", "1.0")) +DEPTH_COLD_LIMIT = float(os.environ.get("VIZION3D_TEST_DEPTH_COLD_LIMIT", "10.0")) +DEPTH_WARM_LIMIT = float(os.environ.get("VIZION3D_TEST_DEPTH_WARM_LIMIT", "1.0")) # ────────────────────────────────────────────────────────────────────────────── # Helpers # ────────────────────────────────────────────────────────────────────────────── + def _save_outputs(response, run_dir: Path, run: int) -> None: run_dir.mkdir(parents=True, exist_ok=True) prefix = str(run_dir / f"run_{run:02d}") @@ -69,9 +70,9 @@ def _run_group( return_point_cloud=True, ) - t0 = time.perf_counter() + t0 = time.perf_counter() response = grpc_client_stub.RunDepthEstimation(request) - elapsed = time.perf_counter() - t0 + elapsed = time.perf_counter() - t0 timings.append(elapsed) _save_outputs(response, run_dir, run) @@ -79,30 +80,30 @@ def _run_group( # ── per-run assertions ──────────────────────────────────────────────── assert len(response.depth_map) > 0, "depth_map rows are missing" - assert response.max_depth > response.min_depth, \ + assert response.max_depth > response.min_depth, ( "Depth variation expected for a real indoor scene" - assert len(response.depth_image) > 0, \ - "depth_image was requested but empty in response" - assert len(response.point_cloud_ply) > 0, \ + ) + assert len(response.depth_image) > 0, "depth_image was requested but empty in response" + assert len(response.point_cloud_ply) > 0, ( "point_cloud_ply was requested but empty in response" + ) # Validate file-format magic bytes - assert response.depth_image[:4] == b"\x89PNG", \ - "depth_image is not a valid PNG" - assert response.point_cloud_ply.startswith(b"ply\n"), \ - "point_cloud_ply is not a valid PLY" + assert response.depth_image[:4] == b"\x89PNG", "depth_image is not a valid PNG" + assert response.point_cloud_ply.startswith(b"ply\n"), "point_cloud_ply is not a valid PLY" - assert len(DepthEstimationHandler._depth_anything_models) > 0, \ + assert len(DepthEstimationHandler._depth_anything_models) > 0, ( "Model should be cached in memory after first gRPC inference" + ) - assert timings[0] < COLD_LIMIT, ( - f"[gRPC / {scenario}] " - f"Cold load took {timings[0]:.3f}s — expected < {COLD_LIMIT}s" + assert timings[0] < DEPTH_COLD_LIMIT, ( + f"[gRPC / {scenario}] Cold load took {timings[0]:.3f}s — expected < {DEPTH_COLD_LIMIT}s" ) for i, t in enumerate(timings[1:], start=2): - assert t < WARM_LIMIT, ( + assert t < DEPTH_WARM_LIMIT, ( f"[gRPC / {scenario}] " - f"Run {i} took {t:.3f}s — expected < {WARM_LIMIT}s (model should be cached)" + f"Run {i} took {t:.3f}s — expected < {DEPTH_WARM_LIMIT}s " + "(model should be cached)" ) return timings @@ -112,9 +113,8 @@ def _run_group( # Tests # ────────────────────────────────────────────────────────────────────────────── -def test_grpc_default_model( - indoor_image_bytes, grpc_client_stub, tmp_path, timing_collector -): + +def test_grpc_default_model(indoor_image_bytes, grpc_client_stub, tmp_path, timing_collector): """5 RPC calls with the default model backend.""" _run_group( model_backend=DEFAULT_DEPTH_MODEL_URL, @@ -138,3 +138,45 @@ def test_grpc_local_model( grpc_client_stub=grpc_client_stub, timing_collector=timing_collector, ) + + +def test_grpc_advanced_config_custom_intrinsics_accepted( + indoor_image_bytes, local_model_path, grpc_client_stub +): + """Custom fx/fy/cx/cy in the proto config are accepted and produce a valid response.""" + DepthEstimationHandler._depth_anything_models.clear() + + request = lifting_pb2.DepthEstimationRequest( + image_bytes=indoor_image_bytes, + model_backend=local_model_path, + return_point_cloud=True, + advanced_config=lifting_pb2.DepthEstimationAdvanceConfig( + fx=615.0, + fy=615.0, + cx=320.0, + cy=240.0, + ), + ) + response = grpc_client_stub.RunDepthEstimation(request) + + assert len(response.depth_map) > 0 + assert response.max_depth >= response.min_depth + assert len(response.point_cloud_ply) > 0 + assert response.point_cloud_ply.startswith(b"ply\n") + + +def test_grpc_advanced_config_partial_override_accepted( + indoor_image_bytes, local_model_path, grpc_client_stub +): + """A proto config with only depth_trunc set is accepted without errors.""" + DepthEstimationHandler._depth_anything_models.clear() + + request = lifting_pb2.DepthEstimationRequest( + image_bytes=indoor_image_bytes, + model_backend=local_model_path, + advanced_config=lifting_pb2.DepthEstimationAdvanceConfig(depth_trunc=5.0), + ) + response = grpc_client_stub.RunDepthEstimation(request) + + assert len(response.depth_map) > 0 + assert response.max_depth >= response.min_depth diff --git a/tests/integration/test_rest.py b/tests/integration/test_depth_estimation_rest.py similarity index 67% rename from tests/integration/test_rest.py rename to tests/integration/test_depth_estimation_rest.py index 6f6f2e7..051b2af 100644 --- a/tests/integration/test_rest.py +++ b/tests/integration/test_depth_estimation_rest.py @@ -26,8 +26,8 @@ from vizion3d.server.rest.app import app # noqa: E402 N_RUNS = 5 -COLD_LIMIT = float(os.environ.get("VIZION3D_TEST_COLD_LIMIT", "10.0")) -WARM_LIMIT = float(os.environ.get("VIZION3D_TEST_WARM_LIMIT", "1.0")) +DEPTH_COLD_LIMIT = float(os.environ.get("VIZION3D_TEST_DEPTH_COLD_LIMIT", "10.0")) +DEPTH_WARM_LIMIT = float(os.environ.get("VIZION3D_TEST_DEPTH_WARM_LIMIT", "1.0")) client = TestClient(app, raise_server_exceptions=True) @@ -36,6 +36,7 @@ # Helpers # ────────────────────────────────────────────────────────────────────────────── + def _save_outputs(data: dict, run_dir: Path, run: int) -> None: run_dir.mkdir(parents=True, exist_ok=True) prefix = str(run_dir / f"run_{run:02d}") @@ -46,9 +47,7 @@ def _save_outputs(data: dict, run_dir: Path, run: int) -> None: Path(prefix + "_depth.png").write_bytes(base64.b64decode(data["depth_image"])) if data.get("point_cloud_ply"): - Path(prefix + "_point_cloud.ply").write_bytes( - base64.b64decode(data["point_cloud_ply"]) - ) + Path(prefix + "_point_cloud.ply").write_bytes(base64.b64decode(data["point_cloud_ply"])) def _run_group( @@ -86,12 +85,15 @@ def _run_group( # ── per-run assertions ──────────────────────────────────────────────── assert isinstance(data["depth_map"], list) and len(data["depth_map"]) > 0 - assert data["max_depth"] > data["min_depth"], \ + assert data["max_depth"] > data["min_depth"], ( "Depth variation expected for a real indoor scene" - assert data["depth_image"] is not None, \ + ) + assert data["depth_image"] is not None, ( "depth_image was requested but missing from response" - assert data["point_cloud_ply"] is not None, \ + ) + assert data["point_cloud_ply"] is not None, ( "point_cloud_ply was requested but missing from response" + ) # Binary fields should decode to valid PLY / PNG png_bytes = base64.b64decode(data["depth_image"]) @@ -100,17 +102,18 @@ def _run_group( ply_bytes = base64.b64decode(data["point_cloud_ply"]) assert ply_bytes.startswith(b"ply\n"), "point_cloud_ply is not a valid PLY" - assert len(DepthEstimationHandler._depth_anything_models) > 0, \ + assert len(DepthEstimationHandler._depth_anything_models) > 0, ( "Model should be cached in memory after first REST inference" + ) - assert timings[0] < COLD_LIMIT, ( - f"[REST / {scenario}] " - f"Cold load took {timings[0]:.3f}s — expected < {COLD_LIMIT}s" + assert timings[0] < DEPTH_COLD_LIMIT, ( + f"[REST / {scenario}] Cold load took {timings[0]:.3f}s — expected < {DEPTH_COLD_LIMIT}s" ) for i, t in enumerate(timings[1:], start=2): - assert t < WARM_LIMIT, ( + assert t < DEPTH_WARM_LIMIT, ( f"[REST / {scenario}] " - f"Run {i} took {t:.3f}s — expected < {WARM_LIMIT}s (model should be cached)" + f"Run {i} took {t:.3f}s — expected < {DEPTH_WARM_LIMIT}s " + "(model should be cached)" ) return timings @@ -120,6 +123,7 @@ def _run_group( # Tests # ────────────────────────────────────────────────────────────────────────────── + def test_rest_default_model(indoor_image_bytes, tmp_path, timing_collector): """5 POST requests with the default model backend.""" _run_group( @@ -131,9 +135,7 @@ def test_rest_default_model(indoor_image_bytes, tmp_path, timing_collector): ) -def test_rest_local_model( - indoor_image_bytes, local_model_path, tmp_path, timing_collector -): +def test_rest_local_model(indoor_image_bytes, local_model_path, tmp_path, timing_collector): """5 POST requests with an explicit local .pth path.""" _run_group( model_backend=local_model_path, @@ -142,3 +144,46 @@ def test_rest_local_model( scenario="Local model", timing_collector=timing_collector, ) + + +def test_rest_advanced_config_custom_intrinsics_accepted(indoor_image_bytes, local_model_path): + """Custom fx/fy/cx/cy form fields are accepted and produce a valid response.""" + from vizion3d.lifting.handlers import DepthEstimationHandler + + DepthEstimationHandler._depth_anything_models.clear() + + response = client.post( + "/lifting/depth-estimation", + files={"image": ("scene.jpg", indoor_image_bytes, "image/jpeg")}, + data={ + "model_backend": local_model_path, + "return_point_cloud": "true", + "fx": "615.0", + "fy": "615.0", + "cx": "320.0", + "cy": "240.0", + }, + ) + assert response.status_code == 200 + data = response.json() + assert isinstance(data["depth_map"], list) and len(data["depth_map"]) > 0 + assert data["point_cloud_ply"] is not None + assert base64.b64decode(data["point_cloud_ply"]).startswith(b"ply\n") + + +def test_rest_advanced_config_custom_depth_trunc_accepted(indoor_image_bytes, local_model_path): + """Custom depth_trunc form field is accepted without errors.""" + from vizion3d.lifting.handlers import DepthEstimationHandler + + DepthEstimationHandler._depth_anything_models.clear() + + response = client.post( + "/lifting/depth-estimation", + files={"image": ("scene.jpg", indoor_image_bytes, "image/jpeg")}, + data={ + "model_backend": local_model_path, + "depth_trunc": "5.0", + }, + ) + assert response.status_code == 200 + assert response.json()["max_depth"] >= response.json()["min_depth"] diff --git a/tests/integration/test_stereo_depth_direct.py b/tests/integration/test_stereo_depth_direct.py new file mode 100644 index 0000000..0bc5cc3 --- /dev/null +++ b/tests/integration/test_stereo_depth_direct.py @@ -0,0 +1,258 @@ +""" +Integration tests — direct Python entry point for Stereo Depth (StereoDepth.run). + +Each test group clears the in-memory model cache before the first inference to +simulate a cold start, then runs N_RUNS iterations. We assert: + - every run produces a valid depth map and disparity map + - depth image and point cloud are returned and non-trivial + - all artefacts are persisted to a session tmp directory + - the model is cached in memory after run 1 (subsequent runs faster) +""" + +from __future__ import annotations + +import json +import os +import time +from pathlib import Path + +import numpy as np +import pytest +from PIL import Image as PILImage + +pytest.importorskip("open3d", reason="open3d required — run: uv python pin 3.12 && uv sync") + +from vizion3d.lifting.utils import create_ply_binary # noqa: E402 +from vizion3d.stereo import ( # noqa: E402 + StereoDepth, + StereoDepthAdvancedConfig, + StereoDepthCommand, +) +from vizion3d.stereo.defaults import DEFAULT_STEREO_MODEL_URL # noqa: E402 +from vizion3d.stereo.handlers import StereoDepthHandler # noqa: E402 + +N_RUNS = 5 +STEREO_COLD_LIMIT = float(os.environ.get("VIZION3D_TEST_STEREO_COLD_LIMIT", "60.0")) +STEREO_WARM_LIMIT = float(os.environ.get("VIZION3D_TEST_STEREO_WARM_LIMIT", "5.0")) + + +# ────────────────────────────────────────────────────────────────────────────── +# Helpers +# ────────────────────────────────────────────────────────────────────────────── + + +def _save_outputs(result, run_dir: Path, run: int) -> None: + run_dir.mkdir(parents=True, exist_ok=True) + prefix = run_dir / f"run_{run:02d}" + + Path(str(prefix) + ".depth_map.json").write_text(json.dumps(result.depth_map)) + Path(str(prefix) + ".disparity_map.json").write_text(json.dumps(result.disparity_map)) + + if result.depth_image is not None: + arr = np.asarray(result.depth_image) + PILImage.fromarray(arr).save(str(prefix) + "_depth.png") + + if result.point_cloud is not None and result.point_cloud.has_points(): + pts = np.asarray(result.point_cloud.points).astype(np.float32) + cols = (np.asarray(result.point_cloud.colors) * 255).astype(np.uint8) + ply = create_ply_binary(pts, cols) + Path(str(prefix) + "_point_cloud.ply").write_bytes(ply) + + +def _run_group( + model_backend: str, + stereo_image_pair: tuple[bytes, bytes], + stereo_advanced_config: StereoDepthAdvancedConfig, + run_dir: Path, + entry_point: str, + scenario: str, + timing_collector, +) -> list[float]: + StereoDepthHandler._stereo_models.clear() + + left_bytes, right_bytes = stereo_image_pair + timings: list[float] = [] + + for run in range(1, N_RUNS + 1): + t0 = time.perf_counter() + result = StereoDepth().run( + StereoDepthCommand( + left_image=left_bytes, + right_image=right_bytes, + model_backend=model_backend, + return_depth_image=True, + return_point_cloud=True, + advanced_config=stereo_advanced_config, + ) + ) + elapsed = time.perf_counter() - t0 + timings.append(elapsed) + + _save_outputs(result, run_dir, run) + timing_collector.add(entry_point, scenario, run, elapsed, str(run_dir)) + + assert isinstance(result.depth_map, list) and len(result.depth_map) > 0 + assert all(isinstance(row, list) for row in result.depth_map) + assert isinstance(result.disparity_map, list) and len(result.disparity_map) > 0 + assert result.max_depth >= result.min_depth + assert result.depth_image is not None, "depth_image was requested but missing" + assert result.point_cloud is not None, "point_cloud was requested but missing" + + assert len(StereoDepthHandler._stereo_models) > 0, ( + "Model should be cached in StereoDepthHandler._stereo_models after first inference" + ) + + assert timings[0] < STEREO_COLD_LIMIT, ( + f"[{entry_point} / {scenario}] Cold load took {timings[0]:.3f}s " + f"— expected < {STEREO_COLD_LIMIT}s" + ) + for i, t in enumerate(timings[1:], start=2): + assert t < STEREO_WARM_LIMIT, ( + f"[{entry_point} / {scenario}] " + f"Run {i} took {t:.3f}s — expected < {STEREO_WARM_LIMIT}s " + "(model should be cached)" + ) + + return timings + + +# ────────────────────────────────────────────────────────────────────────────── +# Tests +# ────────────────────────────────────────────────────────────────────────────── + + +def test_stereo_direct_default_model( + stereo_image_pair, stereo_advanced_config, tmp_path, timing_collector +): + """5 direct inferences with the default model backend (downloads to cache).""" + _run_group( + model_backend=DEFAULT_STEREO_MODEL_URL, + stereo_image_pair=stereo_image_pair, + stereo_advanced_config=stereo_advanced_config, + run_dir=tmp_path / "stereo_direct_default", + entry_point="Stereo Direct", + scenario="Default model", + timing_collector=timing_collector, + ) + + +def test_stereo_direct_local_model( + stereo_image_pair, + stereo_advanced_config, + local_stereo_model_path, + tmp_path, + timing_collector, +): + """5 direct inferences with an explicit local .pth path.""" + _run_group( + model_backend=local_stereo_model_path, + stereo_image_pair=stereo_image_pair, + stereo_advanced_config=stereo_advanced_config, + run_dir=tmp_path / "stereo_direct_local", + entry_point="Stereo Direct", + scenario="Local model", + timing_collector=timing_collector, + ) + + +def test_stereo_direct_advanced_config_custom_intrinsics( + stereo_image_pair, stereo_advanced_config, local_stereo_model_path +): + """The calibrated Middlebury intrinsics produce a valid point cloud.""" + StereoDepthHandler._stereo_models.clear() + left_bytes, right_bytes = stereo_image_pair + + result = StereoDepth().run( + StereoDepthCommand( + left_image=left_bytes, + right_image=right_bytes, + model_backend=local_stereo_model_path, + return_point_cloud=True, + advanced_config=stereo_advanced_config, + ) + ) + assert isinstance(result.depth_map, list) and len(result.depth_map) > 0 + assert result.max_depth >= result.min_depth + assert result.point_cloud is not None + + +def test_stereo_direct_advanced_config_tight_z_far_yields_fewer_points( + stereo_image_pair, stereo_advanced_config, local_stereo_model_path +): + """A tight z_far clips points to a smaller range → fewer points than the default.""" + StereoDepthHandler._stereo_models.clear() + left_bytes, right_bytes = stereo_image_pair + + default_result = StereoDepth().run( + StereoDepthCommand( + left_image=left_bytes, + right_image=right_bytes, + model_backend=local_stereo_model_path, + return_point_cloud=True, + advanced_config=stereo_advanced_config.model_copy(update={"z_far": 100.0}), + ) + ) + + tight_result = StereoDepth().run( + StereoDepthCommand( + left_image=left_bytes, + right_image=right_bytes, + model_backend=local_stereo_model_path, + return_point_cloud=True, + advanced_config=stereo_advanced_config.model_copy(update={"z_far": 0.001}), + ) + ) + + default_pts = len(np.asarray(default_result.point_cloud.points)) + tight_pts = len(np.asarray(tight_result.point_cloud.points)) + assert tight_pts < default_pts, ( + f"Expected tight z_far to yield fewer points ({tight_pts} vs default {default_pts})" + ) + + +def test_stereo_direct_advanced_config_scale_factor( + stereo_image_pair, stereo_advanced_config, local_stereo_model_path +): + """scale_factor=0.5 completes without error and returns a valid depth map.""" + StereoDepthHandler._stereo_models.clear() + left_bytes, right_bytes = stereo_image_pair + + result = StereoDepth().run( + StereoDepthCommand( + left_image=left_bytes, + right_image=right_bytes, + model_backend=local_stereo_model_path, + advanced_config=stereo_advanced_config.model_copy(update={"scale_factor": 0.5}), + ) + ) + assert isinstance(result.depth_map, list) and len(result.depth_map) > 0 + assert result.max_depth >= result.min_depth + + +def test_stereo_direct_all_outputs_returned( + stereo_image_pair, stereo_advanced_config, local_stereo_model_path +): + """depth_image, point_cloud, and mesh are all returned when requested.""" + StereoDepthHandler._stereo_models.clear() + left_bytes, right_bytes = stereo_image_pair + + result = StereoDepth().run( + StereoDepthCommand( + left_image=left_bytes, + right_image=right_bytes, + model_backend=local_stereo_model_path, + return_depth_image=True, + return_point_cloud=True, + return_mesh=True, + advanced_config=stereo_advanced_config, + ) + ) + assert result.depth_image is not None + arr = np.asarray(result.depth_image) + assert arr.dtype == np.uint16 + + assert result.point_cloud is not None + assert result.point_cloud.has_points() + assert result.point_cloud_scale == 1.0 + + assert result.mesh is not None diff --git a/tests/integration/test_stereo_depth_grpc.py b/tests/integration/test_stereo_depth_grpc.py new file mode 100644 index 0000000..47231ab --- /dev/null +++ b/tests/integration/test_stereo_depth_grpc.py @@ -0,0 +1,156 @@ +""" +Integration tests for the Stereo Depth gRPC entry point. + +A live in-process gRPC server is provided by the session fixture. These tests +exercise real S2M2 inference through RunStereoDepth. +""" + +from __future__ import annotations + +import json +import os +import time +from pathlib import Path + +import pytest + +pytest.importorskip("open3d", reason="open3d required - run: uv python pin 3.12 && uv sync") + +from vizion3d.proto import lifting_pb2 # noqa: E402 +from vizion3d.stereo.defaults import DEFAULT_STEREO_MODEL_URL # noqa: E402 +from vizion3d.stereo.handlers import StereoDepthHandler # noqa: E402 + +N_RUNS = 5 +STEREO_COLD_LIMIT = float(os.environ.get("VIZION3D_TEST_STEREO_COLD_LIMIT", "60.0")) +STEREO_WARM_LIMIT = float(os.environ.get("VIZION3D_TEST_STEREO_WARM_LIMIT", "5.0")) + + +def _proto_config(config, **overrides): + values = { + "focal_length": config.focal_length, + "baseline": config.baseline, + "cx": config.cx, + "cy": config.cy, + "doffs": config.doffs, + "z_far": config.z_far, + "conf_threshold": config.conf_threshold, + "occ_threshold": config.occ_threshold, + "scale_factor": config.scale_factor, + } + values.update(overrides) + return lifting_pb2.StereoDepthAdvancedConfig(**values) + + +def _save_outputs(response, run_dir: Path, run: int) -> None: + run_dir.mkdir(parents=True, exist_ok=True) + prefix = str(run_dir / f"run_{run:02d}") + + depth_map = [list(row.values) for row in response.depth_map] + disparity_map = [list(row.values) for row in response.disparity_map] + Path(prefix + ".depth_map.json").write_text(json.dumps(depth_map)) + Path(prefix + ".disparity_map.json").write_text(json.dumps(disparity_map)) + + if response.depth_image: + Path(prefix + "_depth.png").write_bytes(response.depth_image) + + if response.point_cloud_ply: + Path(prefix + "_point_cloud.ply").write_bytes(response.point_cloud_ply) + + +def _run_group( + model_backend: str, + stereo_image_pair: tuple[bytes, bytes], + stereo_advanced_config, + run_dir: Path, + scenario: str, + grpc_client_stub, + timing_collector, +) -> list[float]: + StereoDepthHandler._stereo_models.clear() + left_bytes, right_bytes = stereo_image_pair + timings: list[float] = [] + + for run in range(1, N_RUNS + 1): + request = lifting_pb2.StereoDepthRequest( + left_image_bytes=left_bytes, + right_image_bytes=right_bytes, + model_backend=model_backend, + return_depth_image=True, + return_point_cloud=True, + advanced_config=_proto_config(stereo_advanced_config), + ) + + t0 = time.perf_counter() + response = grpc_client_stub.RunStereoDepth(request) + elapsed = time.perf_counter() - t0 + timings.append(elapsed) + + _save_outputs(response, run_dir, run) + timing_collector.add("Stereo gRPC", scenario, run, elapsed, str(run_dir)) + + assert len(response.depth_map) > 0 + assert len(response.disparity_map) > 0 + assert response.max_depth >= response.min_depth + assert response.depth_image[:4] == b"\x89PNG" + assert response.point_cloud_ply.startswith(b"ply\n") + + assert len(StereoDepthHandler._stereo_models) > 0 + assert timings[0] < STEREO_COLD_LIMIT + for elapsed in timings[1:]: + assert elapsed < STEREO_WARM_LIMIT + + return timings + + +def test_stereo_grpc_default_model( + stereo_image_pair, stereo_advanced_config, grpc_client_stub, tmp_path, timing_collector +): + _run_group( + model_backend=DEFAULT_STEREO_MODEL_URL, + stereo_image_pair=stereo_image_pair, + stereo_advanced_config=stereo_advanced_config, + run_dir=tmp_path / "stereo_grpc_default", + scenario="Default model", + grpc_client_stub=grpc_client_stub, + timing_collector=timing_collector, + ) + + +def test_stereo_grpc_local_model( + stereo_image_pair, + stereo_advanced_config, + local_stereo_model_path, + grpc_client_stub, + tmp_path, + timing_collector, +): + _run_group( + model_backend=local_stereo_model_path, + stereo_image_pair=stereo_image_pair, + stereo_advanced_config=stereo_advanced_config, + run_dir=tmp_path / "stereo_grpc_local", + scenario="Local model", + grpc_client_stub=grpc_client_stub, + timing_collector=timing_collector, + ) + + +def test_stereo_grpc_advanced_config_accepted( + stereo_image_pair, stereo_advanced_config, local_stereo_model_path, grpc_client_stub +): + StereoDepthHandler._stereo_models.clear() + left_bytes, right_bytes = stereo_image_pair + + request = lifting_pb2.StereoDepthRequest( + left_image_bytes=left_bytes, + right_image_bytes=right_bytes, + model_backend=local_stereo_model_path, + return_point_cloud=True, + advanced_config=_proto_config(stereo_advanced_config, scale_factor=0.5), + ) + response = grpc_client_stub.RunStereoDepth(request) + + assert len(response.depth_map) > 0 + assert len(response.disparity_map) > 0 + assert response.max_depth >= response.min_depth + assert response.point_cloud_ply.startswith(b"ply\n") diff --git a/tests/integration/test_stereo_depth_rest.py b/tests/integration/test_stereo_depth_rest.py new file mode 100644 index 0000000..9ce3dff --- /dev/null +++ b/tests/integration/test_stereo_depth_rest.py @@ -0,0 +1,167 @@ +""" +Integration tests for the Stereo Depth REST entry point. + +These tests exercise real S2M2 inference through POST /lifting/stereo-depth, +covering both default URL download/cache behavior and explicit local model paths. +""" + +from __future__ import annotations + +import base64 +import json +import os +import time +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +pytest.importorskip("open3d", reason="open3d required - run: uv python pin 3.12 && uv sync") + +from vizion3d.server.rest.app import app # noqa: E402 +from vizion3d.stereo.defaults import DEFAULT_STEREO_MODEL_URL # noqa: E402 +from vizion3d.stereo.handlers import StereoDepthHandler # noqa: E402 + +N_RUNS = 5 +STEREO_COLD_LIMIT = float(os.environ.get("VIZION3D_TEST_STEREO_COLD_LIMIT", "60.0")) +STEREO_WARM_LIMIT = float(os.environ.get("VIZION3D_TEST_STEREO_WARM_LIMIT", "5.0")) + +client = TestClient(app, raise_server_exceptions=True) + + +def _save_outputs(data: dict, run_dir: Path, run: int) -> None: + run_dir.mkdir(parents=True, exist_ok=True) + prefix = str(run_dir / f"run_{run:02d}") + + Path(prefix + ".depth_map.json").write_text(json.dumps(data["depth_map"])) + Path(prefix + ".disparity_map.json").write_text(json.dumps(data["disparity_map"])) + + if data.get("depth_image"): + Path(prefix + "_depth.png").write_bytes(base64.b64decode(data["depth_image"])) + + if data.get("point_cloud_ply"): + Path(prefix + "_point_cloud.ply").write_bytes(base64.b64decode(data["point_cloud_ply"])) + + +def _run_group( + model_backend: str, + stereo_image_pair: tuple[bytes, bytes], + stereo_advanced_config, + run_dir: Path, + scenario: str, + timing_collector, +) -> list[float]: + StereoDepthHandler._stereo_models.clear() + left_bytes, right_bytes = stereo_image_pair + timings: list[float] = [] + + for run in range(1, N_RUNS + 1): + t0 = time.perf_counter() + response = client.post( + "/lifting/stereo-depth", + files={ + "left_image": ("left.jpg", left_bytes, "image/jpeg"), + "right_image": ("right.jpg", right_bytes, "image/jpeg"), + }, + data={ + "model_backend": model_backend, + "return_depth_image": "true", + "return_point_cloud": "true", + "focal_length": str(stereo_advanced_config.focal_length), + "baseline": str(stereo_advanced_config.baseline), + "cx": str(stereo_advanced_config.cx), + "cy": str(stereo_advanced_config.cy), + "doffs": str(stereo_advanced_config.doffs), + "z_far": str(stereo_advanced_config.z_far), + "conf_threshold": str(stereo_advanced_config.conf_threshold), + "occ_threshold": str(stereo_advanced_config.occ_threshold), + }, + ) + elapsed = time.perf_counter() - t0 + timings.append(elapsed) + + assert response.status_code == 200, ( + f"Run {run} failed: {response.status_code} - {response.text[:300]}" + ) + data = response.json() + _save_outputs(data, run_dir, run) + timing_collector.add("Stereo REST", scenario, run, elapsed, str(run_dir)) + + assert isinstance(data["depth_map"], list) and len(data["depth_map"]) > 0 + assert isinstance(data["disparity_map"], list) and len(data["disparity_map"]) > 0 + assert data["max_depth"] >= data["min_depth"] + assert data["depth_image"] is not None + assert data["point_cloud_ply"] is not None + assert base64.b64decode(data["depth_image"])[:4] == b"\x89PNG" + assert base64.b64decode(data["point_cloud_ply"]).startswith(b"ply\n") + + assert len(StereoDepthHandler._stereo_models) > 0 + assert timings[0] < STEREO_COLD_LIMIT + for elapsed in timings[1:]: + assert elapsed < STEREO_WARM_LIMIT + + return timings + + +def test_stereo_rest_default_model( + stereo_image_pair, stereo_advanced_config, tmp_path, timing_collector +): + _run_group( + model_backend=DEFAULT_STEREO_MODEL_URL, + stereo_image_pair=stereo_image_pair, + stereo_advanced_config=stereo_advanced_config, + run_dir=tmp_path / "stereo_rest_default", + scenario="Default model", + timing_collector=timing_collector, + ) + + +def test_stereo_rest_local_model( + stereo_image_pair, + stereo_advanced_config, + local_stereo_model_path, + tmp_path, + timing_collector, +): + _run_group( + model_backend=local_stereo_model_path, + stereo_image_pair=stereo_image_pair, + stereo_advanced_config=stereo_advanced_config, + run_dir=tmp_path / "stereo_rest_local", + scenario="Local model", + timing_collector=timing_collector, + ) + + +def test_stereo_rest_advanced_config_accepted( + stereo_image_pair, stereo_advanced_config, local_stereo_model_path +): + StereoDepthHandler._stereo_models.clear() + left_bytes, right_bytes = stereo_image_pair + + response = client.post( + "/lifting/stereo-depth", + files={ + "left_image": ("left.jpg", left_bytes, "image/jpeg"), + "right_image": ("right.jpg", right_bytes, "image/jpeg"), + }, + data={ + "model_backend": local_stereo_model_path, + "return_point_cloud": "true", + "focal_length": str(stereo_advanced_config.focal_length), + "baseline": str(stereo_advanced_config.baseline), + "cx": str(stereo_advanced_config.cx), + "cy": str(stereo_advanced_config.cy), + "doffs": str(stereo_advanced_config.doffs), + "z_far": str(stereo_advanced_config.z_far), + "conf_threshold": str(stereo_advanced_config.conf_threshold), + "occ_threshold": str(stereo_advanced_config.occ_threshold), + "scale_factor": "0.5", + }, + ) + + assert response.status_code == 200 + data = response.json() + assert isinstance(data["depth_map"], list) and len(data["depth_map"]) > 0 + assert data["point_cloud_ply"] is not None + assert base64.b64decode(data["point_cloud_ply"]).startswith(b"ply\n") diff --git a/tests/unit/test_depth_estimation_advanced_config.py b/tests/unit/test_depth_estimation_advanced_config.py new file mode 100644 index 0000000..1b6e432 --- /dev/null +++ b/tests/unit/test_depth_estimation_advanced_config.py @@ -0,0 +1,368 @@ +""" +Unit tests for DepthEstimationAdvanceConfig and its propagation through +the command, handler, REST API, and gRPC server. +""" + +import io +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +from PIL import Image + +from vizion3d.lifting.commands import DepthEstimationCommand +from vizion3d.lifting.handlers import DepthEstimationHandler +from vizion3d.lifting.models import DepthEstimationAdvanceConfig + +# ── Fixtures ────────────────────────────────────────────────────────────────── + + +@pytest.fixture +def fake_depth(): + rng = np.random.default_rng(0) + return rng.uniform(0.5, 3.0, (100, 100)) + + +# ── DepthEstimationAdvanceConfig ────────────────────────────────────────────── + + +class TestDepthEstimationAdvanceConfig: + def test_defaults_match_primesense(self): + cfg = DepthEstimationAdvanceConfig() + assert cfg.fx == 525.0 + assert cfg.fy == 525.0 + assert cfg.cx == 319.5 + assert cfg.cy == 239.5 + assert cfg.depth_scale == 1000.0 + assert cfg.depth_trunc == 10.0 + + def test_partial_override_keeps_other_defaults(self): + cfg = DepthEstimationAdvanceConfig(fx=615.0, fy=615.0) + assert cfg.fx == 615.0 + assert cfg.fy == 615.0 + assert cfg.cx == 319.5 + assert cfg.cy == 239.5 + assert cfg.depth_scale == 1000.0 + assert cfg.depth_trunc == 10.0 + + def test_all_fields_overridable(self): + cfg = DepthEstimationAdvanceConfig( + fx=700.0, + fy=701.0, + cx=320.0, + cy=241.0, + depth_scale=500.0, + depth_trunc=5.0, + ) + assert cfg.fx == 700.0 + assert cfg.fy == 701.0 + assert cfg.cx == 320.0 + assert cfg.cy == 241.0 + assert cfg.depth_scale == 500.0 + assert cfg.depth_trunc == 5.0 + + def test_invalid_type_rejected(self): + with pytest.raises(Exception): + DepthEstimationAdvanceConfig(fx="not_a_number") + + +# ── DepthEstimationCommand advanced_config field ────────────────────────────── + + +class TestDepthEstimationCommandAdvancedConfig: + def test_default_advanced_config_is_correct_type(self): + cmd = DepthEstimationCommand(image_input=b"img") + assert isinstance(cmd.advanced_config, DepthEstimationAdvanceConfig) + + def test_default_advanced_config_has_primesense_values(self): + cmd = DepthEstimationCommand(image_input=b"img") + assert cmd.advanced_config.fx == 525.0 + assert cmd.advanced_config.depth_trunc == 10.0 + + def test_each_command_gets_separate_config_instance(self): + cmd1 = DepthEstimationCommand(image_input=b"img") + cmd2 = DepthEstimationCommand(image_input=b"img") + assert cmd1.advanced_config is not cmd2.advanced_config + + def test_mutation_of_one_command_does_not_affect_another(self): + cmd1 = DepthEstimationCommand(image_input=b"img") + cmd2 = DepthEstimationCommand(image_input=b"img") + cmd1.advanced_config.fx = 999.0 + assert cmd2.advanced_config.fx == 525.0 + + def test_custom_config_stored_on_command(self): + cfg = DepthEstimationAdvanceConfig(fx=615.0, depth_trunc=5.0) + cmd = DepthEstimationCommand(image_input=b"img", advanced_config=cfg) + assert cmd.advanced_config.fx == 615.0 + assert cmd.advanced_config.depth_trunc == 5.0 + + +# ── _depth_array_to_rgbd_depth ──────────────────────────────────────────────── + + +class TestDepthArrayToRgbdDepth: + def test_output_dtype_is_uint16(self): + arr = np.array([[0.0, 1.0], [2.0, 3.0]], dtype=float) + result = DepthEstimationHandler._depth_array_to_rgbd_depth(arr, 1000.0, 10.0) + assert result.dtype == np.uint16 + + def test_zero_range_input_returns_zeros(self): + arr = np.ones((4, 4), dtype=float) * 3.0 + result = DepthEstimationHandler._depth_array_to_rgbd_depth(arr, 1000.0, 10.0) + assert np.all(result == 0) + + def test_max_value_equals_trunc_times_scale(self): + arr = np.array([[0.0, 1.0]], dtype=float) + result = DepthEstimationHandler._depth_array_to_rgbd_depth(arr, 1000.0, 10.0) + assert result[0, 1] == pytest.approx(10_000, abs=1) + + def test_custom_depth_scale_halves_output(self): + arr = np.array([[0.0, 1.0]], dtype=float) + full = DepthEstimationHandler._depth_array_to_rgbd_depth(arr, 1000.0, 10.0) + half = DepthEstimationHandler._depth_array_to_rgbd_depth(arr, 500.0, 10.0) + assert half[0, 1] == pytest.approx(full[0, 1] / 2, abs=1) + + def test_custom_depth_trunc_halves_output(self): + arr = np.array([[0.0, 1.0]], dtype=float) + full = DepthEstimationHandler._depth_array_to_rgbd_depth(arr, 1000.0, 10.0) + half = DepthEstimationHandler._depth_array_to_rgbd_depth(arr, 1000.0, 5.0) + assert half[0, 1] == pytest.approx(full[0, 1] / 2, abs=1) + + def test_output_clipped_to_uint16_max(self): + arr = np.array([[0.0, 1.0]], dtype=float) + result = DepthEstimationHandler._depth_array_to_rgbd_depth(arr, 1e6, 1e6) + assert result[0, 1] == np.iinfo(np.uint16).max + + def test_min_value_is_always_zero(self): + arr = np.array([[2.0, 5.0, 10.0]], dtype=float) + result = DepthEstimationHandler._depth_array_to_rgbd_depth(arr, 1000.0, 10.0) + assert result[0, 0] == 0 + + +# ── Handler propagates config to Open3D ────────────────────────────────────── + + +class TestHandlerPropagatesAdvancedConfig: + def test_custom_intrinsics_forwarded_to_pinhole(self, dummy_image_bytes, fake_depth): + open3d = pytest.importorskip( + "open3d", reason="open3d required — run: uv python pin 3.12 && uv sync" + ) + cfg = DepthEstimationAdvanceConfig(fx=615.0, fy=616.0, cx=321.0, cy=241.0) + captured = [] + original_cls = open3d.camera.PinholeCameraIntrinsic + + def capturing(*args, **kwargs): + captured.append(args) + return original_cls(*args, **kwargs) + + with ( + patch.object( + DepthEstimationHandler, + "_run_depth_anything_checkpoint", + return_value=fake_depth, + ), + patch.object(open3d.camera, "PinholeCameraIntrinsic", side_effect=capturing), + ): + DepthEstimationHandler().handle( + DepthEstimationCommand( + image_input=dummy_image_bytes, + model_backend="/fake/model.pth", + return_point_cloud=True, + advanced_config=cfg, + ) + ) + + assert len(captured) == 1 + _w, _h, fx, fy, cx, cy = captured[0] + assert fx == 615.0 + assert fy == 616.0 + assert cx == 321.0 + assert cy == 241.0 + + def test_custom_depth_scale_trunc_forwarded_to_rgbd(self, dummy_image_bytes, fake_depth): + open3d = pytest.importorskip( + "open3d", reason="open3d required — run: uv python pin 3.12 && uv sync" + ) + cfg = DepthEstimationAdvanceConfig(depth_scale=500.0, depth_trunc=3.0) + captured_kwargs = {} + original_fn = open3d.geometry.RGBDImage.create_from_color_and_depth + + def capturing(*args, **kwargs): + captured_kwargs.update(kwargs) + return original_fn(*args, **kwargs) + + with ( + patch.object( + DepthEstimationHandler, + "_run_depth_anything_checkpoint", + return_value=fake_depth, + ), + patch.object( + open3d.geometry.RGBDImage, + "create_from_color_and_depth", + side_effect=capturing, + ), + ): + DepthEstimationHandler().handle( + DepthEstimationCommand( + image_input=dummy_image_bytes, + model_backend="/fake/model.pth", + return_point_cloud=True, + advanced_config=cfg, + ) + ) + + assert captured_kwargs["depth_scale"] == 500.0 + assert captured_kwargs["depth_trunc"] == 3.0 + + +# ── REST API propagates config form fields ──────────────────────────────────── + + +class TestRestAdvancedConfig: + @pytest.fixture(autouse=True) + def _setup(self): + pytest.importorskip("open3d", reason="open3d required — run: uv python pin 3.12 && uv sync") + from fastapi.testclient import TestClient + + from vizion3d.server.rest.app import app + + self.client = TestClient(app) + + def _post(self, extra_data=None): + img = Image.new("RGB", (50, 50), color="blue") + buf = io.BytesIO() + img.save(buf, format="PNG") + buf.seek(0) + + result = MagicMock() + result.depth_map = [[1.0]] + result.min_depth = 1.0 + result.max_depth = 1.0 + result.backend_used = "/fake/model.pth" + result.depth_image = None + result.point_cloud = None + result.mesh = None + + with patch("vizion3d.server.rest.depth_estimation.DepthEstimation") as mock_cls: + mock_cls.return_value.run.return_value = result + self.client.post( + "/lifting/depth-estimation", + files={"image": ("test.png", buf, "image/png")}, + data=extra_data or {}, + ) + return mock_cls.return_value.run.call_args[0][0] + + def test_default_config_used_when_no_form_fields_sent(self): + cmd = self._post() + assert cmd.advanced_config.fx == 525.0 + assert cmd.advanced_config.fy == 525.0 + assert cmd.advanced_config.cx == 319.5 + assert cmd.advanced_config.cy == 239.5 + assert cmd.advanced_config.depth_scale == 1000.0 + assert cmd.advanced_config.depth_trunc == 10.0 + + def test_custom_fx_fy_forwarded(self): + cmd = self._post({"fx": "615.0", "fy": "616.0"}) + assert cmd.advanced_config.fx == pytest.approx(615.0) + assert cmd.advanced_config.fy == pytest.approx(616.0) + assert cmd.advanced_config.cx == 319.5 # unchanged + + def test_custom_cx_cy_forwarded(self): + cmd = self._post({"cx": "321.0", "cy": "241.0"}) + assert cmd.advanced_config.cx == pytest.approx(321.0) + assert cmd.advanced_config.cy == pytest.approx(241.0) + + def test_custom_depth_scale_forwarded(self): + cmd = self._post({"depth_scale": "500.0"}) + assert cmd.advanced_config.depth_scale == pytest.approx(500.0) + assert cmd.advanced_config.depth_trunc == 10.0 # unchanged + + def test_custom_depth_trunc_forwarded(self): + cmd = self._post({"depth_trunc": "5.0"}) + assert cmd.advanced_config.depth_trunc == pytest.approx(5.0) + assert cmd.advanced_config.depth_scale == 1000.0 # unchanged + + def test_partial_override_does_not_affect_other_fields(self): + cmd = self._post({"fx": "700.0"}) + assert cmd.advanced_config.fx == pytest.approx(700.0) + assert cmd.advanced_config.fy == 525.0 + assert cmd.advanced_config.depth_trunc == 10.0 + + +# ── gRPC server propagates config proto fields ──────────────────────────────── + + +class TestGrpcAdvancedConfig: + @pytest.fixture(autouse=True) + def _setup(self): + pytest.importorskip("open3d", reason="open3d required — run: uv python pin 3.12 && uv sync") + from vizion3d.proto import lifting_pb2 + from vizion3d.server.grpc.server import LiftingServiceServicer + + self.pb2 = lifting_pb2 + self.servicer = LiftingServiceServicer() + self.context = MagicMock() + + img = Image.new("RGB", (50, 50), color="red") + buf = io.BytesIO() + img.save(buf, format="PNG") + self.image_bytes = buf.getvalue() + + def _run(self, proto_cfg=None): + kwargs = {"image_bytes": self.image_bytes} + if proto_cfg is not None: + kwargs["advanced_config"] = proto_cfg + request = self.pb2.DepthEstimationRequest(**kwargs) + + result = MagicMock() + result.depth_map = [[1.0]] + result.min_depth = 1.0 + result.max_depth = 1.0 + result.backend_used = "/fake/model.pth" + result.depth_image = None + result.point_cloud = None + result.mesh = None + + with patch("vizion3d.server.grpc.server.DepthEstimation") as mock_cls: + mock_cls.return_value.run.return_value = result + self.servicer.RunDepthEstimation(request, self.context) + return mock_cls.return_value.run.call_args[0][0] + + def test_no_config_in_request_uses_defaults(self): + cmd = self._run() + assert cmd.advanced_config.fx == 525.0 + assert cmd.advanced_config.depth_trunc == 10.0 + + def test_full_config_forwarded(self): + proto_cfg = self.pb2.DepthEstimationAdvanceConfig( + fx=615.0, + fy=616.0, + cx=320.0, + cy=241.0, + depth_scale=500.0, + depth_trunc=3.0, + ) + cmd = self._run(proto_cfg) + assert cmd.advanced_config.fx == pytest.approx(615.0) + assert cmd.advanced_config.fy == pytest.approx(616.0) + assert cmd.advanced_config.cx == pytest.approx(320.0) + assert cmd.advanced_config.cy == pytest.approx(241.0) + assert cmd.advanced_config.depth_scale == pytest.approx(500.0) + assert cmd.advanced_config.depth_trunc == pytest.approx(3.0) + + def test_partial_config_overrides_only_set_fields(self): + proto_cfg = self.pb2.DepthEstimationAdvanceConfig(fx=700.0, depth_trunc=2.0) + cmd = self._run(proto_cfg) + assert cmd.advanced_config.fx == pytest.approx(700.0) + assert cmd.advanced_config.fy == 525.0 # default + assert cmd.advanced_config.cx == 319.5 # default + assert cmd.advanced_config.depth_scale == 1000.0 # default + assert cmd.advanced_config.depth_trunc == pytest.approx(2.0) + + def test_empty_config_message_uses_all_defaults(self): + proto_cfg = self.pb2.DepthEstimationAdvanceConfig() + cmd = self._run(proto_cfg) + assert cmd.advanced_config.fx == 525.0 + assert cmd.advanced_config.fy == 525.0 + assert cmd.advanced_config.depth_trunc == 10.0 diff --git a/tests/unit/test_lifting_defaults.py b/tests/unit/test_depth_estimation_defaults.py similarity index 92% rename from tests/unit/test_lifting_defaults.py rename to tests/unit/test_depth_estimation_defaults.py index 65ca3af..4ad7101 100644 --- a/tests/unit/test_lifting_defaults.py +++ b/tests/unit/test_depth_estimation_defaults.py @@ -38,14 +38,8 @@ def test_loaded_checkpoint_cache_is_shared_across_handler_instances(): try: DepthEstimationHandler._depth_anything_models = {cache_key: cached_model} - assert ( - DepthEstimationHandler()._load_depth_anything_checkpoint(cache_key) - is cached_model - ) - assert ( - DepthEstimationHandler()._load_depth_anything_checkpoint(cache_key) - is cached_model - ) + assert DepthEstimationHandler()._load_depth_anything_checkpoint(cache_key) is cached_model + assert DepthEstimationHandler()._load_depth_anything_checkpoint(cache_key) is cached_model finally: DepthEstimationHandler._depth_anything_models = original_cache diff --git a/tests/unit/test_facade.py b/tests/unit/test_depth_estimation_facade.py similarity index 99% rename from tests/unit/test_facade.py rename to tests/unit/test_depth_estimation_facade.py index 3764f2f..9782d64 100644 --- a/tests/unit/test_facade.py +++ b/tests/unit/test_depth_estimation_facade.py @@ -32,6 +32,7 @@ def test_depth_estimation_returns_depth_image(dummy_image_bytes): assert isinstance(res.depth_image, o3d.geometry.Image) import numpy as np + arr = np.asarray(res.depth_image) assert arr.dtype == np.uint16 assert arr.shape == (100, 100) diff --git a/tests/unit/test_grpc.py b/tests/unit/test_depth_estimation_grpc.py similarity index 100% rename from tests/unit/test_grpc.py rename to tests/unit/test_depth_estimation_grpc.py diff --git a/tests/unit/test_depth_estimation_stereo_depth_rest_api.py b/tests/unit/test_depth_estimation_stereo_depth_rest_api.py new file mode 100644 index 0000000..b080f47 --- /dev/null +++ b/tests/unit/test_depth_estimation_stereo_depth_rest_api.py @@ -0,0 +1,335 @@ +import io +from unittest.mock import MagicMock, patch + +import pytest +from fastapi.testclient import TestClient +from PIL import Image + +pytest.importorskip("open3d", reason="open3d required — run: uv python pin 3.12 && uv sync") + +from vizion3d.lifting.defaults import DEFAULT_DEPTH_MODEL_URL # noqa: E402 +from vizion3d.server.rest.app import app, create_app, run # noqa: E402 +from vizion3d.stereo.defaults import DEFAULT_STEREO_MODEL_URL # noqa: E402 + +client = TestClient(app) + + +@pytest.fixture +def image_file(): + img = Image.new("RGB", (50, 50), color="green") + buf = io.BytesIO() + img.save(buf, format="PNG") + buf.seek(0) + return buf + + +@pytest.fixture +def stereo_files(): + img = Image.new("RGB", (50, 50), color="green") + left = io.BytesIO() + right = io.BytesIO() + img.save(left, format="PNG") + img.save(right, format="PNG") + left.seek(0) + right.seek(0) + return left, right + + +def _fake_result(depth_map=None): + result = MagicMock() + result.depth_map = depth_map or [[1.0, 2.0], [3.0, 4.0]] + result.min_depth = 1.0 + result.max_depth = 4.0 + result.backend_used = "/fake/model.pth" + result.depth_image = None + result.point_cloud = None + result.mesh = None + return result + + +def _fake_stereo_result(): + result = _fake_result() + result.disparity_map = [[4.0, 5.0], [6.0, 7.0]] + result.backend_used = "/fake/stereo.pth" + return result + + +def _uvicorn_app_routes(uvicorn_run) -> set[str]: + app_arg = uvicorn_run.call_args.args[0] + return {route.path for route in app_arg.routes} + + +def test_depth_estimation_returns_200(image_file): + with patch("vizion3d.server.rest.depth_estimation.DepthEstimation") as mock_cls: + mock_cls.return_value.run.return_value = _fake_result() + response = client.post( + "/lifting/depth-estimation", + files={"image": ("test.png", image_file, "image/png")}, + ) + assert response.status_code == 200 + + +def test_depth_estimation_response_has_expected_keys(image_file): + with patch("vizion3d.server.rest.depth_estimation.DepthEstimation") as mock_cls: + mock_cls.return_value.run.return_value = _fake_result() + response = client.post( + "/lifting/depth-estimation", + files={"image": ("test.png", image_file, "image/png")}, + ) + data = response.json() + assert "depth_map" in data + assert "min_depth" in data + assert "max_depth" in data + assert "backend_used" in data + + +def test_depth_estimation_depth_map_is_nested_list(image_file): + with patch("vizion3d.server.rest.depth_estimation.DepthEstimation") as mock_cls: + mock_cls.return_value.run.return_value = _fake_result() + response = client.post( + "/lifting/depth-estimation", + files={"image": ("test.png", image_file, "image/png")}, + ) + data = response.json() + assert isinstance(data["depth_map"], list) + assert all(isinstance(row, list) for row in data["depth_map"]) + + +def test_depth_estimation_max_depth_gte_min_depth(image_file): + with patch("vizion3d.server.rest.depth_estimation.DepthEstimation") as mock_cls: + mock_cls.return_value.run.return_value = _fake_result() + response = client.post( + "/lifting/depth-estimation", + files={"image": ("test.png", image_file, "image/png")}, + ) + data = response.json() + assert data["max_depth"] >= data["min_depth"] + + +def test_depth_estimation_optional_outputs_null_by_default(image_file): + with patch("vizion3d.server.rest.depth_estimation.DepthEstimation") as mock_cls: + mock_cls.return_value.run.return_value = _fake_result() + response = client.post( + "/lifting/depth-estimation", + files={"image": ("test.png", image_file, "image/png")}, + ) + data = response.json() + assert data["depth_image"] is None + assert data["point_cloud_ply"] is None + assert data["mesh_ply"] is None + + +def test_depth_estimation_backend_used_is_returned(image_file): + with patch("vizion3d.server.rest.depth_estimation.DepthEstimation") as mock_cls: + mock_cls.return_value.run.return_value = _fake_result() + response = client.post( + "/lifting/depth-estimation", + files={"image": ("test.png", image_file, "image/png")}, + ) + data = response.json() + assert data["backend_used"] == "/fake/model.pth" + + +def test_depth_estimation_passes_form_fields_to_command(image_file): + with patch("vizion3d.server.rest.depth_estimation.DepthEstimation") as mock_cls: + mock_cls.return_value.run.return_value = _fake_result() + client.post( + "/lifting/depth-estimation", + files={"image": ("test.png", image_file, "image/png")}, + data={"model_backend": "/my/model.pth", "return_depth_image": "false"}, + ) + called_cmd = mock_cls.return_value.run.call_args[0][0] + assert called_cmd.model_backend == "/my/model.pth" + assert called_cmd.return_depth_image is False + + +def test_depth_estimation_uses_default_backend_when_omitted(image_file): + with patch("vizion3d.server.rest.depth_estimation.DepthEstimation") as mock_cls: + mock_cls.return_value.run.return_value = _fake_result() + client.post( + "/lifting/depth-estimation", + files={"image": ("test.png", image_file, "image/png")}, + ) + called_cmd = mock_cls.return_value.run.call_args[0][0] + assert called_cmd.model_backend == DEFAULT_DEPTH_MODEL_URL + + +def test_depth_estimation_missing_image_returns_422(): + response = client.post("/lifting/depth-estimation") + assert response.status_code == 422 + + +def test_stereo_depth_returns_200(stereo_files): + left, right = stereo_files + with patch("vizion3d.server.rest.stereo_depth.StereoDepth") as mock_cls: + mock_cls.return_value.run.return_value = _fake_stereo_result() + response = client.post( + "/lifting/stereo-depth", + files={ + "left_image": ("left.png", left, "image/png"), + "right_image": ("right.png", right, "image/png"), + }, + ) + assert response.status_code == 200 + + +def test_stereo_depth_response_has_expected_keys(stereo_files): + left, right = stereo_files + with patch("vizion3d.server.rest.stereo_depth.StereoDepth") as mock_cls: + mock_cls.return_value.run.return_value = _fake_stereo_result() + response = client.post( + "/lifting/stereo-depth", + files={ + "left_image": ("left.png", left, "image/png"), + "right_image": ("right.png", right, "image/png"), + }, + ) + data = response.json() + assert "depth_map" in data + assert "disparity_map" in data + assert "min_depth" in data + assert "max_depth" in data + assert "backend_used" in data + + +def test_stereo_depth_uses_default_backend_when_omitted(stereo_files): + left, right = stereo_files + with patch("vizion3d.server.rest.stereo_depth.StereoDepth") as mock_cls: + mock_cls.return_value.run.return_value = _fake_stereo_result() + client.post( + "/lifting/stereo-depth", + files={ + "left_image": ("left.png", left, "image/png"), + "right_image": ("right.png", right, "image/png"), + }, + ) + called_cmd = mock_cls.return_value.run.call_args[0][0] + assert called_cmd.model_backend == DEFAULT_STEREO_MODEL_URL + + +def test_stereo_depth_passes_form_fields_to_command(stereo_files): + left, right = stereo_files + with patch("vizion3d.server.rest.stereo_depth.StereoDepth") as mock_cls: + mock_cls.return_value.run.return_value = _fake_stereo_result() + client.post( + "/lifting/stereo-depth", + files={ + "left_image": ("left.png", left, "image/png"), + "right_image": ("right.png", right, "image/png"), + }, + data={ + "model_backend": "/my/stereo.pth", + "return_depth_image": "true", + "focal_length": "1733.74", + "baseline": "536.62", + "scale_factor": "0.5", + }, + ) + called_cmd = mock_cls.return_value.run.call_args[0][0] + assert called_cmd.model_backend == "/my/stereo.pth" + assert called_cmd.return_depth_image is True + assert called_cmd.advanced_config.focal_length == pytest.approx(1733.74) + assert called_cmd.advanced_config.baseline == pytest.approx(536.62) + assert called_cmd.advanced_config.scale_factor == pytest.approx(0.5) + + +def test_create_app_depth_only_disables_stereo_endpoint(stereo_files, image_file): + depth_only = TestClient(create_app(enable_depth_estimation=True, enable_stereo_depth=False)) + response = depth_only.post( + "/lifting/stereo-depth", + files={ + "left_image": ("left.png", stereo_files[0], "image/png"), + "right_image": ("right.png", stereo_files[1], "image/png"), + }, + ) + assert response.status_code == 404 + + with patch("vizion3d.server.rest.depth_estimation.DepthEstimation") as mock_cls: + mock_cls.return_value.run.return_value = _fake_result() + response = depth_only.post( + "/lifting/depth-estimation", + files={"image": ("test.png", image_file, "image/png")}, + ) + assert response.status_code == 200 + + +def test_create_app_stereo_only_disables_depth_endpoint(stereo_files, image_file): + stereo_only = TestClient(create_app(enable_depth_estimation=False, enable_stereo_depth=True)) + response = stereo_only.post( + "/lifting/depth-estimation", + files={"image": ("test.png", image_file, "image/png")}, + ) + assert response.status_code == 404 + + left, right = stereo_files + with patch("vizion3d.server.rest.stereo_depth.StereoDepth") as mock_cls: + mock_cls.return_value.run.return_value = _fake_stereo_result() + response = stereo_only.post( + "/lifting/stereo-depth", + files={ + "left_image": ("left.png", left, "image/png"), + "right_image": ("right.png", right, "image/png"), + }, + ) + assert response.status_code == 200 + + +def test_run_preloads_enabled_feature_model_paths(): + with ( + patch("vizion3d.server.rest.app.depth_estimation.configure_model") as depth_cfg, + patch("vizion3d.server.rest.app.stereo_depth.configure_model") as stereo_cfg, + patch("vizion3d.server.rest.app.uvicorn.run") as uvicorn_run, + ): + run( + [ + "--host", + "127.0.0.1", + "--port", + "9000", + "--depth_model", + "/models/depth.pth", + "--stereo_model", + "/models/stereo.pth", + ] + ) + + depth_cfg.assert_called_once_with("/models/depth.pth") + stereo_cfg.assert_called_once_with("/models/stereo.pth") + uvicorn_run.assert_called_once() + assert uvicorn_run.call_args.kwargs["host"] == "127.0.0.1" + assert uvicorn_run.call_args.kwargs["port"] == 9000 + assert "/lifting/depth-estimation" in _uvicorn_app_routes(uvicorn_run) + assert "/lifting/stereo-depth" in _uvicorn_app_routes(uvicorn_run) + + +def test_run_model_path_enables_and_preloads_feature(): + with ( + patch("vizion3d.server.rest.app.depth_estimation.configure_model") as depth_cfg, + patch("vizion3d.server.rest.app.stereo_depth.configure_model") as stereo_cfg, + patch("vizion3d.server.rest.app.uvicorn.run") as uvicorn_run, + ): + run( + [ + "--depth_model", + "/models/depth.pth", + ] + ) + + depth_cfg.assert_called_once_with("/models/depth.pth") + stereo_cfg.assert_not_called() + assert "/lifting/depth-estimation" in _uvicorn_app_routes(uvicorn_run) + assert "/lifting/stereo-depth" not in _uvicorn_app_routes(uvicorn_run) + + +def test_run_only_selected_feature_does_not_enable_other_endpoint(): + with ( + patch("vizion3d.server.rest.app.depth_estimation.configure_model") as depth_cfg, + patch("vizion3d.server.rest.app.stereo_depth.configure_model") as stereo_cfg, + patch("vizion3d.server.rest.app.uvicorn.run") as uvicorn_run, + ): + run(["--stereo_depth", "--stereo_model", "/models/stereo.pth"]) + + depth_cfg.assert_not_called() + stereo_cfg.assert_called_once_with("/models/stereo.pth") + assert "/lifting/depth-estimation" not in _uvicorn_app_routes(uvicorn_run) + assert "/lifting/stereo-depth" in _uvicorn_app_routes(uvicorn_run) diff --git a/tests/unit/test_rest_api.py b/tests/unit/test_rest_api.py deleted file mode 100644 index 7d480b4..0000000 --- a/tests/unit/test_rest_api.py +++ /dev/null @@ -1,122 +0,0 @@ -import io -from unittest.mock import MagicMock, patch - -import pytest -from fastapi.testclient import TestClient -from PIL import Image - -pytest.importorskip("open3d", reason="open3d required — run: uv python pin 3.12 && uv sync") - -from vizion3d.server.rest.app import app # noqa: E402 - -client = TestClient(app) - - -@pytest.fixture -def image_file(): - img = Image.new("RGB", (50, 50), color="green") - buf = io.BytesIO() - img.save(buf, format="PNG") - buf.seek(0) - return buf - - -def _fake_result(depth_map=None): - result = MagicMock() - result.depth_map = depth_map or [[1.0, 2.0], [3.0, 4.0]] - result.min_depth = 1.0 - result.max_depth = 4.0 - result.backend_used = "/fake/model.pth" - result.depth_image = None - result.point_cloud = None - result.mesh = None - return result - - -def test_depth_estimation_returns_200(image_file): - with patch("vizion3d.server.rest.app.DepthEstimation") as mock_cls: - mock_cls.return_value.run.return_value = _fake_result() - response = client.post( - "/lifting/depth-estimation", - files={"image": ("test.png", image_file, "image/png")}, - ) - assert response.status_code == 200 - - -def test_depth_estimation_response_has_expected_keys(image_file): - with patch("vizion3d.server.rest.app.DepthEstimation") as mock_cls: - mock_cls.return_value.run.return_value = _fake_result() - response = client.post( - "/lifting/depth-estimation", - files={"image": ("test.png", image_file, "image/png")}, - ) - data = response.json() - assert "depth_map" in data - assert "min_depth" in data - assert "max_depth" in data - assert "backend_used" in data - - -def test_depth_estimation_depth_map_is_nested_list(image_file): - with patch("vizion3d.server.rest.app.DepthEstimation") as mock_cls: - mock_cls.return_value.run.return_value = _fake_result() - response = client.post( - "/lifting/depth-estimation", - files={"image": ("test.png", image_file, "image/png")}, - ) - data = response.json() - assert isinstance(data["depth_map"], list) - assert all(isinstance(row, list) for row in data["depth_map"]) - - -def test_depth_estimation_max_depth_gte_min_depth(image_file): - with patch("vizion3d.server.rest.app.DepthEstimation") as mock_cls: - mock_cls.return_value.run.return_value = _fake_result() - response = client.post( - "/lifting/depth-estimation", - files={"image": ("test.png", image_file, "image/png")}, - ) - data = response.json() - assert data["max_depth"] >= data["min_depth"] - - -def test_depth_estimation_optional_outputs_null_by_default(image_file): - with patch("vizion3d.server.rest.app.DepthEstimation") as mock_cls: - mock_cls.return_value.run.return_value = _fake_result() - response = client.post( - "/lifting/depth-estimation", - files={"image": ("test.png", image_file, "image/png")}, - ) - data = response.json() - assert data["depth_image"] is None - assert data["point_cloud_ply"] is None - assert data["mesh_ply"] is None - - -def test_depth_estimation_backend_used_is_returned(image_file): - with patch("vizion3d.server.rest.app.DepthEstimation") as mock_cls: - mock_cls.return_value.run.return_value = _fake_result() - response = client.post( - "/lifting/depth-estimation", - files={"image": ("test.png", image_file, "image/png")}, - ) - data = response.json() - assert data["backend_used"] == "/fake/model.pth" - - -def test_depth_estimation_passes_form_fields_to_command(image_file): - with patch("vizion3d.server.rest.app.DepthEstimation") as mock_cls: - mock_cls.return_value.run.return_value = _fake_result() - client.post( - "/lifting/depth-estimation", - files={"image": ("test.png", image_file, "image/png")}, - data={"model_backend": "/my/model.pth", "return_depth_image": "false"}, - ) - called_cmd = mock_cls.return_value.run.call_args[0][0] - assert called_cmd.model_backend == "/my/model.pth" - assert called_cmd.return_depth_image is False - - -def test_depth_estimation_missing_image_returns_422(): - response = client.post("/lifting/depth-estimation") - assert response.status_code == 422 diff --git a/tests/unit/test_stereo_depth_config.py b/tests/unit/test_stereo_depth_config.py new file mode 100644 index 0000000..191b311 --- /dev/null +++ b/tests/unit/test_stereo_depth_config.py @@ -0,0 +1,317 @@ +""" +Unit tests for StereoDepthAdvancedConfig and its propagation through +the command, handler, REST API, and gRPC server. +""" + +import io +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +from PIL import Image + +from vizion3d.stereo.commands import StereoDepthCommand +from vizion3d.stereo.models import StereoDepthAdvancedConfig + +# ── Fixtures ────────────────────────────────────────────────────────────────── + + +@pytest.fixture +def dummy_stereo_bytes(): + """Return a pair of identical small PNG images as bytes.""" + img = Image.new("RGB", (64, 48), color=(100, 150, 200)) + buf = io.BytesIO() + img.save(buf, format="PNG") + data = buf.getvalue() + return data, data # (left_bytes, right_bytes) + + +@pytest.fixture +def fake_disp(): + """Small synthetic disparity map.""" + rng = np.random.default_rng(42) + return rng.uniform(0.5, 50.0, (48, 64)).astype(np.float32) + + +# ── StereoDepthAdvancedConfig ───────────────────────────────────────────────── + + +class TestStereoDepthAdvancedConfig: + def test_defaults(self): + cfg = StereoDepthAdvancedConfig() + assert cfg.focal_length == 1000.0 + assert cfg.cx == 640.0 + assert cfg.cy == 360.0 + assert cfg.baseline == 100.0 + assert cfg.doffs == 0.0 + assert cfg.z_far == 10.0 + assert cfg.conf_threshold == 0.1 + assert cfg.occ_threshold == 0.5 + assert cfg.scale_factor == 1.0 + + def test_partial_override_keeps_other_defaults(self): + cfg = StereoDepthAdvancedConfig(focal_length=1733.74, baseline=536.62) + assert cfg.focal_length == pytest.approx(1733.74) + assert cfg.baseline == pytest.approx(536.62) + assert cfg.cx == 640.0 # unchanged + assert cfg.cy == 360.0 # unchanged + + def test_all_fields_overridable(self): + cfg = StereoDepthAdvancedConfig( + focal_length=800.0, + cx=320.0, + cy=240.0, + baseline=200.0, + doffs=5.0, + z_far=15.0, + conf_threshold=0.3, + occ_threshold=0.7, + scale_factor=0.5, + ) + assert cfg.focal_length == 800.0 + assert cfg.cx == 320.0 + assert cfg.cy == 240.0 + assert cfg.baseline == 200.0 + assert cfg.doffs == 5.0 + assert cfg.z_far == 15.0 + assert cfg.conf_threshold == 0.3 + assert cfg.occ_threshold == 0.7 + assert cfg.scale_factor == 0.5 + + def test_invalid_type_rejected(self): + with pytest.raises(Exception): + StereoDepthAdvancedConfig(focal_length="not_a_float") + + +# ── StereoDepthCommand advanced_config field ────────────────────────────────── + + +class TestStereoDepthCommandAdvancedConfig: + def test_default_advanced_config_is_correct_type(self, dummy_stereo_bytes): + left, right = dummy_stereo_bytes + cmd = StereoDepthCommand(left_image=left, right_image=right) + assert isinstance(cmd.advanced_config, StereoDepthAdvancedConfig) + + def test_default_advanced_config_has_expected_values(self, dummy_stereo_bytes): + left, right = dummy_stereo_bytes + cmd = StereoDepthCommand(left_image=left, right_image=right) + assert cmd.advanced_config.focal_length == 1000.0 + assert cmd.advanced_config.baseline == 100.0 + + def test_each_command_gets_separate_config_instance(self, dummy_stereo_bytes): + left, right = dummy_stereo_bytes + cmd1 = StereoDepthCommand(left_image=left, right_image=right) + cmd2 = StereoDepthCommand(left_image=left, right_image=right) + assert cmd1.advanced_config is not cmd2.advanced_config + + def test_mutation_of_one_command_does_not_affect_another(self, dummy_stereo_bytes): + left, right = dummy_stereo_bytes + cmd1 = StereoDepthCommand(left_image=left, right_image=right) + cmd2 = StereoDepthCommand(left_image=left, right_image=right) + cmd1.advanced_config.focal_length = 999.0 + assert cmd2.advanced_config.focal_length == 1000.0 + + def test_custom_config_stored_on_command(self, dummy_stereo_bytes): + left, right = dummy_stereo_bytes + cfg = StereoDepthAdvancedConfig(focal_length=1733.74, baseline=536.62) + cmd = StereoDepthCommand(left_image=left, right_image=right, advanced_config=cfg) + assert cmd.advanced_config.focal_length == pytest.approx(1733.74) + assert cmd.advanced_config.baseline == pytest.approx(536.62) + + +# ── Depth formula correctness ───────────────────────────────────────────────── + + +class TestDepthFormula: + """Validate the disparity→depth formula: depth_m = baseline * focal / disp / 1000.""" + + def test_depth_from_known_disparity(self): + # baseline=100mm, focal=1000px, disp=10px → depth = 100*1000/10/1000 = 10m + baseline = 100.0 + focal = 1000.0 + disp = 10.0 + expected_depth = baseline * focal / disp / 1000.0 + assert expected_depth == pytest.approx(10.0) + + def test_zero_disparity_gives_zero_depth(self): + # Handler sets depth_mm[disp<=0]=0.0, so depth_m=0 + disp = np.array([0.0, 1.0, 2.0]) + baseline = 100.0 + focal = 1000.0 + with np.errstate(divide="ignore", invalid="ignore"): + depth_mm = baseline * focal / (disp + 0.0) + depth_mm[disp <= 0] = 0.0 + depth_m = depth_mm / 1000.0 + assert depth_m[0] == 0.0 + + def test_doffs_shifts_depth(self): + # doffs shifts the effective disparity: depth = baseline * focal / (disp + doffs) + baseline = 100.0 + focal = 1000.0 + disp = 10.0 + doffs = 5.0 + depth = baseline * focal / (disp + doffs) / 1000.0 + expected = 100.0 * 1000.0 / 15.0 / 1000.0 + assert depth == pytest.approx(expected) + + +# ── REST API propagates config form fields ──────────────────────────────────── + + +class TestStereoRestAdvancedConfig: + @pytest.fixture(autouse=True) + def _setup(self): + pytest.importorskip("open3d", reason="open3d required — run: uv python pin 3.12 && uv sync") + from fastapi.testclient import TestClient + + from vizion3d.server.rest.app import app + + self.client = TestClient(app) + + def _post(self, extra_data=None): + img = Image.new("RGB", (32, 32), color="green") + buf = io.BytesIO() + img.save(buf, format="PNG") + img_bytes = buf.getvalue() + + result = MagicMock() + result.depth_map = [[1.0]] + result.disparity_map = [[5.0]] + result.min_depth = 1.0 + result.max_depth = 1.0 + result.backend_used = "/fake/stereo.pth" + result.depth_image = None + result.point_cloud = None + result.mesh = None + + with patch("vizion3d.server.rest.stereo_depth.StereoDepth") as mock_cls: + mock_cls.return_value.run.return_value = result + self.client.post( + "/lifting/stereo-depth", + files={ + "left_image": ("left.png", io.BytesIO(img_bytes), "image/png"), + "right_image": ("right.png", io.BytesIO(img_bytes), "image/png"), + }, + data=extra_data or {}, + ) + return mock_cls.return_value.run.call_args[0][0] + + def test_default_config_used_when_no_form_fields_sent(self): + cmd = self._post() + assert cmd.advanced_config.focal_length == 1000.0 + assert cmd.advanced_config.baseline == 100.0 + assert cmd.advanced_config.cx == 640.0 + assert cmd.advanced_config.cy == 360.0 + assert cmd.advanced_config.z_far == 10.0 + + def test_custom_focal_length_forwarded(self): + cmd = self._post({"focal_length": "1733.74"}) + assert cmd.advanced_config.focal_length == pytest.approx(1733.74) + assert cmd.advanced_config.baseline == 100.0 # unchanged + + def test_custom_baseline_forwarded(self): + cmd = self._post({"baseline": "536.62"}) + assert cmd.advanced_config.baseline == pytest.approx(536.62) + + def test_custom_cx_cy_forwarded(self): + cmd = self._post({"cx": "792.27", "cy": "541.89"}) + assert cmd.advanced_config.cx == pytest.approx(792.27) + assert cmd.advanced_config.cy == pytest.approx(541.89) + + def test_custom_z_far_forwarded(self): + cmd = self._post({"z_far": "5.0"}) + assert cmd.advanced_config.z_far == pytest.approx(5.0) + + def test_custom_scale_factor_forwarded(self): + cmd = self._post({"scale_factor": "0.5"}) + assert cmd.advanced_config.scale_factor == pytest.approx(0.5) + + def test_partial_override_does_not_affect_other_fields(self): + cmd = self._post({"focal_length": "800.0"}) + assert cmd.advanced_config.focal_length == pytest.approx(800.0) + assert cmd.advanced_config.cy == 360.0 # unchanged + + +# ── gRPC server propagates config proto fields ──────────────────────────────── + + +class TestStereoGrpcAdvancedConfig: + @pytest.fixture(autouse=True) + def _setup(self): + pytest.importorskip("open3d", reason="open3d required — run: uv python pin 3.12 && uv sync") + from vizion3d.proto import lifting_pb2 + from vizion3d.server.grpc.server import LiftingServiceServicer + + self.pb2 = lifting_pb2 + self.servicer = LiftingServiceServicer() + self.context = MagicMock() + + img = Image.new("RGB", (32, 32), color=(50, 100, 150)) + buf = io.BytesIO() + img.save(buf, format="PNG") + self.image_bytes = buf.getvalue() + + def _run(self, proto_cfg=None): + kwargs = { + "left_image_bytes": self.image_bytes, + "right_image_bytes": self.image_bytes, + } + if proto_cfg is not None: + kwargs["advanced_config"] = proto_cfg + request = self.pb2.StereoDepthRequest(**kwargs) + + result = MagicMock() + result.depth_map = [[1.0]] + result.disparity_map = [[5.0]] + result.min_depth = 1.0 + result.max_depth = 1.0 + result.backend_used = "/fake/stereo.pth" + result.depth_image = None + result.point_cloud = None + result.mesh = None + + with patch("vizion3d.server.grpc.server.StereoDepth") as mock_cls: + mock_cls.return_value.run.return_value = result + self.servicer.RunStereoDepth(request, self.context) + return mock_cls.return_value.run.call_args[0][0] + + def test_no_config_in_request_uses_defaults(self): + cmd = self._run() + assert cmd.advanced_config.focal_length == 1000.0 + assert cmd.advanced_config.baseline == 100.0 + + def test_full_config_forwarded(self): + proto_cfg = self.pb2.StereoDepthAdvancedConfig( + focal_length=1733.74, + cx=792.27, + cy=541.89, + baseline=536.62, + doffs=0.0, + z_far=8.0, + conf_threshold=0.2, + occ_threshold=0.6, + scale_factor=0.5, + ) + cmd = self._run(proto_cfg) + assert cmd.advanced_config.focal_length == pytest.approx(1733.74) + assert cmd.advanced_config.cx == pytest.approx(792.27) + assert cmd.advanced_config.cy == pytest.approx(541.89) + assert cmd.advanced_config.baseline == pytest.approx(536.62) + assert cmd.advanced_config.z_far == pytest.approx(8.0) + assert cmd.advanced_config.conf_threshold == pytest.approx(0.2) + assert cmd.advanced_config.occ_threshold == pytest.approx(0.6) + assert cmd.advanced_config.scale_factor == pytest.approx(0.5) + + def test_partial_config_overrides_only_set_fields(self): + proto_cfg = self.pb2.StereoDepthAdvancedConfig(focal_length=800.0, z_far=5.0) + cmd = self._run(proto_cfg) + assert cmd.advanced_config.focal_length == pytest.approx(800.0) + assert cmd.advanced_config.baseline == 100.0 # default + assert cmd.advanced_config.cx == 640.0 # default + assert cmd.advanced_config.z_far == pytest.approx(5.0) + + def test_empty_config_message_uses_all_defaults(self): + proto_cfg = self.pb2.StereoDepthAdvancedConfig() + cmd = self._run(proto_cfg) + assert cmd.advanced_config.focal_length == 1000.0 + assert cmd.advanced_config.baseline == 100.0 diff --git a/tests/unit/test_stereo_depth_handler.py b/tests/unit/test_stereo_depth_handler.py new file mode 100644 index 0000000..f2029cb --- /dev/null +++ b/tests/unit/test_stereo_depth_handler.py @@ -0,0 +1,300 @@ +""" +Unit tests for StereoDepthHandler internals and the StereoDepth facade. + +These tests mock the S2M2 model so no checkpoint file is needed on disk. +""" + +import io +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +from PIL import Image + +from vizion3d.stereo.commands import StereoDepthCommand +from vizion3d.stereo.handlers import StereoDepthHandler +from vizion3d.stereo.models import StereoDepthAdvancedConfig + +# ── Fixtures ────────────────────────────────────────────────────────────────── + + +@pytest.fixture +def dummy_image_bytes(): + img = Image.new("RGB", (64, 48), color=(80, 120, 160)) + buf = io.BytesIO() + img.save(buf, format="PNG") + return buf.getvalue() + + +@pytest.fixture +def fake_disp(): + """Realistic disparity field: positive, 48×64 float32 array.""" + rng = np.random.default_rng(7) + return rng.uniform(1.0, 50.0, (48, 64)).astype(np.float32) + + +# ── Handler device helper ───────────────────────────────────────────────────── + + +class TestHandlerDeviceHelper: + def test_returns_cuda_when_available(self): + torch_mock = MagicMock() + torch_mock.cuda.is_available.return_value = True + assert StereoDepthHandler._torch_device(torch_mock) == "cuda" + + def test_returns_mps_when_cuda_unavailable(self): + torch_mock = MagicMock() + torch_mock.cuda.is_available.return_value = False + torch_mock.backends.mps.is_available.return_value = True + assert StereoDepthHandler._torch_device(torch_mock) == "mps" + + def test_returns_cpu_as_fallback(self): + torch_mock = MagicMock() + torch_mock.cuda.is_available.return_value = False + torch_mock.backends.mps.is_available.return_value = False + assert StereoDepthHandler._torch_device(torch_mock) == "cpu" + + +# ── Handler output shapes ───────────────────────────────────────────────────── + + +class TestHandlerOutputShapes: + """Verify depth_map, disparity_map shapes and metric depth computation.""" + + def test_depth_map_is_2d_list_of_floats(self, dummy_image_bytes, fake_disp): + with patch.object(StereoDepthHandler, "_run_s2m2", return_value=fake_disp): + cmd = StereoDepthCommand( + left_image=dummy_image_bytes, + right_image=dummy_image_bytes, + model_backend="/fake/model.pth", + ) + result = StereoDepthHandler().handle(cmd) + + assert isinstance(result.depth_map, list) + assert isinstance(result.depth_map[0], list) + assert isinstance(result.depth_map[0][0], float) + + def test_disparity_map_matches_input_shape(self, dummy_image_bytes, fake_disp): + with patch.object(StereoDepthHandler, "_run_s2m2", return_value=fake_disp): + cmd = StereoDepthCommand( + left_image=dummy_image_bytes, + right_image=dummy_image_bytes, + model_backend="/fake/model.pth", + ) + result = StereoDepthHandler().handle(cmd) + + H, W = fake_disp.shape + assert len(result.disparity_map) == H + assert len(result.disparity_map[0]) == W + + def test_min_max_depth_ordering(self, dummy_image_bytes, fake_disp): + with patch.object(StereoDepthHandler, "_run_s2m2", return_value=fake_disp): + result = StereoDepthHandler().handle( + StereoDepthCommand( + left_image=dummy_image_bytes, + right_image=dummy_image_bytes, + model_backend="/fake/model.pth", + ) + ) + assert result.max_depth >= result.min_depth + + def test_point_cloud_scale_is_1(self, dummy_image_bytes, fake_disp): + with patch.object(StereoDepthHandler, "_run_s2m2", return_value=fake_disp): + result = StereoDepthHandler().handle( + StereoDepthCommand( + left_image=dummy_image_bytes, + right_image=dummy_image_bytes, + model_backend="/fake/model.pth", + ) + ) + assert result.point_cloud_scale == 1.0 + + def test_backend_used_propagated(self, dummy_image_bytes, fake_disp): + with patch.object(StereoDepthHandler, "_run_s2m2", return_value=fake_disp): + result = StereoDepthHandler().handle( + StereoDepthCommand( + left_image=dummy_image_bytes, + right_image=dummy_image_bytes, + model_backend="/fake/model.pth", + ) + ) + assert result.backend_used == "/fake/model.pth" + + +# ── Depth formula integration ───────────────────────────────────────────────── + + +class TestDepthMetricConversion: + """Verify the handler converts disparity to real metric depth correctly.""" + + def test_high_disparity_gives_low_depth(self, dummy_image_bytes): + # All pixels at disp=100px: depth = 100 * 1000 / 100 / 1000 = 1m + uniform_disp = np.full((48, 64), 100.0, dtype=np.float32) + with patch.object(StereoDepthHandler, "_run_s2m2", return_value=uniform_disp): + result = StereoDepthHandler().handle( + StereoDepthCommand( + left_image=dummy_image_bytes, + right_image=dummy_image_bytes, + model_backend="/fake/model.pth", + advanced_config=StereoDepthAdvancedConfig(focal_length=1000.0, baseline=100.0), + ) + ) + # depth = baseline * focal / disp / 1000 = 100*1000/100/1000 = 1.0 m + assert result.min_depth == pytest.approx(1.0, abs=1e-3) + assert result.max_depth == pytest.approx(1.0, abs=1e-3) + + def test_zero_disparity_pixels_excluded_from_depth(self, dummy_image_bytes): + disp = np.zeros((48, 64), dtype=np.float32) + disp[0, 0] = 10.0 # one valid pixel + with patch.object(StereoDepthHandler, "_run_s2m2", return_value=disp): + result = StereoDepthHandler().handle( + StereoDepthCommand( + left_image=dummy_image_bytes, + right_image=dummy_image_bytes, + model_backend="/fake/model.pth", + advanced_config=StereoDepthAdvancedConfig(focal_length=1000.0, baseline=100.0), + ) + ) + # Zero disparity → zero depth; non-zero depth present for (0,0) + assert result.min_depth == pytest.approx(0.0) + assert result.max_depth > 0.0 + + +# ── Optional outputs ────────────────────────────────────────────────────────── + + +class TestHandlerOptionalOutputs: + def test_no_optional_outputs_by_default(self, dummy_image_bytes, fake_disp): + with patch.object(StereoDepthHandler, "_run_s2m2", return_value=fake_disp): + result = StereoDepthHandler().handle( + StereoDepthCommand( + left_image=dummy_image_bytes, + right_image=dummy_image_bytes, + model_backend="/fake/model.pth", + ) + ) + assert result.depth_image is None + assert result.point_cloud is None + assert result.mesh is None + + def test_return_depth_image_requires_open3d(self, dummy_image_bytes, fake_disp): + pytest.importorskip("open3d", reason="open3d required") + with patch.object(StereoDepthHandler, "_run_s2m2", return_value=fake_disp): + result = StereoDepthHandler().handle( + StereoDepthCommand( + left_image=dummy_image_bytes, + right_image=dummy_image_bytes, + model_backend="/fake/model.pth", + return_depth_image=True, + ) + ) + assert result.depth_image is not None + arr = np.asarray(result.depth_image) + assert arr.dtype == np.uint16 + + def test_return_point_cloud_requires_open3d(self, dummy_image_bytes, fake_disp): + pytest.importorskip("open3d", reason="open3d required") + with patch.object(StereoDepthHandler, "_run_s2m2", return_value=fake_disp): + result = StereoDepthHandler().handle( + StereoDepthCommand( + left_image=dummy_image_bytes, + right_image=dummy_image_bytes, + model_backend="/fake/model.pth", + return_point_cloud=True, + advanced_config=StereoDepthAdvancedConfig( + focal_length=1000.0, baseline=100.0, z_far=10.0 + ), + ) + ) + assert result.point_cloud is not None + assert result.point_cloud.has_points() + assert result.point_cloud.has_colors() + + def test_conf_occ_thresholds_filter_point_cloud(self, dummy_image_bytes, fake_disp): + pytest.importorskip("open3d", reason="open3d required") + occ = np.ones_like(fake_disp, dtype=np.float32) + conf = np.ones_like(fake_disp, dtype=np.float32) + conf[:, :32] = 0.0 + + with patch.object(StereoDepthHandler, "_run_s2m2", return_value=(fake_disp, occ, conf)): + result = StereoDepthHandler().handle( + StereoDepthCommand( + left_image=dummy_image_bytes, + right_image=dummy_image_bytes, + model_backend="/fake/model.pth", + return_point_cloud=True, + advanced_config=StereoDepthAdvancedConfig( + focal_length=1000.0, + baseline=100.0, + z_far=100.0, + conf_threshold=0.5, + occ_threshold=0.5, + ), + ) + ) + + assert result.point_cloud is not None + assert len(np.asarray(result.point_cloud.points)) < fake_disp.size + + def test_return_mesh_requires_open3d(self, dummy_image_bytes, fake_disp): + open3d = pytest.importorskip("open3d", reason="open3d required") + with patch.object(StereoDepthHandler, "_run_s2m2", return_value=fake_disp): + result = StereoDepthHandler().handle( + StereoDepthCommand( + left_image=dummy_image_bytes, + right_image=dummy_image_bytes, + model_backend="/fake/model.pth", + return_mesh=True, + advanced_config=StereoDepthAdvancedConfig( + focal_length=1000.0, baseline=100.0, z_far=10.0 + ), + ) + ) + assert result.mesh is not None + assert isinstance(result.mesh, open3d.geometry.TriangleMesh) + + +# ── S2M2 variant detection ──────────────────────────────────────────────────── + + +class TestS2M2VariantDetection: + def test_l_variant_detected(self): + from vizion3d.stereo.arch.s2m2 import s2m2_config_from_checkpoint + + cfg = s2m2_config_from_checkpoint("/models/stereo-depth-s2m2-L.pth") + assert cfg["feature_channels"] == 256 + assert cfg["num_transformer"] == 3 + + def test_s_variant_detected(self): + from vizion3d.stereo.arch.s2m2 import s2m2_config_from_checkpoint + + cfg = s2m2_config_from_checkpoint("/models/stereo-depth-s2m2-S.pth") + assert cfg["feature_channels"] == 128 + assert cfg["num_transformer"] == 1 + + def test_m_variant_detected(self): + from vizion3d.stereo.arch.s2m2 import s2m2_config_from_checkpoint + + cfg = s2m2_config_from_checkpoint("/models/stereo-depth-s2m2-M.pth") + assert cfg["feature_channels"] == 192 + assert cfg["num_transformer"] == 2 + + def test_xl_variant_detected(self): + from vizion3d.stereo.arch.s2m2 import s2m2_config_from_checkpoint + + cfg = s2m2_config_from_checkpoint("/models/stereo-depth-s2m2-XL.pth") + assert cfg["feature_channels"] == 384 + assert cfg["num_transformer"] == 3 + + def test_unknown_variant_falls_back_to_l(self): + from vizion3d.stereo.arch.s2m2 import s2m2_config_from_checkpoint + + cfg = s2m2_config_from_checkpoint("/models/my-custom-stereo-model.pth") + assert cfg["feature_channels"] == 256 # L default + assert cfg["num_transformer"] == 3 + + def test_xl_not_confused_with_l(self): + from vizion3d.stereo.arch.s2m2 import s2m2_config_from_checkpoint + + cfg = s2m2_config_from_checkpoint("/models/model-XL.pth") + assert cfg["feature_channels"] == 384 # XL, not L diff --git a/vizion3d/lifting/__init__.py b/vizion3d/lifting/__init__.py index 3da1674..7f8f2da 100644 --- a/vizion3d/lifting/__init__.py +++ b/vizion3d/lifting/__init__.py @@ -2,7 +2,7 @@ from .commands import DepthEstimationCommand from .handlers import DepthEstimationHandler -from .models import DepthEstimationResult +from .models import DepthEstimationAdvanceConfig, DepthEstimationResult # Register handlers on import or application startup register_command_handler(DepthEstimationCommand, DepthEstimationHandler) @@ -17,9 +17,21 @@ class DepthEstimation: Example: ```python - cmd = DepthEstimationCommand(image_input=b"...", return_mesh=True) - task = DepthEstimation() - result = task.run(cmd) + from vizion3d.lifting import ( + DepthEstimation, + DepthEstimationAdvanceConfig, + DepthEstimationCommand, + ) + + cmd = DepthEstimationCommand( + image_input=b"...", + return_point_cloud=True, + return_mesh=True, + advanced_config=DepthEstimationAdvanceConfig( + fx=615.0, fy=615.0, cx=320.0, cy=240.0, depth_trunc=5.0 + ), + ) + result = DepthEstimation().run(cmd) ``` """ @@ -38,4 +50,9 @@ def run(self, command: DepthEstimationCommand) -> DepthEstimationResult: return command_bus.dispatch(command) -__all__ = ["DepthEstimation", "DepthEstimationCommand", "DepthEstimationResult"] +__all__ = [ + "DepthEstimation", + "DepthEstimationAdvanceConfig", + "DepthEstimationCommand", + "DepthEstimationResult", +] diff --git a/vizion3d/lifting/commands.py b/vizion3d/lifting/commands.py index 6528b67..6abe366 100644 --- a/vizion3d/lifting/commands.py +++ b/vizion3d/lifting/commands.py @@ -1,9 +1,9 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field from vizion3d.core.cqrs import Command from .defaults import DEFAULT_DEPTH_MODEL_URL -from .models import DepthEstimationResult +from .models import DepthEstimationAdvanceConfig, DepthEstimationResult @dataclass @@ -30,11 +30,15 @@ class DepthEstimationCommand(Command[DepthEstimationResult]): to the full 0–65535 range. Requires Open3D (Python 3.12). return_point_cloud: When `True`, the result includes an `open3d.geometry.PointCloud` unprojected from the RGB-D image using - PrimeSense default camera intrinsics. Point coordinates are in metres. + the camera intrinsics in `advanced_config`. Point coordinates are in metres. Requires Open3D (Python 3.12). return_mesh: When `True`, the result includes an `open3d.geometry.TriangleMesh` reconstructed from the point cloud via ball-pivoting. Includes vertex colours. Requires Open3D (Python 3.12). + advanced_config: Camera intrinsics and depth range settings. Override any + field to customise — e.g. + ``advanced_config=DepthEstimationAdvanceConfig(fx=615.0, fy=615.0)``. + Unspecified fields keep their defaults (PrimeSense values). """ image_input: str | bytes @@ -42,3 +46,6 @@ class DepthEstimationCommand(Command[DepthEstimationResult]): return_depth_image: bool = False return_point_cloud: bool = False return_mesh: bool = False + advanced_config: DepthEstimationAdvanceConfig = field( + default_factory=DepthEstimationAdvanceConfig + ) diff --git a/vizion3d/lifting/depth_anything.py b/vizion3d/lifting/depth_anything.py index 16bbc3b..5462829 100644 --- a/vizion3d/lifting/depth_anything.py +++ b/vizion3d/lifting/depth_anything.py @@ -94,11 +94,14 @@ def add(target: str, source: str): "pretrained.patch_embed.proj.bias", ) - layer_count = max( - int(key.split(".")[2]) - for key in state_dict - if key.startswith("pretrained.blocks.") and key.endswith(".norm1.weight") - ) + 1 + layer_count = ( + max( + int(key.split(".")[2]) + for key in state_dict + if key.startswith("pretrained.blocks.") and key.endswith(".norm1.weight") + ) + + 1 + ) for idx in range(layer_count): original = f"pretrained.blocks.{idx}" @@ -107,12 +110,10 @@ def add(target: str, source: str): add(f"{target}.norm1.weight", f"{original}.norm1.weight") add(f"{target}.norm1.bias", f"{original}.norm1.bias") - query_weight, key_weight, value_weight = state_dict[ - f"{original}.attn.qkv.weight" - ].chunk(3, dim=0) - query_bias, key_bias, value_bias = state_dict[f"{original}.attn.qkv.bias"].chunk( + query_weight, key_weight, value_weight = state_dict[f"{original}.attn.qkv.weight"].chunk( 3, dim=0 ) + query_bias, key_bias, value_bias = state_dict[f"{original}.attn.qkv.bias"].chunk(3, dim=0) converted[f"{target}.attention.attention.query.weight"] = query_weight converted[f"{target}.attention.attention.key.weight"] = key_weight diff --git a/vizion3d/lifting/handlers.py b/vizion3d/lifting/handlers.py index a4cc7d6..660fbc5 100644 --- a/vizion3d/lifting/handlers.py +++ b/vizion3d/lifting/handlers.py @@ -1,3 +1,4 @@ +import contextlib import io import threading @@ -11,8 +12,10 @@ from .depth_anything import convert_depth_anything_v2_state_dict, depth_anything_v2_config from .models import DepthEstimationResult -RGBD_DEPTH_SCALE = 1000.0 # default: uint16 millimetres (RealSense / Kinect / PrimeSense) -RGBD_DEPTH_TRUNC = 10.0 # default: discard points beyond 10 m +# Legacy module-level constants kept for reference — actual values come from +# DepthEstimationAdvanceConfig on each command. +_DEFAULT_DEPTH_SCALE = 1000.0 +_DEFAULT_DEPTH_TRUNC = 10.0 OPEN3D_CAMERA_TO_IMAGE_VIEW_TRANSFORM = np.array( [ [1.0, 0.0, 0.0, 0.0], @@ -40,7 +43,7 @@ def handle(self, command: DepthEstimationCommand) -> DepthEstimationResult: min_depth = float(np.min(depth_array)) max_depth = float(np.max(depth_array)) - depth_map = depth_array.tolist() + depth_map = depth_array.astype(np.float32).tolist() depth_image = None if command.return_depth_image: @@ -48,8 +51,7 @@ def handle(self, command: DepthEstimationCommand) -> DepthEstimationResult: import open3d as o3d except ImportError: raise ImportError( - "open3d is required for depth image output. " - "Pin to Python 3.12 and run: uv sync" + "open3d is required for depth image output. Pin to Python 3.12 and run: uv sync" ) range_depth = max_depth - min_depth normalized = ( @@ -71,21 +73,22 @@ def handle(self, command: DepthEstimationCommand) -> DepthEstimationResult: "Pin to Python 3.12 and run: uv sync" ) + cfg = command.advanced_config color_o3d = o3d.geometry.Image(np.asarray(image).copy()) depth_o3d = o3d.geometry.Image( - self._depth_array_to_rgbd_depth(depth_array) + self._depth_array_to_rgbd_depth(depth_array, cfg.depth_scale, cfg.depth_trunc) ) rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth( color_o3d, depth_o3d, - depth_scale=RGBD_DEPTH_SCALE, - depth_trunc=RGBD_DEPTH_TRUNC, + depth_scale=cfg.depth_scale, + depth_trunc=cfg.depth_trunc, convert_rgb_to_intensity=False, ) generated_point_cloud = o3d.geometry.PointCloud.create_from_rgbd_image( rgbd_image, o3d.camera.PinholeCameraIntrinsic( - o3d.camera.PinholeCameraIntrinsicParameters.PrimeSenseDefault + image.width, image.height, cfg.fx, cfg.fy, cfg.cx, cfg.cy ), ) self._orient_point_cloud_like_image(generated_point_cloud) @@ -113,7 +116,9 @@ def _orient_point_cloud_like_image(point_cloud): return point_cloud @staticmethod - def _depth_array_to_rgbd_depth(depth_array: np.ndarray) -> np.ndarray: + def _depth_array_to_rgbd_depth( + depth_array: np.ndarray, depth_scale: float, depth_trunc: float + ) -> np.ndarray: min_depth = float(np.nanmin(depth_array)) max_depth = float(np.nanmax(depth_array)) depth_range = max_depth - min_depth @@ -121,7 +126,7 @@ def _depth_array_to_rgbd_depth(depth_array: np.ndarray) -> np.ndarray: return np.zeros_like(depth_array, dtype=np.uint16) normalized = (depth_array - min_depth) / depth_range - scaled_depth = normalized * RGBD_DEPTH_TRUNC * RGBD_DEPTH_SCALE + scaled_depth = normalized * depth_trunc * depth_scale return np.clip(scaled_depth, 0, np.iinfo(np.uint16).max).astype(np.uint16) @staticmethod @@ -148,14 +153,22 @@ def _mesh_from_point_cloud(point_cloud, o3d): o3d.utility.DoubleVector([radius, radius * 2.0]), ) + @classmethod + def preload(cls, model_path: str) -> None: + """Resolve *model_path* (downloading if a URL) and load it into the class-level cache. + + Call this at server startup to ensure the model is in memory before the first request. + """ + from .defaults import resolve_model_backend + + resolved = resolve_model_backend(model_path) + cls()._load_depth_anything_checkpoint(resolved) + @staticmethod def _torch_device(torch_module) -> str: if torch_module.cuda.is_available(): return "cuda" - if ( - hasattr(torch_module.backends, "mps") - and torch_module.backends.mps.is_available() - ): + if hasattr(torch_module.backends, "mps") and torch_module.backends.mps.is_available(): return "mps" return "cpu" @@ -172,8 +185,7 @@ def _load_depth_anything_checkpoint(self, model_path: str): from transformers import DepthAnythingForDepthEstimation, DPTImageProcessor except ImportError as exc: raise ImportError( - "Depth Anything V2 checkpoints require torch and transformers. " - "Run: uv sync" + "Depth Anything V2 checkpoints require torch and transformers. Run: uv sync" ) from exc model = DepthAnythingForDepthEstimation(depth_anything_v2_config(model_path)) @@ -200,16 +212,23 @@ def _load_depth_anything_checkpoint(self, model_path: str): image_std=[0.229, 0.224, 0.225], ) - self._depth_anything_models[model_path] = (model, processor, torch) + self._depth_anything_models[model_path] = (model, processor, torch, device) return self._depth_anything_models[model_path] def _run_depth_anything_checkpoint(self, model_path: str, image: Image.Image) -> np.ndarray: - model, processor, torch = self._load_depth_anything_checkpoint(model_path) - device = next(model.parameters()).device + model, processor, torch, device = self._load_depth_anything_checkpoint(model_path) inputs = processor(images=image, return_tensors="pt") - inputs = {name: value.to(device) for name, value in inputs.items()} + inputs = {name: value.to(device, non_blocking=True) for name, value in inputs.items()} - with torch.inference_mode(): + device_type = device if isinstance(device, str) else device.type + if device_type == "cuda": + autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=True) + elif device_type == "mps": + autocast_ctx = torch.amp.autocast(device_type="mps", dtype=torch.float16, enabled=True) + else: + autocast_ctx = contextlib.nullcontext() + + with torch.inference_mode(), autocast_ctx: outputs = model(**inputs) post_processed = processor.post_process_depth_estimation( @@ -217,5 +236,9 @@ def _run_depth_anything_checkpoint(self, model_path: str, image: Image.Image) -> target_sizes=[(image.height, image.width)], ) depth = post_processed[0]["predicted_depth"] + result = depth.detach().cpu().numpy().astype(np.float32) + + if device_type == "mps": + torch.mps.empty_cache() - return depth.detach().cpu().numpy().astype(float) + return result diff --git a/vizion3d/lifting/models.py b/vizion3d/lifting/models.py index 2990379..e1bd73b 100644 --- a/vizion3d/lifting/models.py +++ b/vizion3d/lifting/models.py @@ -4,6 +4,37 @@ from pydantic import BaseModel, ConfigDict +class DepthEstimationAdvanceConfig(BaseModel): + """ + Camera intrinsics and depth range settings for depth estimation. + + All fields are optional overrides — unspecified fields retain their defaults, + which match the Open3D PrimeSense preset (640×480 RGB-D sensor). + + Attributes: + fx: Horizontal focal length in pixels. Controls the horizontal field of + view: a larger value means a narrower FOV and more perspective compression. + fy: Vertical focal length in pixels. Usually equal to ``fx`` for square + pixels; differs on sensors with non-square pixels. + cx: Principal point x — the pixel column of the optical axis, typically + near the horizontal image centre. + cy: Principal point y — the pixel row of the optical axis, typically near + the vertical image centre. + depth_scale: Divisor applied to raw uint16 depth values to convert them to + metres. ``1000`` means the raw values are in millimetres (the standard + for RealSense, Kinect, and PrimeSense sensors). + depth_trunc: Maximum depth in metres. Points beyond this distance are + discarded from the point cloud and mesh. + """ + + fx: float = 525.0 + fy: float = 525.0 + cx: float = 319.5 + cy: float = 239.5 + depth_scale: float = 1000.0 + depth_trunc: float = 10.0 + + class DepthEstimationResult(BaseModel): """ Result payload returned after a depth estimation inference task. diff --git a/vizion3d/proto/__init__.py b/vizion3d/proto/__init__.py index e69de29..6bffbb8 100644 --- a/vizion3d/proto/__init__.py +++ b/vizion3d/proto/__init__.py @@ -0,0 +1,7 @@ +import os +import sys + +# grpc_tools generates `import lifting_pb2` (bare) inside lifting_pb2_grpc.py. +# Adding this directory to sys.path makes that import resolvable when the +# package is used as `from vizion3d.proto import ...`. +sys.path.insert(0, os.path.dirname(__file__)) diff --git a/vizion3d/proto/lifting.proto b/vizion3d/proto/lifting.proto index 8314751..9848b26 100644 --- a/vizion3d/proto/lifting.proto +++ b/vizion3d/proto/lifting.proto @@ -4,6 +4,18 @@ package vizion3d.lifting; service LiftingService { rpc RunDepthEstimation (DepthEstimationRequest) returns (DepthEstimationResponse); + rpc RunStereoDepth (StereoDepthRequest) returns (StereoDepthResponse); +} + +// ── Depth Estimation ────────────────────────────────────────────────────────── + +message DepthEstimationAdvanceConfig { + optional float fx = 1; + optional float fy = 2; + optional float cx = 3; + optional float cy = 4; + optional float depth_scale = 5; + optional float depth_trunc = 6; } message DepthEstimationRequest { @@ -13,6 +25,7 @@ message DepthEstimationRequest { bool return_point_cloud = 4; bool return_mesh = 5; optional string local_model_path = 6; + optional DepthEstimationAdvanceConfig advanced_config = 7; } message FloatRow { @@ -28,3 +41,38 @@ message DepthEstimationResponse { bytes point_cloud_ply = 6; bytes mesh_ply = 7; } + +// ── Stereo Depth ────────────────────────────────────────────────────────────── + +message StereoDepthAdvancedConfig { + optional float focal_length = 1; + optional float cx = 2; + optional float cy = 3; + optional float baseline = 4; + optional float doffs = 5; + optional float z_far = 6; + optional float conf_threshold = 7; + optional float occ_threshold = 8; + optional float scale_factor = 9; +} + +message StereoDepthRequest { + bytes left_image_bytes = 1; + bytes right_image_bytes = 2; + string model_backend = 3; + bool return_depth_image = 4; + bool return_point_cloud = 5; + bool return_mesh = 6; + optional StereoDepthAdvancedConfig advanced_config = 7; +} + +message StereoDepthResponse { + repeated FloatRow depth_map = 1; + repeated FloatRow disparity_map = 2; + float min_depth = 3; + float max_depth = 4; + string backend_used = 5; + bytes depth_image = 6; + bytes point_cloud_ply = 7; + bytes mesh_ply = 8; +} diff --git a/vizion3d/proto/lifting_pb2.py b/vizion3d/proto/lifting_pb2.py index ee69c74..d43ef44 100644 --- a/vizion3d/proto/lifting_pb2.py +++ b/vizion3d/proto/lifting_pb2.py @@ -24,19 +24,27 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1cvizion3d/proto/lifting.proto\x12\x10vizion3d.lifting\"\xc5\x01\n\x16\x44\x65pthEstimationRequest\x12\x13\n\x0bimage_bytes\x18\x01 \x01(\x0c\x12\x15\n\rmodel_backend\x18\x02 \x01(\t\x12\x1a\n\x12return_depth_image\x18\x03 \x01(\x08\x12\x1a\n\x12return_point_cloud\x18\x04 \x01(\x08\x12\x13\n\x0breturn_mesh\x18\x05 \x01(\x08\x12\x1d\n\x10local_model_path\x18\x06 \x01(\tH\x00\x88\x01\x01\x42\x13\n\x11_local_model_path\"\x1a\n\x08\x46loatRow\x12\x0e\n\x06values\x18\x01 \x03(\x02\"\xc4\x01\n\x17\x44\x65pthEstimationResponse\x12-\n\tdepth_map\x18\x01 \x03(\x0b\x32\x1a.vizion3d.lifting.FloatRow\x12\x11\n\tmin_depth\x18\x02 \x01(\x02\x12\x11\n\tmax_depth\x18\x03 \x01(\x02\x12\x14\n\x0c\x62\x61\x63kend_used\x18\x04 \x01(\t\x12\x13\n\x0b\x64\x65pth_image\x18\x05 \x01(\x0c\x12\x17\n\x0fpoint_cloud_ply\x18\x06 \x01(\x0c\x12\x10\n\x08mesh_ply\x18\x07 \x01(\x0c\x32{\n\x0eLiftingService\x12i\n\x12RunDepthEstimation\x12(.vizion3d.lifting.DepthEstimationRequest\x1a).vizion3d.lifting.DepthEstimationResponseb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1cvizion3d/proto/lifting.proto\x12\x10vizion3d.lifting\"\xd2\x01\n\x1c\x44\x65pthEstimationAdvanceConfig\x12\x0f\n\x02\x66x\x18\x01 \x01(\x02H\x00\x88\x01\x01\x12\x0f\n\x02\x66y\x18\x02 \x01(\x02H\x01\x88\x01\x01\x12\x0f\n\x02\x63x\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x0f\n\x02\x63y\x18\x04 \x01(\x02H\x03\x88\x01\x01\x12\x18\n\x0b\x64\x65pth_scale\x18\x05 \x01(\x02H\x04\x88\x01\x01\x12\x18\n\x0b\x64\x65pth_trunc\x18\x06 \x01(\x02H\x05\x88\x01\x01\x42\x05\n\x03_fxB\x05\n\x03_fyB\x05\n\x03_cxB\x05\n\x03_cyB\x0e\n\x0c_depth_scaleB\x0e\n\x0c_depth_trunc\"\xa7\x02\n\x16\x44\x65pthEstimationRequest\x12\x13\n\x0bimage_bytes\x18\x01 \x01(\x0c\x12\x15\n\rmodel_backend\x18\x02 \x01(\t\x12\x1a\n\x12return_depth_image\x18\x03 \x01(\x08\x12\x1a\n\x12return_point_cloud\x18\x04 \x01(\x08\x12\x13\n\x0breturn_mesh\x18\x05 \x01(\x08\x12\x1d\n\x10local_model_path\x18\x06 \x01(\tH\x00\x88\x01\x01\x12L\n\x0f\x61\x64vanced_config\x18\x07 \x01(\x0b\x32..vizion3d.lifting.DepthEstimationAdvanceConfigH\x01\x88\x01\x01\x42\x13\n\x11_local_model_pathB\x12\n\x10_advanced_config\"\x1a\n\x08\x46loatRow\x12\x0e\n\x06values\x18\x01 \x03(\x02\"\xc4\x01\n\x17\x44\x65pthEstimationResponse\x12-\n\tdepth_map\x18\x01 \x03(\x0b\x32\x1a.vizion3d.lifting.FloatRow\x12\x11\n\tmin_depth\x18\x02 \x01(\x02\x12\x11\n\tmax_depth\x18\x03 \x01(\x02\x12\x14\n\x0c\x62\x61\x63kend_used\x18\x04 \x01(\t\x12\x13\n\x0b\x64\x65pth_image\x18\x05 \x01(\x0c\x12\x17\n\x0fpoint_cloud_ply\x18\x06 \x01(\x0c\x12\x10\n\x08mesh_ply\x18\x07 \x01(\x0c\"\xe1\x02\n\x19StereoDepthAdvancedConfig\x12\x19\n\x0c\x66ocal_length\x18\x01 \x01(\x02H\x00\x88\x01\x01\x12\x0f\n\x02\x63x\x18\x02 \x01(\x02H\x01\x88\x01\x01\x12\x0f\n\x02\x63y\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x15\n\x08\x62\x61seline\x18\x04 \x01(\x02H\x03\x88\x01\x01\x12\x12\n\x05\x64offs\x18\x05 \x01(\x02H\x04\x88\x01\x01\x12\x12\n\x05z_far\x18\x06 \x01(\x02H\x05\x88\x01\x01\x12\x1b\n\x0e\x63onf_threshold\x18\x07 \x01(\x02H\x06\x88\x01\x01\x12\x1a\n\rocc_threshold\x18\x08 \x01(\x02H\x07\x88\x01\x01\x12\x19\n\x0cscale_factor\x18\t \x01(\x02H\x08\x88\x01\x01\x42\x0f\n\r_focal_lengthB\x05\n\x03_cxB\x05\n\x03_cyB\x0b\n\t_baselineB\x08\n\x06_doffsB\x08\n\x06_z_farB\x11\n\x0f_conf_thresholdB\x10\n\x0e_occ_thresholdB\x0f\n\r_scale_factor\"\x8c\x02\n\x12StereoDepthRequest\x12\x18\n\x10left_image_bytes\x18\x01 \x01(\x0c\x12\x19\n\x11right_image_bytes\x18\x02 \x01(\x0c\x12\x15\n\rmodel_backend\x18\x03 \x01(\t\x12\x1a\n\x12return_depth_image\x18\x04 \x01(\x08\x12\x1a\n\x12return_point_cloud\x18\x05 \x01(\x08\x12\x13\n\x0breturn_mesh\x18\x06 \x01(\x08\x12I\n\x0f\x61\x64vanced_config\x18\x07 \x01(\x0b\x32+.vizion3d.lifting.StereoDepthAdvancedConfigH\x00\x88\x01\x01\x42\x12\n\x10_advanced_config\"\xf3\x01\n\x13StereoDepthResponse\x12-\n\tdepth_map\x18\x01 \x03(\x0b\x32\x1a.vizion3d.lifting.FloatRow\x12\x31\n\rdisparity_map\x18\x02 \x03(\x0b\x32\x1a.vizion3d.lifting.FloatRow\x12\x11\n\tmin_depth\x18\x03 \x01(\x02\x12\x11\n\tmax_depth\x18\x04 \x01(\x02\x12\x14\n\x0c\x62\x61\x63kend_used\x18\x05 \x01(\t\x12\x13\n\x0b\x64\x65pth_image\x18\x06 \x01(\x0c\x12\x17\n\x0fpoint_cloud_ply\x18\x07 \x01(\x0c\x12\x10\n\x08mesh_ply\x18\x08 \x01(\x0c\x32\xda\x01\n\x0eLiftingService\x12i\n\x12RunDepthEstimation\x12(.vizion3d.lifting.DepthEstimationRequest\x1a).vizion3d.lifting.DepthEstimationResponse\x12]\n\x0eRunStereoDepth\x12$.vizion3d.lifting.StereoDepthRequest\x1a%.vizion3d.lifting.StereoDepthResponseb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'vizion3d.proto.lifting_pb2', _globals) if not _descriptor._USE_C_DESCRIPTORS: DESCRIPTOR._loaded_options = None - _globals['_DEPTHESTIMATIONREQUEST']._serialized_start=51 - _globals['_DEPTHESTIMATIONREQUEST']._serialized_end=248 - _globals['_FLOATROW']._serialized_start=250 - _globals['_FLOATROW']._serialized_end=276 - _globals['_DEPTHESTIMATIONRESPONSE']._serialized_start=279 - _globals['_DEPTHESTIMATIONRESPONSE']._serialized_end=475 - _globals['_LIFTINGSERVICE']._serialized_start=477 - _globals['_LIFTINGSERVICE']._serialized_end=600 + _globals['_DEPTHESTIMATIONADVANCECONFIG']._serialized_start=51 + _globals['_DEPTHESTIMATIONADVANCECONFIG']._serialized_end=261 + _globals['_DEPTHESTIMATIONREQUEST']._serialized_start=264 + _globals['_DEPTHESTIMATIONREQUEST']._serialized_end=559 + _globals['_FLOATROW']._serialized_start=561 + _globals['_FLOATROW']._serialized_end=587 + _globals['_DEPTHESTIMATIONRESPONSE']._serialized_start=590 + _globals['_DEPTHESTIMATIONRESPONSE']._serialized_end=786 + _globals['_STEREODEPTHADVANCEDCONFIG']._serialized_start=789 + _globals['_STEREODEPTHADVANCEDCONFIG']._serialized_end=1142 + _globals['_STEREODEPTHREQUEST']._serialized_start=1145 + _globals['_STEREODEPTHREQUEST']._serialized_end=1413 + _globals['_STEREODEPTHRESPONSE']._serialized_start=1416 + _globals['_STEREODEPTHRESPONSE']._serialized_end=1659 + _globals['_LIFTINGSERVICE']._serialized_start=1662 + _globals['_LIFTINGSERVICE']._serialized_end=1880 # @@protoc_insertion_point(module_scope) diff --git a/vizion3d/proto/lifting_pb2_grpc.py b/vizion3d/proto/lifting_pb2_grpc.py index b6b6e1d..bf0be99 100644 --- a/vizion3d/proto/lifting_pb2_grpc.py +++ b/vizion3d/proto/lifting_pb2_grpc.py @@ -39,6 +39,11 @@ def __init__(self, channel): request_serializer=vizion3d_dot_proto_dot_lifting__pb2.DepthEstimationRequest.SerializeToString, response_deserializer=vizion3d_dot_proto_dot_lifting__pb2.DepthEstimationResponse.FromString, _registered_method=True) + self.RunStereoDepth = channel.unary_unary( + '/vizion3d.lifting.LiftingService/RunStereoDepth', + request_serializer=vizion3d_dot_proto_dot_lifting__pb2.StereoDepthRequest.SerializeToString, + response_deserializer=vizion3d_dot_proto_dot_lifting__pb2.StereoDepthResponse.FromString, + _registered_method=True) class LiftingServiceServicer(object): @@ -50,6 +55,12 @@ def RunDepthEstimation(self, request, context): context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') + def RunStereoDepth(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + def add_LiftingServiceServicer_to_server(servicer, server): rpc_method_handlers = { @@ -58,6 +69,11 @@ def add_LiftingServiceServicer_to_server(servicer, server): request_deserializer=vizion3d_dot_proto_dot_lifting__pb2.DepthEstimationRequest.FromString, response_serializer=vizion3d_dot_proto_dot_lifting__pb2.DepthEstimationResponse.SerializeToString, ), + 'RunStereoDepth': grpc.unary_unary_rpc_method_handler( + servicer.RunStereoDepth, + request_deserializer=vizion3d_dot_proto_dot_lifting__pb2.StereoDepthRequest.FromString, + response_serializer=vizion3d_dot_proto_dot_lifting__pb2.StereoDepthResponse.SerializeToString, + ), } generic_handler = grpc.method_handlers_generic_handler( 'vizion3d.lifting.LiftingService', rpc_method_handlers) @@ -95,3 +111,30 @@ def RunDepthEstimation(request, timeout, metadata, _registered_method=True) + + @staticmethod + def RunStereoDepth(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/vizion3d.lifting.LiftingService/RunStereoDepth', + vizion3d_dot_proto_dot_lifting__pb2.StereoDepthRequest.SerializeToString, + vizion3d_dot_proto_dot_lifting__pb2.StereoDepthResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) diff --git a/vizion3d/server/grpc/server.py b/vizion3d/server/grpc/server.py index ebd091c..36d2d5a 100644 --- a/vizion3d/server/grpc/server.py +++ b/vizion3d/server/grpc/server.py @@ -1,3 +1,17 @@ +""" +gRPC server for the vizion3d Lifting service. + +Exposes two RPC methods: +- ``RunDepthEstimation`` — monocular depth from a single image. +- ``RunStereoDepth`` — metric depth from a rectified stereo image pair. + +Start with:: + + uv run vizion3d-serve-grpc + # or + python -m vizion3d.server.grpc.server +""" + import io import logging from concurrent import futures @@ -8,11 +22,18 @@ from vizion3d.lifting import DepthEstimation, DepthEstimationCommand from vizion3d.lifting.defaults import DEFAULT_DEPTH_MODEL_URL +from vizion3d.lifting.models import DepthEstimationAdvanceConfig from vizion3d.lifting.utils import create_mesh_ply_binary, create_ply_binary from vizion3d.proto import lifting_pb2, lifting_pb2_grpc +from vizion3d.stereo import StereoDepth, StereoDepthCommand +from vizion3d.stereo.defaults import DEFAULT_STEREO_MODEL_URL +from vizion3d.stereo.models import StereoDepthAdvancedConfig + +# ── Shared serialisation helpers ────────────────────────────────────────────── 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") @@ -20,26 +41,67 @@ def _o3d_depth_image_to_png_bytes(o3d_image) -> bytes: def _o3d_point_cloud_to_ply_bytes(pcd) -> bytes: + """Serialise an Open3D PointCloud to binary PLY bytes.""" points = np.asarray(pcd.points).astype(np.float32) colors = (np.asarray(pcd.colors) * 255).astype(np.uint8) return create_ply_binary(points, colors) def _o3d_mesh_to_ply_bytes(mesh) -> bytes: + """Serialise an Open3D TriangleMesh to binary PLY bytes.""" points = np.asarray(mesh.vertices).astype(np.float32) colors = (np.asarray(mesh.vertex_colors) * 255).astype(np.uint8) faces = np.asarray(mesh.triangles).astype(np.int32) return create_mesh_ply_binary(points, colors, faces) +# ── gRPC Servicer ───────────────────────────────────────────────────────────── + + class LiftingServiceServicer(lifting_pb2_grpc.LiftingServiceServicer): + """Implements the LiftingService proto RPC methods.""" + + # ── RunDepthEstimation ──────────────────────────────────────────────────── + def RunDepthEstimation(self, request, context): + """Handle a monocular depth estimation request. + + Unmarshals the proto config, dispatches through the CQRS command bus, and + packs the result back into a proto response message. + + Args: + request: ``DepthEstimationRequest`` proto message. + context: gRPC server context. + + Returns: + ``DepthEstimationResponse`` proto message. + """ + base_cfg = DepthEstimationAdvanceConfig() + if request.HasField("advanced_config"): + proto_cfg = request.advanced_config + base_cfg = DepthEstimationAdvanceConfig( + fx=proto_cfg.fx if proto_cfg.HasField("fx") else base_cfg.fx, + fy=proto_cfg.fy if proto_cfg.HasField("fy") else base_cfg.fy, + cx=proto_cfg.cx if proto_cfg.HasField("cx") else base_cfg.cx, + cy=proto_cfg.cy if proto_cfg.HasField("cy") else base_cfg.cy, + depth_scale=( + proto_cfg.depth_scale + if proto_cfg.HasField("depth_scale") + else base_cfg.depth_scale + ), + depth_trunc=( + proto_cfg.depth_trunc + if proto_cfg.HasField("depth_trunc") + else base_cfg.depth_trunc + ), + ) cmd = DepthEstimationCommand( image_input=request.image_bytes, model_backend=request.model_backend or DEFAULT_DEPTH_MODEL_URL, return_depth_image=request.return_depth_image, return_point_cloud=request.return_point_cloud, return_mesh=request.return_mesh, + advanced_config=base_cfg, ) result = DepthEstimation().run(cmd) @@ -48,31 +110,91 @@ def RunDepthEstimation(self, request, context): max_depth=result.max_depth, backend_used=result.backend_used, ) - for row in result.depth_map: response.depth_map.append(lifting_pb2.FloatRow(values=row)) - if result.depth_image is not None: response.depth_image = _o3d_depth_image_to_png_bytes(result.depth_image) - if result.point_cloud is not None: response.point_cloud_ply = _o3d_point_cloud_to_ply_bytes(result.point_cloud) - if result.mesh is not None: response.mesh_ply = _o3d_mesh_to_ply_bytes(result.mesh) + return response + + # ── RunStereoDepth ──────────────────────────────────────────────────────── + + def RunStereoDepth(self, request, context): + """Handle a stereo depth estimation request. + + Unmarshals the proto config, dispatches through the CQRS command bus, and + packs the result back into a proto response message. + + Args: + request: ``StereoDepthRequest`` proto message. + context: gRPC server context. + + Returns: + ``StereoDepthResponse`` proto message. + """ + base_cfg = StereoDepthAdvancedConfig() + if request.HasField("advanced_config"): + proto_cfg = request.advanced_config + + def _f(field: str, default): + return getattr(proto_cfg, field) if proto_cfg.HasField(field) else default + + base_cfg = StereoDepthAdvancedConfig( + focal_length=_f("focal_length", base_cfg.focal_length), + cx=_f("cx", base_cfg.cx), + cy=_f("cy", base_cfg.cy), + baseline=_f("baseline", base_cfg.baseline), + doffs=_f("doffs", base_cfg.doffs), + z_far=_f("z_far", base_cfg.z_far), + conf_threshold=_f("conf_threshold", base_cfg.conf_threshold), + occ_threshold=_f("occ_threshold", base_cfg.occ_threshold), + scale_factor=_f("scale_factor", base_cfg.scale_factor), + ) + + cmd = StereoDepthCommand( + left_image=request.left_image_bytes, + right_image=request.right_image_bytes, + model_backend=request.model_backend or DEFAULT_STEREO_MODEL_URL, + return_depth_image=request.return_depth_image, + return_point_cloud=request.return_point_cloud, + return_mesh=request.return_mesh, + advanced_config=base_cfg, + ) + result = StereoDepth().run(cmd) + response = lifting_pb2.StereoDepthResponse( + min_depth=result.min_depth, + max_depth=result.max_depth, + backend_used=result.backend_used, + ) + for row in result.depth_map: + response.depth_map.append(lifting_pb2.FloatRow(values=row)) + for row in result.disparity_map: + response.disparity_map.append(lifting_pb2.FloatRow(values=row)) + if result.depth_image is not None: + response.depth_image = _o3d_depth_image_to_png_bytes(result.depth_image) + if result.point_cloud is not None: + response.point_cloud_ply = _o3d_point_cloud_to_ply_bytes(result.point_cloud) + if result.mesh is not None: + response.mesh_ply = _o3d_mesh_to_ply_bytes(result.mesh) return response -_MAX_MSG = 500 * 1024 * 1024 # 500 MB +# ── Server bootstrap ────────────────────────────────────────────────────────── + +_MAX_MSG = 500 * 1024 * 1024 # 500 MB _GRPC_OPTIONS = [ - ("grpc.max_send_message_length", _MAX_MSG), + ("grpc.max_send_message_length", _MAX_MSG), ("grpc.max_receive_message_length", _MAX_MSG), ] def serve(): + """Start the gRPC server on port 50051 and block until terminated.""" server = grpc.server( futures.ThreadPoolExecutor(max_workers=10), options=_GRPC_OPTIONS, diff --git a/vizion3d/server/rest/app.py b/vizion3d/server/rest/app.py index 3291cc7..b20c6b1 100644 --- a/vizion3d/server/rest/app.py +++ b/vizion3d/server/rest/app.py @@ -1,102 +1,142 @@ -import base64 -import io +""" +FastAPI REST server for the vizion3d Lifting service. -import numpy as np -import uvicorn -from fastapi import APIRouter, FastAPI, File, Form, Request, UploadFile -from fastapi.responses import JSONResponse -from PIL import Image +Registers feature routers under the ``/lifting`` prefix: +- ``POST /lifting/depth-estimation`` — monocular depth from a single image. +- ``POST /lifting/stereo-depth`` — metric depth from a rectified stereo pair. -from vizion3d.lifting import DepthEstimation, DepthEstimationCommand -from vizion3d.lifting.defaults import DEFAULT_DEPTH_MODEL_URL -from vizion3d.lifting.utils import create_mesh_ply_binary, create_ply_binary +Start with:: -_MAX_BODY = 500 * 1024 * 1024 # 500 MB + uv run vizion3d-serve-rest -app = FastAPI(title="vizion3d REST API", version="1.0.0") + # Enable only depth estimation: + uv run vizion3d-serve-rest --depth_estimation + # Enable only stereo depth with a pre-loaded local model: + uv run vizion3d-serve-rest --stereo_depth --stereo_model /path/to/stereo-depth-s2m2-L.pth -@app.middleware("http") -async def limit_body_size(request: Request, call_next): - content_length = request.headers.get("content-length") - if content_length and int(content_length) > _MAX_BODY: - return JSONResponse( - {"detail": "Request body exceeds the 500 MB limit."}, - status_code=413, - ) - return await call_next(request) + # Both features, custom model paths (pre-loaded at startup): + uv run vizion3d-serve-rest \\ + --depth_model /path/to/depth_anything_v2_vitb.pth \\ + --stereo_model /path/to/stereo-depth-s2m2-L.pth +""" -lifting_router = APIRouter(prefix="/lifting", tags=["Lifting (2D -> 3D)"]) +import argparse +import uvicorn +from fastapi import APIRouter, FastAPI, Request +from fastapi.responses import JSONResponse -def _o3d_depth_image_to_png_bytes(o3d_image) -> bytes: - arr = np.asarray(o3d_image) - buf = io.BytesIO() - Image.fromarray(arr).save(buf, format="PNG") - return buf.getvalue() - +from vizion3d.server.rest import depth_estimation, stereo_depth + +_MAX_BODY = 500 * 1024 * 1024 # 500 MB + + +def create_app( + *, + enable_depth_estimation: bool = True, + enable_stereo_depth: bool = True, +) -> FastAPI: + """Build and return a FastAPI application with the selected feature routers. + + Args: + enable_depth_estimation: When ``True``, register ``POST /lifting/depth-estimation``. + enable_stereo_depth: When ``True``, register ``POST /lifting/stereo-depth``. + + Returns: + A fully configured :class:`FastAPI` instance ready to pass to ``uvicorn.run``. + """ + _app = FastAPI(title="vizion3d REST API", version="1.0.0") + + @_app.middleware("http") + async def limit_body_size(request: Request, call_next): + """Reject requests whose Content-Length header exceeds 500 MB.""" + content_length = request.headers.get("content-length") + if content_length and int(content_length) > _MAX_BODY: + return JSONResponse( + {"detail": "Request body exceeds the 500 MB limit."}, + status_code=413, + ) + return await call_next(request) + + lifting_router = APIRouter(prefix="/lifting", tags=["Lifting (2D -> 3D)"]) + if enable_depth_estimation: + lifting_router.include_router(depth_estimation.router) + if enable_stereo_depth: + lifting_router.include_router(stereo_depth.router) + _app.include_router(lifting_router) + + return _app + + +# Module-level app with all features enabled — used by tests and direct imports. +app = create_app() + + +def _parse_args(argv=None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + prog="vizion3d-serve-rest", + description="Start the vizion3d REST API server.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +feature/model flags (omit all four to enable all features): + --depth_estimation enable POST /lifting/depth-estimation + --stereo_depth enable POST /lifting/stereo-depth + +model pre-loading (also enables that feature; downloads if a URL is given): + --depth_model PATH use PATH as the default depth-estimation model; + the model is loaded into memory at startup + --stereo_model PATH use PATH as the default stereo-depth model; + the model is loaded into memory at startup +""", + ) + parser.add_argument( + "--host", default="0.0.0.0", metavar="HOST", help="Bind host (default: 0.0.0.0)" + ) + parser.add_argument( + "--port", type=int, default=8000, metavar="PORT", help="Bind port (default: 8000)" + ) + parser.add_argument( + "--depth_model", + default=None, + metavar="PATH", + help="Local path or URL for the depth-estimation checkpoint", + ) + parser.add_argument( + "--stereo_model", + default=None, + metavar="PATH", + help="Local path or URL for the stereo-depth checkpoint", + ) + parser.add_argument( + "--depth_estimation", action="store_true", help="Enable only the depth-estimation endpoint" + ) + parser.add_argument( + "--stereo_depth", action="store_true", help="Enable only the stereo-depth endpoint" + ) + return parser.parse_args(argv) -def _o3d_point_cloud_to_ply_bytes(pcd) -> bytes: - points = np.asarray(pcd.points).astype(np.float32) - colors = (np.asarray(pcd.colors) * 255).astype(np.uint8) - return create_ply_binary(points, colors) +def run(argv=None) -> None: + """Parse CLI arguments, configure models, and start the uvicorn server.""" + args = _parse_args(argv) -def _o3d_mesh_to_ply_bytes(mesh) -> bytes: - points = np.asarray(mesh.vertices).astype(np.float32) - colors = (np.asarray(mesh.vertex_colors) * 255).astype(np.uint8) - faces = np.asarray(mesh.triangles).astype(np.int32) - return create_mesh_ply_binary(points, colors, faces) + any_selector = ( + args.depth_estimation or args.stereo_depth or args.depth_model or args.stereo_model + ) + enable_depth = args.depth_estimation or bool(args.depth_model) or not any_selector + enable_stereo = args.stereo_depth or bool(args.stereo_model) or not any_selector + if args.depth_model and enable_depth: + depth_estimation.configure_model(args.depth_model) + if args.stereo_model and enable_stereo: + stereo_depth.configure_model(args.stereo_model) -@lifting_router.post("/depth-estimation") -async def depth_estimation( - image: UploadFile = File(...), - model_backend: str = Form(DEFAULT_DEPTH_MODEL_URL), - return_depth_image: bool = Form(False), - return_point_cloud: bool = Form(False), - return_mesh: bool = Form(False), -): - image_bytes = await image.read() - cmd = DepthEstimationCommand( - image_input=image_bytes, - model_backend=model_backend, - return_depth_image=return_depth_image, - return_point_cloud=return_point_cloud, - return_mesh=return_mesh, + _app = create_app( + enable_depth_estimation=enable_depth, + enable_stereo_depth=enable_stereo, ) - - result = DepthEstimation().run(cmd) - - def _b64(data: bytes | None) -> str | None: - return base64.b64encode(data).decode() if data is not None else None - - return { - "depth_map": result.depth_map, - "min_depth": result.min_depth, - "max_depth": result.max_depth, - "backend_used": result.backend_used, - "depth_image": _b64( - _o3d_depth_image_to_png_bytes(result.depth_image) - if result.depth_image is not None - else None - ), - "point_cloud_ply": _b64( - _o3d_point_cloud_to_ply_bytes(result.point_cloud) - if result.point_cloud is not None - else None - ), - "mesh_ply": _b64( - _o3d_mesh_to_ply_bytes(result.mesh) if result.mesh is not None else None - ), - } - - -app.include_router(lifting_router) - - -def run(): - uvicorn.run(app, host="0.0.0.0", port=8000) + uvicorn.run(_app, host=args.host, port=args.port) if __name__ == "__main__": diff --git a/vizion3d/server/rest/depth_estimation.py b/vizion3d/server/rest/depth_estimation.py new file mode 100644 index 0000000..d195e7f --- /dev/null +++ b/vizion3d/server/rest/depth_estimation.py @@ -0,0 +1,108 @@ +""" +REST endpoint for the Depth Estimation feature. + +Registers ``POST /lifting/depth-estimation`` on the ``lifting_router`` exported +from this module. Import the router in ``app.py`` and call +``app.include_router(lifting_router)``. +""" + +from fastapi import APIRouter, File, Form, UploadFile + +from vizion3d.lifting import DepthEstimation, DepthEstimationCommand +from vizion3d.lifting.defaults import DEFAULT_DEPTH_MODEL_URL +from vizion3d.lifting.models import DepthEstimationAdvanceConfig + +from .serialisation import ( + b64, + o3d_depth_image_to_png_bytes, + o3d_mesh_to_ply_bytes, + o3d_point_cloud_to_ply_bytes, +) + +router = APIRouter() + +# Set by configure_model() at server startup when --depth_model is passed on the CLI. +_model_override: str | None = None + + +def configure_model(path: str) -> None: + """Set a server-wide default model path and pre-load it into handler memory. + + Called from ``app.run()`` when ``--depth_model`` is supplied on the command line. + After this call the endpoint uses *path* whenever the caller omits + ``model_backend`` from the form data. + """ + global _model_override + _model_override = path + from vizion3d.lifting.handlers import DepthEstimationHandler + + DepthEstimationHandler.preload(path) + + +@router.post("/depth-estimation") +async def depth_estimation( + image: UploadFile = File(...), + model_backend: str | None = Form(None), + return_depth_image: bool = Form(False), + return_point_cloud: bool = Form(False), + return_mesh: bool = Form(False), + fx: float | None = Form(None), + fy: float | None = Form(None), + cx: float | None = Form(None), + cy: float | None = Form(None), + depth_scale: float | None = Form(None), + depth_trunc: float | None = Form(None), +): + """Run monocular depth estimation on a single uploaded image. + + Args: + image: The image file (any PIL-supported format). + model_backend: Checkpoint URL or local path (defaults to the vizion3D release). + return_depth_image: Include a base64-encoded 16-bit PNG depth image. + return_point_cloud: Include a base64-encoded binary PLY point cloud. + return_mesh: Include a base64-encoded binary PLY surface mesh. + fx, fy, cx, cy: Camera intrinsics (uses PrimeSense defaults if omitted). + depth_scale: Raw uint16 → metres divisor (default 1000). + depth_trunc: Maximum depth in metres (default 10). + + Returns: + JSON with ``depth_map``, ``min_depth``, ``max_depth``, ``backend_used``, + and optional ``depth_image``, ``point_cloud_ply``, ``mesh_ply`` (base64). + """ + image_bytes = await image.read() + effective_backend = model_backend or _model_override or DEFAULT_DEPTH_MODEL_URL + base_cfg = DepthEstimationAdvanceConfig() + advanced_config = DepthEstimationAdvanceConfig( + fx=fx if fx is not None else base_cfg.fx, + fy=fy if fy is not None else base_cfg.fy, + cx=cx if cx is not None else base_cfg.cx, + cy=cy if cy is not None else base_cfg.cy, + depth_scale=depth_scale if depth_scale is not None else base_cfg.depth_scale, + depth_trunc=depth_trunc if depth_trunc is not None else base_cfg.depth_trunc, + ) + cmd = DepthEstimationCommand( + image_input=image_bytes, + model_backend=effective_backend, + return_depth_image=return_depth_image, + return_point_cloud=return_point_cloud, + return_mesh=return_mesh, + advanced_config=advanced_config, + ) + result = DepthEstimation().run(cmd) + return { + "depth_map": result.depth_map, + "min_depth": result.min_depth, + "max_depth": result.max_depth, + "backend_used": result.backend_used, + "depth_image": b64( + o3d_depth_image_to_png_bytes(result.depth_image) + if result.depth_image is not None + else None + ), + "point_cloud_ply": b64( + o3d_point_cloud_to_ply_bytes(result.point_cloud) + if result.point_cloud is not None + else None + ), + "mesh_ply": b64(o3d_mesh_to_ply_bytes(result.mesh) if result.mesh is not None else None), + } diff --git a/vizion3d/server/rest/serialisation.py b/vizion3d/server/rest/serialisation.py new file mode 100644 index 0000000..ff375ec --- /dev/null +++ b/vizion3d/server/rest/serialisation.py @@ -0,0 +1,42 @@ +""" +Shared serialisation helpers for the vizion3d REST server. + +Converts Open3D geometry objects to wire-format bytes (PNG or binary PLY) +and base64-encodes arbitrary byte payloads for JSON transport. +""" + +import base64 +import io + +import numpy as np +from PIL import Image + +from vizion3d.lifting.utils import create_mesh_ply_binary, create_ply_binary + + +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") + return buf.getvalue() + + +def o3d_point_cloud_to_ply_bytes(pcd) -> bytes: + """Serialise an Open3D PointCloud to binary PLY bytes.""" + points = np.asarray(pcd.points).astype(np.float32) + colors = (np.asarray(pcd.colors) * 255).astype(np.uint8) + return create_ply_binary(points, colors) + + +def o3d_mesh_to_ply_bytes(mesh) -> bytes: + """Serialise an Open3D TriangleMesh to binary PLY bytes.""" + points = np.asarray(mesh.vertices).astype(np.float32) + colors = (np.asarray(mesh.vertex_colors) * 255).astype(np.uint8) + faces = np.asarray(mesh.triangles).astype(np.int32) + return create_mesh_ply_binary(points, colors, faces) + + +def b64(data: bytes | None) -> str | None: + """Base64-encode *data*, or return ``None`` if *data* is ``None``.""" + return base64.b64encode(data).decode() if data is not None else None diff --git a/vizion3d/server/rest/stereo_depth.py b/vizion3d/server/rest/stereo_depth.py new file mode 100644 index 0000000..6794556 --- /dev/null +++ b/vizion3d/server/rest/stereo_depth.py @@ -0,0 +1,126 @@ +""" +REST endpoint for the Stereo Depth feature. + +Registers ``POST /lifting/stereo-depth`` on the ``router`` exported from this +module. Import the router in ``app.py`` and call +``app.include_router(router)``. +""" + +from fastapi import APIRouter, File, Form, UploadFile + +from vizion3d.stereo import StereoDepth, StereoDepthCommand +from vizion3d.stereo.defaults import DEFAULT_STEREO_MODEL_URL +from vizion3d.stereo.models import StereoDepthAdvancedConfig + +from .serialisation import ( + b64, + o3d_depth_image_to_png_bytes, + o3d_mesh_to_ply_bytes, + o3d_point_cloud_to_ply_bytes, +) + +router = APIRouter() + +# Set by configure_model() at server startup when --stereo_model is passed on the CLI. +_model_override: str | None = None + + +def configure_model(path: str) -> None: + """Set a server-wide default stereo model path and pre-load it into handler memory. + + Called from ``app.run()`` when ``--stereo_model`` is supplied on the command line. + After this call the endpoint uses *path* whenever the caller omits + ``model_backend`` from the form data. + """ + global _model_override + _model_override = path + from vizion3d.stereo.handlers import StereoDepthHandler + + StereoDepthHandler.preload(path) + + +@router.post("/stereo-depth") +async def stereo_depth( + left_image: UploadFile = File(...), + right_image: UploadFile = File(...), + model_backend: str | None = Form(None), + return_depth_image: bool = Form(False), + return_point_cloud: bool = Form(False), + return_mesh: bool = Form(False), + focal_length: float | None = Form(None), + cx: float | None = Form(None), + cy: float | None = Form(None), + baseline: float | None = Form(None), + doffs: float | None = Form(None), + z_far: float | None = Form(None), + conf_threshold: float | None = Form(None), + occ_threshold: float | None = Form(None), + scale_factor: float | None = Form(None), +): + """Run stereo depth estimation on a rectified left/right image pair. + + Args: + left_image: Left-camera image file (any PIL-supported format). + right_image: Right-camera image file (same resolution, horizontally offset). + model_backend: S2M2 checkpoint URL or local path (defaults to the vizion3D release). + return_depth_image: Include a base64-encoded 16-bit PNG depth image. + return_point_cloud: Include a base64-encoded binary PLY point cloud. + return_mesh: Include a base64-encoded binary PLY surface mesh. + focal_length: Focal length in pixels (default 1000.0 — override with your calibration). + cx, cy: Principal point in pixels. + baseline: Stereo baseline in millimetres (default 100.0). + doffs: Disparity offset (default 0.0). + z_far: Max depth in metres for point cloud (default 10.0). + conf_threshold: Minimum confidence for point inclusion (default 0.1). + occ_threshold: Minimum occlusion score for point inclusion (default 0.5). + scale_factor: Input downscale factor for speed/quality tradeoff (default 1.0). + + Returns: + JSON with ``depth_map``, ``disparity_map``, ``min_depth``, ``max_depth``, + ``backend_used``, and optional ``depth_image``, ``point_cloud_ply``, + ``mesh_ply`` (base64). + """ + left_bytes = await left_image.read() + right_bytes = await right_image.read() + effective_backend = model_backend or _model_override or DEFAULT_STEREO_MODEL_URL + + base_cfg = StereoDepthAdvancedConfig() + advanced_config = StereoDepthAdvancedConfig( + focal_length=focal_length if focal_length is not None else base_cfg.focal_length, + cx=cx if cx is not None else base_cfg.cx, + cy=cy if cy is not None else base_cfg.cy, + baseline=baseline if baseline is not None else base_cfg.baseline, + doffs=doffs if doffs is not None else base_cfg.doffs, + z_far=z_far if z_far is not None else base_cfg.z_far, + conf_threshold=conf_threshold if conf_threshold is not None else base_cfg.conf_threshold, + occ_threshold=occ_threshold if occ_threshold is not None else base_cfg.occ_threshold, + scale_factor=scale_factor if scale_factor is not None else base_cfg.scale_factor, + ) + cmd = StereoDepthCommand( + left_image=left_bytes, + right_image=right_bytes, + model_backend=effective_backend, + return_depth_image=return_depth_image, + return_point_cloud=return_point_cloud, + return_mesh=return_mesh, + advanced_config=advanced_config, + ) + result = StereoDepth().run(cmd) + return { + "depth_map": result.depth_map, + "disparity_map": result.disparity_map, + "min_depth": result.min_depth, + "max_depth": result.max_depth, + "backend_used": result.backend_used, + "depth_image": b64( + o3d_depth_image_to_png_bytes(result.depth_image) + if result.depth_image is not None + else None + ), + "point_cloud_ply": b64( + o3d_point_cloud_to_ply_bytes(result.point_cloud) + if result.point_cloud is not None + else None + ), + "mesh_ply": b64(o3d_mesh_to_ply_bytes(result.mesh) if result.mesh is not None else None), + } diff --git a/vizion3d/stereo/__init__.py b/vizion3d/stereo/__init__.py new file mode 100644 index 0000000..e32a043 --- /dev/null +++ b/vizion3d/stereo/__init__.py @@ -0,0 +1,74 @@ +""" +Stereo Depth task — direct Python entry point. + +Import :class:`StereoDepth` and run it with a :class:`StereoDepthCommand` to +obtain metric depth maps, disparity maps, point clouds, and meshes from +rectified left/right stereo image pairs. + +Example:: + + from vizion3d.stereo import ( + StereoDepth, + StereoDepthAdvancedConfig, + StereoDepthCommand, + ) + + cmd = StereoDepthCommand( + left_image="left.png", + right_image="right.png", + return_point_cloud=True, + return_mesh=True, + advanced_config=StereoDepthAdvancedConfig( + focal_length=1733.74, + cx=792.27, + cy=541.89, + baseline=536.62, + ), + ) + result = StereoDepth().run(cmd) + print(f"Depth range: {result.min_depth:.2f} – {result.max_depth:.2f} m") +""" + +from vizion3d.core.container import command_bus, register_command_handler + +from .commands import StereoDepthCommand +from .handlers import StereoDepthHandler +from .models import StereoDepthAdvancedConfig, StereoDepthResult + +register_command_handler(StereoDepthCommand, StereoDepthHandler) + + +class StereoDepth: + """Facade for the Stereo Depth task. + + Serves as the primary entry point for stereo depth inference via direct + Python import. Internally dispatches through the CQRS command bus to + :class:`~vizion3d.stereo.handlers.StereoDepthHandler`. + + Example:: + + from vizion3d.stereo import StereoDepth, StereoDepthCommand + + cmd = StereoDepthCommand(left_image=b"...", right_image=b"...") + result = StereoDepth().run(cmd) + """ + + def run(self, command: StereoDepthCommand) -> StereoDepthResult: + """Dispatch *command* through the CQRS bus to the registered handler. + + Args: + command: The stereo inference parameters and flags. + + Returns: + :class:`StereoDepthResult` with metric depth, disparity, and optional + depth image, point cloud, and mesh. + """ + return command_bus.dispatch(command) + + +__all__ = [ + "StereoDepth", + "StereoDepthAdvancedConfig", + "StereoDepthCommand", + "StereoDepthResult", +] diff --git a/vizion3d/stereo/arch/__init__.py b/vizion3d/stereo/arch/__init__.py new file mode 100644 index 0000000..ee88309 --- /dev/null +++ b/vizion3d/stereo/arch/__init__.py @@ -0,0 +1,5 @@ +"""S2M2 stereo model architecture package.""" + +from .s2m2 import S2M2, build_s2m2, s2m2_config_from_checkpoint + +__all__ = ["S2M2", "build_s2m2", "s2m2_config_from_checkpoint"] diff --git a/vizion3d/stereo/arch/attention.py b/vizion3d/stereo/arch/attention.py new file mode 100644 index 0000000..dfd88cb --- /dev/null +++ b/vizion3d/stereo/arch/attention.py @@ -0,0 +1,289 @@ +""" +Multi-head attention modules for the S2M2 stereo matching transformer. + +Provides self-attention and bidirectional cross-attention in both 1D (row-wise) +and 2D (global/flattened) variants, plus feed-forward network layers and the +combined ``GlobalAttnBlock`` / ``BasicAttnBlock`` building blocks used by the +feature pyramid and transformer stages. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + + +class SelfAttn(nn.Module): + """Multi-head self-attention with optional sinc-based relative positional encoding. + + Args: + dim: Input and output channel count. + num_heads: Number of attention heads. + dim_expansion: Channel expansion factor inside Q/K/V projections. + use_pe: If ``True``, adds a learned projection of the 2D relative PE to + each head's output. + """ + + def __init__(self, dim: int, num_heads: int, dim_expansion: int, use_pe: bool): + super().__init__() + self.num_heads = num_heads + self.head_dim = dim_expansion * dim // num_heads + self.scale = self.head_dim**-0.5 + self.use_pe = use_pe + self.q = nn.Linear(dim, dim_expansion * dim, bias=False) + self.k = nn.Linear(dim, dim_expansion * dim, bias=False) + self.v = nn.Linear(dim, dim_expansion * dim, bias=True) + self.proj = nn.Linear(dim_expansion * dim, dim, bias=False) + if use_pe: + self.pe_proj = nn.Linear(32, self.head_dim) + + def forward(self, x: Tensor, pe: Tensor = None) -> Tensor: + """ + Args: + x: Token sequence ``(B, N, C)``. + pe: Relative positional encoding ``(N, N, 32)`` (only used when + ``use_pe=True``). + + Returns: + Attended output ``(B, N, C)``. + """ + B, N, C = x.shape + q = self.q(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2) + k = self.k(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2) + v = self.v(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2) + if self.use_pe: + score = torch.einsum("...ic,...jc->...ij", self.scale * q, k) + attn = score.softmax(dim=-1) + out = torch.einsum("...ij,...jc->...ic", attn, v) + pe_sum = torch.einsum("...nij,ijc->...nic", attn, pe) + out = out + self.pe_proj(pe_sum) + else: + out = F.scaled_dot_product_attention(q, k, v) + return self.proj(out.transpose(1, 2).reshape(B, N, self.num_heads * self.head_dim)) + + +class CrossAttn(nn.Module): + """Bidirectional cross-attention that updates both left and right feature sequences. + + Each direction attends the other's keys/values, allowing both images to + exchange information in a single block. + + Args: + dim: Input and output channel count. + num_heads: Number of attention heads. + dim_expansion: Channel expansion factor inside Q/K/V projections. + """ + + def __init__(self, dim: int, num_heads: int, dim_expansion: int): + super().__init__() + self.num_heads = num_heads + self.head_dim = dim_expansion * dim // num_heads + self.q = nn.Linear(dim, dim_expansion * dim, bias=False) + self.k = nn.Linear(dim, dim_expansion * dim, bias=False) + self.v = nn.Linear(dim, dim_expansion * dim, bias=True) + self.proj = nn.Linear(dim_expansion * dim, dim, bias=False) + + def forward(self, x: Tensor, y: Tensor): + """ + Args: + x: Left feature sequence ``(B, N, C)``. + y: Right feature sequence ``(B, N, C)``. + + Returns: + ``(x_out, y_out)`` — updated left and right sequences ``(B, N, C)``. + """ + B, N, _ = x.shape + qx = self.q(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2) + ky = self.k(y).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2) + vy = self.v(y).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2) + x_out = F.scaled_dot_product_attention(qx, ky, vy) + kx = self.k(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2) + qy = self.q(y).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2) + vx = self.v(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2) + y_out = F.scaled_dot_product_attention(qy, kx, vx) + x_out = self.proj(x_out.transpose(1, 2).reshape(B, N, self.num_heads * self.head_dim)) + y_out = self.proj(y_out.transpose(1, 2).reshape(B, N, self.num_heads * self.head_dim)) + return x_out, y_out + + +class FFN(nn.Module): + """Position-wise two-layer feed-forward network with pre-LayerNorm and residual. + + Args: + dim: Input and output channel count. + dim_expansion: Hidden layer size multiplier. + """ + + def __init__(self, dim: int, dim_expansion: int): + super().__init__() + self.ffn = nn.Sequential( + nn.Linear(dim, dim_expansion * dim), + nn.GELU(), + nn.Linear(dim_expansion * dim, dim), + ) + self.norm_pre = nn.LayerNorm(dim, elementwise_affine=False) + + def forward(self, z: Tensor) -> Tensor: + return self.ffn(self.norm_pre(z)) + z + + +class SelfAttnBlock1D(nn.Module): + """Row-wise self-attention: treats each image row as an independent sequence. + + Reshapes ``(B, H, W, C)`` → ``(B·H, W, C)`` so attention operates along W. + + Args: + dim: Channel count. + num_heads: Attention heads. + dim_expansion: Q/K/V expansion factor. + use_pe: Whether to apply relative PE. + """ + + def __init__(self, dim: int, num_heads: int, dim_expansion: int, use_pe: bool): + super().__init__() + self.attn = SelfAttn(dim, num_heads, dim_expansion, use_pe) + self.norm_pre = nn.LayerNorm(dim, elementwise_affine=False) + + def forward(self, z: Tensor, pe: Tensor = None) -> Tensor: + B, H, W, C = z.shape + z = z.reshape(B * H, W, C) + z = self.attn(self.norm_pre(z), pe) + z + return z.reshape(B, H, W, C) + + +class CrossAttnBlock1D(nn.Module): + """Row-wise bidirectional cross-attention for paired left/right feature maps. + + Expects *z* to hold ``[left, right]`` concatenated along the batch dimension, + i.e. ``z.shape[0]`` is ``2·B``. + + Args: + dim: Channel count. + num_heads: Attention heads. + dim_expansion: Q/K/V expansion factor. + """ + + def __init__(self, dim: int, num_heads: int, dim_expansion: int): + super().__init__() + self.attn = CrossAttn(dim, num_heads, dim_expansion) + self.norm_pre = nn.LayerNorm(dim, elementwise_affine=False) + + def forward(self, z: Tensor) -> Tensor: + x, y = self.norm_pre(z).chunk(2, dim=0) + B, H, W, C = x.shape + x, y = x.reshape(B * H, W, C), y.reshape(B * H, W, C) + x, y = self.attn(x, y) + x, y = x.reshape(B, H, W, C), y.reshape(B, H, W, C) + return torch.cat([x, y], dim=0) + z + + +class SelfAttnBlock2D(nn.Module): + """Global self-attention: flattens H×W into one long sequence per image. + + Args: + dim: Channel count. + num_heads: Attention heads. + dim_expansion: Q/K/V expansion factor. + use_pe: Whether to apply relative PE. + """ + + def __init__(self, dim: int, num_heads: int, dim_expansion: int, use_pe: bool): + super().__init__() + self.attn = SelfAttn(dim, num_heads, dim_expansion, use_pe) + self.norm_pre = nn.LayerNorm(dim, elementwise_affine=False) + + def forward(self, z: Tensor, pe: Tensor = None) -> Tensor: + B, H, W, C = z.shape + z = z.reshape(B, H * W, C) + z = self.attn(self.norm_pre(z), pe) + z + return z.reshape(B, H, W, C).contiguous() + + +class CrossAttnBlock2D(nn.Module): + """Global bidirectional cross-attention for paired left/right feature maps. + + Expects *z* to hold ``[left, right]`` concatenated along batch. + + Args: + dim: Channel count. + num_heads: Attention heads. + dim_expansion: Q/K/V expansion factor. + """ + + def __init__(self, dim: int, num_heads: int, dim_expansion: int): + super().__init__() + self.attn = CrossAttn(dim, num_heads, dim_expansion) + self.norm_pre = nn.LayerNorm(dim, elementwise_affine=False) + + def forward(self, z: Tensor) -> Tensor: + x, y = self.norm_pre(z).chunk(2, dim=0) + B, H, W, C = x.shape + x, y = x.reshape(B, H * W, C), y.reshape(B, H * W, C) + x, y = self.attn(x, y) + x, y = x.reshape(B, H, W, C), y.reshape(B, H, W, C) + return torch.cat([x, y], dim=0) + z + + +class GlobalAttnBlock(nn.Module): + """Global self-attention (+ optional cross-attention) block used in U-Net stages. + + Operates on ``(B, C, H, W)`` conv-format tensors by permuting to + ``(B, H, W, C)`` internally. + + Args: + dim: Channel count. + num_heads: Attention heads. + dim_expansion: Expansion factor. + use_cross_attn: If ``True``, applies 2D cross-attention before self-attention. + use_pe: Whether to apply relative PE in self-attention. + """ + + def __init__( + self, + dim: int, + num_heads: int, + dim_expansion: int, + use_cross_attn: bool = False, + use_pe: bool = False, + ): + super().__init__() + self.self_attn = SelfAttnBlock2D(dim, num_heads, dim_expansion, use_pe) + if use_cross_attn: + self.cross_attn = CrossAttnBlock2D(dim, num_heads, dim_expansion) + self.ffn_c = FFN(dim, dim_expansion) + else: + self.cross_attn = None + self.ffn = FFN(dim, dim_expansion) + + def forward(self, z: Tensor, pe: Tensor = None) -> Tensor: + z = z.permute(0, 2, 3, 1) + if self.cross_attn is not None: + z = self.ffn_c(self.cross_attn(z)) + z = self.ffn(self.self_attn(z, pe)) + return z.permute(0, 3, 1, 2).contiguous() + + +class BasicAttnBlock(nn.Module): + """Row-wise cross-attention then row-wise self-attention, with FFN after each. + + The standard building block for the MRT transformer stages. + + Args: + dim: Channel count. + num_heads: Attention heads. + dim_expansion: Expansion factor. + use_pe: Whether to apply relative PE in self-attention. + """ + + def __init__(self, dim: int, num_heads: int, dim_expansion: int, use_pe: bool = False): + super().__init__() + self.cross_attn = CrossAttnBlock1D(dim, num_heads, dim_expansion) + self.self_attn = SelfAttnBlock1D(dim, num_heads, dim_expansion, use_pe) + self.ffn_c = FFN(dim, dim_expansion) + self.ffn = FFN(dim, dim_expansion) + + def forward(self, z: Tensor, pe: Tensor = None) -> Tensor: + z = z.permute(0, 2, 3, 1) + z = self.ffn_c(self.cross_attn(z)) + z = self.ffn(self.self_attn(z, pe)) + return z.permute(0, 3, 1, 2) diff --git a/vizion3d/stereo/arch/components.py b/vizion3d/stereo/arch/components.py new file mode 100644 index 0000000..427db77 --- /dev/null +++ b/vizion3d/stereo/arch/components.py @@ -0,0 +1,341 @@ +""" +Feature extraction, U-Net pyramid, cost volume, and disparity initialisation for S2M2. + +Contains: +- :class:`FeatureFusion` — gated blending of two same-dimension feature maps. +- :class:`ConvBlock2D` — residual conv block used throughout the U-Net. +- :class:`CNNEncoder` — 4× downsampling CNN backbone. +- :class:`Unet` — multi-scale feature pyramid with bottleneck attention. +- :class:`CostVolume` — efficient stereo cost-volume sampler. +- :class:`DispInit` — optimal-transport disparity initialisation from correlations. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +from .attention import GlobalAttnBlock +from .utils import bilinear_sampler, get_pe, logsumexp_stable + + +class FeatureFusion(nn.Module): + """Gated fusion of two same-resolution feature maps. + + Learns a soft gate *w ∈ (0, 1)* per spatial position so the output is + ``fusion(z0, z1) + w·z0 + (1−w)·z1`` when ``use_gate=True``, or just + ``fusion(z0, z1)`` otherwise. + + Args: + dim: Input and output channel count for both *z0* and *z1*. + kernel_size: Convolution kernel size for the gate and fusion networks. + use_gate: Whether to include the learnable gate (default ``True``). + """ + + def __init__(self, dim: int, kernel_size: int, use_gate: bool = True): + super().__init__() + pad = kernel_size // 2 + self.use_gate = use_gate + if use_gate: + self.feature_gate = nn.Sequential( + nn.Conv2d(2 * dim, dim, kernel_size=kernel_size, padding=pad), + nn.GELU(), + nn.Conv2d(dim, dim, kernel_size=1), + nn.Sigmoid(), + ) + self.feature_fusion = nn.Sequential( + nn.Conv2d(2 * dim, 2 * dim, kernel_size=kernel_size, padding=pad), + nn.GELU(), + nn.Conv2d(2 * dim, dim, kernel_size=1), + ) + + def forward(self, z0: Tensor, z1: Tensor) -> Tensor: + z = torch.cat([z0, z1], dim=1) + if self.use_gate: + eps = 0.01 + w = self.feature_gate(z).clamp(min=eps, max=1 - eps) + return self.feature_fusion(z) + w * z0 + (1 - w) * z1 + return self.feature_fusion(z) + + +class ConvBlock2D(nn.Module): + """Residual 2D conv block: kxk conv path fused with a parallel 1×1 skip path. + + Args: + dim: Channel count (input = output). + kernel_size: Kernel size for the main conv path. + dim_expansion: Hidden-width multiplier for both paths. + """ + + def __init__(self, dim: int, kernel_size: int, dim_expansion: int): + super().__init__() + p = kernel_size // 2 + self.convs = nn.Sequential( + nn.Conv2d(dim, dim_expansion * dim, kernel_size, padding=p), + nn.GELU(), + nn.Conv2d(dim_expansion * dim, dim, kernel_size, padding=p), + ) + self.convs_1x = nn.Sequential( + nn.Conv2d(dim, dim_expansion * dim, 1), + nn.ReLU(), + nn.Conv2d(dim_expansion * dim, dim, 1), + ) + + def forward(self, z: Tensor) -> Tensor: + return self.convs(z) + self.convs_1x(z) + + +class CNNEncoder(nn.Module): + """4× downsampling CNN backbone producing features at 2× and 4× strides. + + Processes concatenated left+right images (along the batch dim) so both + images share weights. The two outputs are split externally. + + Args: + output_dim: Output channel count at both stride levels. + """ + + def __init__(self, output_dim: int): + super().__init__() + self.conv0 = nn.Sequential(nn.Conv2d(3, 16, 1), nn.GELU(), nn.Conv2d(16, 16, 1)) + self.conv1_down = nn.Sequential( + nn.Conv2d(16, 64, 5, stride=2, padding=2), + nn.GELU(), + nn.Conv2d(64, output_dim, 3, padding=1), + ) + self.norm1 = nn.GroupNorm(8, output_dim) + self.conv2 = nn.Sequential( + nn.Conv2d(output_dim, output_dim, 3, padding=1), + nn.GELU(), + nn.Conv2d(output_dim, output_dim, 3, padding=1), + ) + self.conv2_down = nn.Sequential(nn.Conv2d(output_dim, output_dim, 3, stride=2, padding=1)) + + def forward(self, x: Tensor): + """ + Args: + x: ``(2B, 3, H, W)`` — left and right images stacked along batch. + + Returns: + ``(feature_4x, feature_2x)`` — downsampled 4× and 2× feature maps, + both ``(2B, output_dim, H/4, W/4)`` and ``(2B, output_dim, H/2, W/2)``. + """ + x = self.conv0(x) + x_2x = self.norm1(self.conv1_down(x)) + x_2x = self.conv2(x_2x) + x_2x + return self.conv2_down(x_2x), x_2x + + +class Unet(nn.Module): + """Multi-scale U-Net feature pyramid with global attention at the bottleneck. + + Encoder downsamples three times via average pooling; decoder upsamples back + with gated skip connections. Global attention blocks sit at the lowest + resolution for long-range context. + + Args: + dims: Channel counts at the three pyramid levels ``[d0, d1, d2]``. + dim_expansion: Conv/attention expansion factor. + use_pe: Whether to use relative PE in bottleneck attention. + n_attn: Number of global-attention blocks at the bottleneck. + use_gate_fusion: Whether to use gated skip connections. + """ + + def __init__( + self, + dims: list, + dim_expansion: int, + use_pe: bool, + n_attn: int = 1, + use_gate_fusion: bool = True, + ): + super().__init__() + self.use_pe = use_pe + self.down_conv0 = nn.Sequential(nn.AvgPool2d(2), nn.Conv2d(dims[0], dims[1], 1)) + self.down_conv1 = nn.Sequential(nn.AvgPool2d(2), nn.Conv2d(dims[1], dims[2], 1)) + self.down_conv2 = nn.Sequential(nn.AvgPool2d(2), nn.Conv2d(dims[2], dims[2], 1)) + self.up_conv0 = nn.Sequential( + nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False), + nn.Conv2d(dims[1], dims[0], 1), + ) + self.up_conv1 = nn.Sequential( + nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False), + nn.Conv2d(dims[2], dims[1], 1), + ) + self.up_conv2 = nn.Sequential( + nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False), + nn.Conv2d(dims[2], dims[2], 1), + ) + self.concat_conv0 = FeatureFusion(dims[0], 1, use_gate_fusion) + self.concat_conv1 = FeatureFusion(dims[1], 1, use_gate_fusion) + self.concat_conv2 = FeatureFusion(dims[2], 1, use_gate_fusion) + self.enc0 = ConvBlock2D(dims[0], 3, dim_expansion) + self.enc1 = ConvBlock2D(dims[1], 3, dim_expansion) + self.enc2 = ConvBlock2D(dims[2], 3, dim_expansion) + self.enc3s = nn.ModuleList( + [GlobalAttnBlock(dims[2], 8, dim_expansion, False, use_pe) for _ in range(n_attn)] + ) + self.dec0 = ConvBlock2D(dims[0], 3, dim_expansion) + self.dec1 = ConvBlock2D(dims[1], 3, dim_expansion) + self.dec2 = ConvBlock2D(dims[2], 3, dim_expansion) + self.dec3s = nn.ModuleList( + [GlobalAttnBlock(dims[2], 8, dim_expansion, False, False) for _ in range(n_attn)] + ) + + def forward(self, z: Tensor): + """ + Args: + z: Input feature map ``(B, dims[0], H, W)``. + + Returns: + ``(z0, z1, z2, z3)`` — multi-scale outputs at strides 1×, 2×, 4×, 8× + relative to the input spatial size. + """ + pe = None + if self.use_pe: + H, W = z.shape[-2:] + pe = get_pe(H // 8, W // 8, 32, z.dtype, z.device) + z0 = self.enc0(z) + z1 = self.enc1(self.down_conv0(z0)) + z2 = self.enc2(self.down_conv1(z1)) + z3 = self.down_conv2(z2) + for blk in self.enc3s: + z3 = blk(z3, pe) + for blk in self.dec3s: + z3 = blk(z3, pe) + z2_new = self.dec2(self.concat_conv2(z2, self.up_conv2(z3))) + z1_new = self.dec1(self.concat_conv1(z1, self.up_conv1(z2_new))) + z0_new = self.dec0(self.concat_conv0(z0, self.up_conv0(z1_new))) + return z0_new, z1_new, z2_new, z3 + + +class CostVolume: + """Efficient stereo cost-volume sampler for local refinement. + + Pre-computes the dense correlation matrix and supports fast local lookups + around a running disparity estimate during GRU-based refinement. + + Args: + cv: Dense correlation volume ``(B, H, W, W)`` — left × right feature dot-products. + coords: Left-image x-coordinates ``(B, H, W, 1)`` used to anchor lookups. + radius: Disparity search radius; lookups cover ``[disp−radius, disp+radius]``. + """ + + def __init__(self, cv: Tensor, coords: Tensor, radius: int): + self.radius = radius + dx = torch.linspace(-radius, radius, 2 * radius + 1, device=cv.device, dtype=cv.dtype) + self.dx = dx.reshape(1, 1, 2 * radius + 1, 1) + b, h, w, w2 = cv.shape + self.cv = cv.reshape(b * h * w, 1, 1, w2) + self.cv_2x = F.avg_pool2d(self.cv, kernel_size=[1, 2]) + self.cv = self.cv.reshape(b * h, 1, w, w2) + self.cv_2x = self.cv_2x.reshape(b * h, 1, w, w2 // 2) + self.coords = coords.reshape(b * h * w, 1, 1, 1) + + def __call__(self, disp: Tensor): + """Sample full-res and half-res cost values around the current disparity. + + Args: + disp: Current disparity estimate ``(B, 1, H, W)``. + + Returns: + ``(corrs, corrs_2x)`` — sampled cost volumes at full and half resolution, + each ``(B, 2·radius+1, H, W)``. + """ + b, _, h, w = disp.shape + dx = self.dx + x0 = (self.coords - disp.reshape(b * h * w, 1, 1, 1) + dx).reshape(b * h, w, -1, 1) + y0 = (self.coords + 0 * dx).reshape(b * h, w, -1, 1) + corrs = bilinear_sampler(self.cv, torch.cat([x0, y0], dim=-1)) + corrs = corrs.reshape(b, h, w, 2 * self.radius + 1).permute(0, 3, 1, 2) + x0_2 = (self.coords / 2 - disp.reshape(b * h * w, 1, 1, 1) / 2 + dx).reshape( + b * h, w, -1, 1 + ) + corrs_2x = bilinear_sampler(self.cv_2x, torch.cat([x0_2, y0], dim=-1)) + corrs_2x = corrs_2x.reshape(b, h, w, 2 * self.radius + 1).permute(0, 3, 1, 2) + return corrs, corrs_2x + + +class DispInit(nn.Module): + """Optimal-transport disparity initialisation from dense feature correlations. + + Computes the dense left-right correlation volume, applies a Sinkhorn OT solver + to produce a soft correspondence distribution, and extracts the expected disparity + via a soft-argmax over a local window around the peak. + + Args: + dim: Feature channel count. + ot_iter: Number of Sinkhorn iterations. + use_positivity: Mask the upper triangle so only positive (left→right) disparities + are considered — appropriate for well-rectified stereo pairs. + """ + + def __init__(self, dim: int, ot_iter: int, use_positivity: bool): + super().__init__() + self.layer_norm = nn.LayerNorm(dim, elementwise_affine=True) + self.ot_iter = ot_iter + self.use_positivity = use_positivity + + def _sinkhorn(self, attn: Tensor, log_mu: Tensor, log_nu: Tensor) -> Tensor: + """Run *ot_iter* steps of the Sinkhorn algorithm in log-space.""" + v = log_nu - logsumexp_stable(attn, dim=2) + u = log_mu - logsumexp_stable(attn + v.unsqueeze(2), dim=3) + for _ in range(self.ot_iter - 1): + v = log_nu - logsumexp_stable(attn + u.unsqueeze(3), dim=2) + u = log_mu - logsumexp_stable(attn + v.unsqueeze(2), dim=3) + return attn + u.unsqueeze(3) + v.unsqueeze(2) + + def _optimal_transport(self, attn: Tensor) -> Tensor: + """Convert a raw attention matrix to a normalised transport plan.""" + bs, h, w, _ = attn.shape + dtype = attn.dtype + marginal = torch.cat( + [torch.ones([w], device=attn.device), torch.tensor([w], device=attn.device)] + ) / (2 * w) + log_mu = marginal.log().reshape(1, 1, w + 1) + log_nu = marginal.log().reshape(1, 1, w + 1) + attn = F.pad(attn, (0, 1, 0, 1), "constant", 0) + attn = self._sinkhorn(attn, log_mu, log_nu) + log_const = torch.log(torch.tensor(w, dtype=dtype, device=attn.device) * 2) + return (attn[:, :, :-1, :-1] + log_const).exp().to(dtype) + + def forward(self, feature: Tensor): + """ + Args: + feature: Joint left+right feature tensor ``(2B, C, H, W)``. + + Returns: + ``(disparity, confidence, occlusion, cost_volume)`` — all at the same + spatial resolution as *feature*. + """ + dtype = feature.dtype + device = feature.device + w = feature.shape[-1] + x_grid = torch.linspace(0, w - 1, w, device=device, dtype=dtype) + mask = ( + torch.triu(torch.ones((w, w), dtype=torch.bool, device=device), diagonal=1) + if self.use_positivity + else torch.zeros((w, w), dtype=torch.bool, device=device) + ) + + feature0, feature1 = self.layer_norm(feature.permute(0, 2, 3, 1)).chunk(2, dim=0) + cv = torch.einsum("...hic,...hjc->...hij", feature0, feature1) + cv_mask = cv.masked_fill(mask, -1e4) + prob = self._optimal_transport(cv_mask) + masked_prob = prob.masked_fill(mask, 0) + + prob_max_ind = masked_prob.argmax(dim=3).unsqueeze(3) + prob_l = 2 + masked_prob_pad = F.pad(masked_prob, (prob_l, prob_l), "constant", 0) + conf = 0 + correspondence_left = 0 + for idx in range(2 * prob_l + 1): + weight = torch.gather(masked_prob_pad, index=prob_max_ind + idx, dim=3) + conf += weight + correspondence_left += weight * (prob_max_ind + idx - prob_l) + eps = 1e-4 + correspondence_left = (correspondence_left + eps) / (conf + eps) + disparity = (x_grid.reshape(1, 1, w) - correspondence_left.squeeze(3)).unsqueeze(1) + conf = conf.unsqueeze(1).squeeze(-1) + occ = masked_prob.sum(dim=3).unsqueeze(1) + return disparity, conf, occ, cv diff --git a/vizion3d/stereo/arch/refiners.py b/vizion3d/stereo/arch/refiners.py new file mode 100644 index 0000000..53fdd76 --- /dev/null +++ b/vizion3d/stereo/arch/refiners.py @@ -0,0 +1,281 @@ +""" +Disparity refinement modules for S2M2. + +Contains: +- :class:`ConvGRU` — separable convolutional GRU for iterative updates. +- :class:`UpsampleMask4x` — learned 4× disparity upsampler. +- :class:`UpsampleMask1x` — learned 1× (or 2×) guided upsampler. +- :class:`GlobalRefiner` — U-Net-based global gap-filling pass. +- :class:`LocalRefiner` — cost-volume-guided local GRU refinement. +""" + +from typing import Callable + +import torch +import torch.nn as nn +from torch import Tensor + +from .components import Unet + + +class ConvGRU(nn.Module): + """Separable convolutional GRU: horizontal then vertical 1D convolutions. + + Runs two sequential GRU steps — one with a vertical kernel then one with a + horizontal kernel — to propagate context efficiently across 2D feature maps. + + Args: + hidden_dim: Recurrent state channel count. + input_dim: Input feature channel count. + kernel_size: Kernel size along the non-unit spatial dimension. + """ + + def __init__(self, hidden_dim: int = 128, input_dim: int = 128, kernel_size: int = 3): + super().__init__() + self.convz1 = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, [kernel_size, 1], padding=[kernel_size // 2, 0] + ) + self.convr1 = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, [kernel_size, 1], padding=[kernel_size // 2, 0] + ) + self.convq1 = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, [kernel_size, 1], padding=[kernel_size // 2, 0] + ) + self.convz2 = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, [1, kernel_size], padding=[0, kernel_size // 2] + ) + self.convr2 = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, [1, kernel_size], padding=[0, kernel_size // 2] + ) + self.convq2 = nn.Conv2d( + hidden_dim + input_dim, hidden_dim, [1, kernel_size], padding=[0, kernel_size // 2] + ) + + def forward(self, h: Tensor, x: Tensor) -> Tensor: + """ + Args: + h: Current hidden state ``(B, hidden_dim, H, W)``. + x: Input features ``(B, input_dim, H, W)``. + + Returns: + Updated hidden state ``(B, hidden_dim, H, W)``, dtype matches *x*. + """ + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz1(hx)) + r = torch.sigmoid(self.convr1(hx)) + h = (1 - z) * h + z * torch.tanh(self.convq1(torch.cat([r * h, x], dim=1))) + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz2(hx)) + r = torch.sigmoid(self.convr2(hx)) + h = (1 - z) * h + z * torch.tanh(self.convq2(torch.cat([r * h, x], dim=1))) + return h.to(x.dtype) + + +class UpsampleMask4x(nn.Module): + """Learned convex-combination 4× upsampler using multi-scale feature context. + + Produces a ``(B, 9, H·4, W·4)`` weight map (via a transposed-conv step) that is + used to combine 3×3 neighbourhood patches of the 4×-downsampled disparity map. + + Args: + dim: Context feature channel count. + """ + + def __init__(self, dim: int): + super().__init__() + self.conv_x = nn.ConvTranspose2d(dim, 64, 2, stride=2) + self.conv_y = nn.Conv2d(dim, 64, 3, padding=1) + self.conv_concat = nn.Sequential( + nn.Conv2d(128, 128, 3, padding=1), + nn.ReLU(inplace=False), + nn.ConvTranspose2d(128, 9, 2, stride=2), + ) + + def forward(self, feat_x: Tensor, feat_y: Tensor) -> Tensor: + """ + Args: + feat_x: GRU hidden state ``(B, dim, H, W)``. + feat_y: 2× feature map ``(B, dim, H·2, W·2)``. + + Returns: + 9-channel upsample weight map ``(B, 9, H·4, W·4)``. + """ + return self.conv_concat(torch.cat([self.conv_x(feat_x), self.conv_y(feat_y)], dim=1)) + + +class UpsampleMask1x(nn.Module): + """Guided 1× (or 2×) disparity upsampler using disparity, RGB, and context. + + Produces a ``(B, 9, H, W)`` (or ``(B, 9, H·2, W·2)`` when ``output_upsample`` + is enabled in S2M2) weight map for subpixel-accurate disparity refinement. + + Args: + dim: Context feature channel count. + """ + + def __init__(self, dim: int): + super().__init__() + self.conv_disp = nn.Sequential( + nn.ConvTranspose2d(1, 16, 3, padding=1), nn.ReLU(inplace=False) + ) + self.conv_rgb = nn.Sequential( + nn.ConvTranspose2d(3, 16, 3, padding=1), nn.ReLU(inplace=False) + ) + self.conv_ctx = nn.ConvTranspose2d(dim, 16, 2, stride=2) + self.conv_concat = nn.Sequential( + nn.Conv2d(48, 48, 3, padding=1), + nn.ReLU(inplace=False), + nn.ConvTranspose2d(48, 9, 1), + ) + + def forward(self, disp: Tensor, rgb: Tensor, ctx: Tensor) -> Tensor: + """ + Args: + disp: 4×-upsampled disparity ``(B, 1, H, W)``. + rgb: Normalised left image ``(B, 3, H, W)``. + ctx: 2× feature map ``(B, dim, H/2, W/2)``. + + Returns: + 9-channel upsample filter weights ``(B, 9, H, W)``. + """ + return self.conv_concat( + torch.cat([self.conv_disp(disp), self.conv_rgb(rgb), self.conv_ctx(ctx)], dim=1) + ) + + +class GlobalRefiner(nn.Module): + """U-Net-based global gap-filling pass that corrects low-confidence regions. + + Uses the confidence mask to suppress reliable disparities and predict + corrections only where the OT initialisation was uncertain. + + Args: + feature_channels: Context feature channel count. + """ + + def __init__(self, feature_channels: int): + super().__init__() + ch = feature_channels + self.init_feat = nn.Sequential( + nn.Conv2d(2 + ch, ch, 3, padding=1), nn.GELU(), nn.Conv2d(ch, ch, 1) + ) + self.refine_unet = Unet([ch, ch, ch], 1, False, n_attn=1, use_gate_fusion=True) + self.out_feat = nn.Sequential(nn.Conv2d(ch, 1, 3, padding=1)) + + def forward(self, ctx: Tensor, disp: Tensor, conf: Tensor) -> Tensor: + """ + Args: + ctx: Left context features ``(B, C, H, W)``. + disp: Initial disparity estimate ``(B, 1, H, W)``. + conf: Per-pixel confidence ``(B, 1, H, W)`` in ``[0, 1]``. + + Returns: + Refined disparity ``(B, 1, H, W)`` — confident regions are kept, + uncertain regions are replaced by the U-Net prediction. + """ + mask = 1.0 * (conf > 0.2) + conf_logit = (mask * conf).logit(eps=1e-1) + feat = self.init_feat(torch.cat([disp / 1e2 * mask, conf_logit, ctx], dim=1).to(disp.dtype)) + disp_update = self.out_feat(self.refine_unet(feat)[0]) * 1e2 + return (mask * disp + (1 - mask) * disp_update).to(disp.dtype) + + +class LocalRefiner(nn.Module): + """Cost-volume-guided iterative local refinement using a ConvGRU. + + Each call to :meth:`forward` performs one GRU update step, reading cost-volume + correlations around the current disparity estimate and jointly updating the + disparity, confidence, and occlusion maps. + + Args: + feature_channels: Context feature channel count. + dim_expansion: U-Net expansion factor. + radius: Cost-volume lookup radius. + use_gate_fusion: Whether to use gated U-Net skip connections. + """ + + def __init__( + self, + feature_channels: int, + dim_expansion: int, + radius: int, + use_gate_fusion: bool, + ): + super().__init__() + ch = feature_channels + r = radius + self.disp_feat = nn.Sequential( + nn.Conv2d(1, 96, 3, padding=1), nn.GELU(), nn.Conv2d(96, 96, 3, padding=1) + ) + self.corr_feat1 = nn.Sequential( + nn.Conv2d((2 * r + 1), 96, 1), nn.GELU(), nn.Conv2d(96, 64, 1) + ) + self.corr_feat2 = nn.Sequential( + nn.Conv2d((2 * r + 1), 96, 1), nn.GELU(), nn.Conv2d(96, 64, 1) + ) + self.conf_occ_feat = nn.Sequential( + nn.Conv2d(2, 64, 3, padding=1), nn.GELU(), nn.Conv2d(64, 32, 1) + ) + self.disp_corr_ctx_cat = nn.Sequential( + nn.Conv2d(256 + ch, 2 * ch, 1), nn.GELU(), nn.Conv2d(2 * ch, ch, 3, padding=1) + ) + self.refine_unet = Unet( + [ch, ch, 2 * ch], dim_expansion, False, n_attn=1, use_gate_fusion=use_gate_fusion + ) + self.disp_update = nn.Sequential( + nn.Conv2d(ch, ch, 3, padding=1), + nn.GELU(), + nn.Conv2d(ch, 1, 3, padding=1, bias=False), + ) + self.conf_occ_update = nn.Sequential( + nn.Conv2d(ch, ch, 3, padding=1), + nn.GELU(), + nn.Conv2d(ch, 2, 3, padding=1, bias=False), + ) + self.gru = ConvGRU(ch, ch, 3) + + def forward( + self, + hidden: Tensor, + ctx: Tensor, + disp: Tensor, + conf: Tensor, + occ: Tensor, + cv_fn: Callable, + ): + """Perform one refinement step. + + Args: + hidden: GRU hidden state ``(B, C, H, W)``. + ctx: Left context features ``(B, C, H, W)``. + disp: Current disparity ``(B, 1, H, W)``. + conf: Current confidence ``(B, 1, H, W)``. + occ: Current occlusion ``(B, 1, H, W)``. + cv_fn: Callable that accepts *disp* and returns ``(corrs, corrs_2x)``. + + Returns: + ``(hidden_new, disp_new, conf_new, occ_new)`` — all updated tensors. + """ + conf_logit = conf.logit(eps=1e-2) + occ_logit = occ.logit(eps=1e-2) + corr1, corr2 = cv_fn(disp) + cat_in = torch.cat( + [ + self.disp_feat(disp / 1e2), + self.corr_feat1(corr1 / 16), + self.corr_feat2(corr2 / 16), + ctx, + self.conf_occ_feat(torch.cat([conf_logit, occ_logit], dim=1).to(disp.dtype)), + ], + dim=1, + ).to(disp.dtype) + refine_feat = self.refine_unet(self.disp_corr_ctx_cat(cat_in))[0] + hidden_new = self.gru(hidden, refine_feat) + disp_update = self.disp_update(hidden_new) + conf_update, occ_update = self.conf_occ_update(hidden_new).chunk(2, dim=1) + return ( + hidden_new.to(disp.dtype), + (disp + disp_update).to(disp.dtype), + torch.sigmoid(conf_update + conf_logit).to(disp.dtype), + torch.sigmoid(occ_update + occ_logit).to(disp.dtype), + ) diff --git a/vizion3d/stereo/arch/s2m2.py b/vizion3d/stereo/arch/s2m2.py new file mode 100644 index 0000000..6af860e --- /dev/null +++ b/vizion3d/stereo/arch/s2m2.py @@ -0,0 +1,350 @@ +""" +S2M2 stereo matching transformer — model definition and checkpoint helpers. + +Model architecture constants are hardcoded here and must not be exposed as user +configuration. The only external knob is the checkpoint variant (S / M / L / XL), +which is detected automatically from the checkpoint filename. + +Architecture constants (fixed for all inference): + _DIM_EXPANSION = 1 — Q/K/V width multiplier (not a tunable hyperparameter). + _OT_ITER = 3 — Sinkhorn iterations for disparity initialisation. + _USE_POSITIVITY = True — clamp disparity ≥ 0 (correct for rectified pairs). + _OUTPUT_UPSAMPLE = False — training-only 2× flag; never needed at inference. + _REFINE_ITER = 3 — local GRU refinement steps per inference call. + +Variant configs (feature_channels, num_transformer): + S → (128, 1) + M → (192, 2) + L → (256, 3) ← default / recommended + XL → (384, 3) +""" + +from __future__ import annotations + +from pathlib import Path + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +from .attention import BasicAttnBlock, GlobalAttnBlock +from .components import CNNEncoder, CostVolume, DispInit, FeatureFusion, Unet +from .refiners import GlobalRefiner, LocalRefiner, UpsampleMask1x, UpsampleMask4x +from .utils import custom_unfold + +# ── Model-level constants (must not be user-configurable) ───────────────────── + +_DIM_EXPANSION: int = 1 +_OT_ITER: int = 3 +_USE_POSITIVITY: bool = True +_OUTPUT_UPSAMPLE: bool = False +_REFINE_ITER: int = 3 + +# Per-variant channel / transformer counts — detected from checkpoint filename +_VARIANT_CONFIGS: dict[str, dict] = { + "S": {"feature_channels": 128, "num_transformer": 1}, + "M": {"feature_channels": 192, "num_transformer": 2}, + "L": {"feature_channels": 256, "num_transformer": 3}, + "XL": {"feature_channels": 384, "num_transformer": 3}, +} +_DEFAULT_VARIANT: str = "L" + + +def s2m2_config_from_checkpoint(model_path: str) -> dict: + """Return the architecture config for the S2M2 variant inferred from *model_path*. + + Detection checks the filename (case-insensitive) for ``-S``, ``-M``, ``-L``, + or ``-XL`` suffixes (with or without the ``.pth``/``.pt`` extension). + Falls back to the L variant when no match is found. + + Args: + model_path: Local file path to the checkpoint. + + Returns: + Dict with ``"feature_channels"`` and ``"num_transformer"`` keys. + """ + stem = Path(model_path).stem.upper() + for variant in ("XL", "L", "M", "S"): # XL before L to avoid partial match + if stem.endswith(f"-{variant}") or stem.endswith(f"_{variant}"): + return _VARIANT_CONFIGS[variant] + return _VARIANT_CONFIGS[_DEFAULT_VARIANT] + + +# ── MRT transformer block ───────────────────────────────────────────────────── + + +class MRT(nn.Module): + """Multi-scale recurrent transformer block — one stage of the StackedMRT. + + Implements a U-Net-shaped encoder-decoder path where each level applies + one :class:`BasicAttnBlock` (cross + self attention) and the bottleneck + applies two :class:`GlobalAttnBlock` steps with cross-attention. + + Args: + dims: Channel counts at three resolution levels ``[d0, d1, d2]``. + num_heads: Base head count (multiplied per level: 1×, 2×, 4×, 8×). + dim_expansion: Attention expansion factor. + use_gate_fusion: Whether to use gated skip connections. + """ + + def __init__(self, dims: list, num_heads: int, dim_expansion: int, use_gate_fusion: bool): + super().__init__() + self.down_conv0 = nn.Sequential(nn.AvgPool2d(2), nn.Conv2d(dims[0], dims[1], 1)) + self.down_conv1 = nn.Sequential(nn.AvgPool2d(2), nn.Conv2d(dims[1], dims[2], 1)) + self.down_conv2 = nn.Sequential(nn.AvgPool2d(2), nn.Conv2d(dims[2], dims[2], 1)) + self.up_conv0 = nn.Sequential( + nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False), + nn.Conv2d(dims[1], dims[0], 1), + ) + self.up_conv1 = nn.Sequential( + nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False), + nn.Conv2d(dims[2], dims[1], 1), + ) + self.up_conv2 = nn.Sequential( + nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False), + nn.Conv2d(dims[2], dims[2], 1), + ) + self.down_concat1 = FeatureFusion(dims[1], 1, use_gate_fusion) + self.down_concat2 = FeatureFusion(dims[2], 1, use_gate_fusion) + self.down_concat3 = FeatureFusion(dims[2], 1, use_gate_fusion) + self.up_concat0 = FeatureFusion(dims[0], 1, use_gate_fusion) + self.up_concat1 = FeatureFusion(dims[1], 1, use_gate_fusion) + self.up_concat2 = FeatureFusion(dims[2], 1, use_gate_fusion) + self.enc_attn0 = BasicAttnBlock(dims[0], 1 * num_heads, dim_expansion) + self.enc_attn1 = BasicAttnBlock(dims[1], 2 * num_heads, dim_expansion) + self.enc_attn2 = BasicAttnBlock(dims[2], 4 * num_heads, dim_expansion) + self.enc_attn3s = nn.ModuleList( + [GlobalAttnBlock(dims[2], 8 * num_heads, dim_expansion, True) for _ in range(2)] + ) + self.dec_attn0 = BasicAttnBlock(dims[0], 1 * num_heads, dim_expansion) + self.dec_attn1 = BasicAttnBlock(dims[1], 2 * num_heads, dim_expansion) + self.dec_attn2 = BasicAttnBlock(dims[2], 4 * num_heads, dim_expansion) + self.dec_attn3s = nn.ModuleList( + [GlobalAttnBlock(dims[2], 8 * num_heads, dim_expansion, True) for _ in range(2)] + ) + + def forward(self, z0, z1, z2, z3): + """Apply one MRT encoder-decoder pass to four multi-scale feature maps. + + Args: + z0, z1, z2, z3: Feature maps at strides 1×, 2×, 4×, 8× (all ``(2B, C, H, W)``). + + Returns: + Updated ``(z0, z1, z2, z3)`` at the same shapes. + """ + z0 = self.enc_attn0(z0) + z1 = self.enc_attn1(self.down_concat1(z1, self.down_conv0(z0))) + z2 = self.enc_attn2(self.down_concat2(z2, self.down_conv1(z1))) + z3 = self.down_concat3(z3, self.down_conv2(z2)) + for blk in self.enc_attn3s: + z3 = blk(z3) + for blk in self.dec_attn3s: + z3 = blk(z3) + z2 = self.dec_attn2(self.up_concat2(z2, self.up_conv2(z3))) + z1 = self.dec_attn1(self.up_concat1(z1, self.up_conv1(z2))) + z0 = self.dec_attn0(self.up_concat0(z0, self.up_conv0(z1))) + return z0, z1, z2, z3 + + +class StackedMRT(nn.Module): + """Stack of *num_transformer* MRT blocks that refine multi-scale features iteratively. + + Args: + num_transformer: Number of MRT stages to stack. + dims: Channel counts per pyramid level ``[d0, d1, d2]``. + num_heads: Base attention head count. + dim_expansion: Attention expansion factor. + use_gate_fusion: Whether to use gated skip connections. + """ + + def __init__( + self, + num_transformer: int, + dims: list, + num_heads: int, + dim_expansion: int, + use_gate_fusion: bool, + ): + super().__init__() + self.uformer_list = nn.ModuleList( + [MRT(dims, num_heads, dim_expansion, use_gate_fusion) for _ in range(num_transformer)] + ) + + def forward(self, z0, z1, z2, z3) -> Tensor: + """Run all MRT stages in sequence and return the finest-resolution output. + + Returns: + ``z0`` after the final MRT stage, shape ``(2B, dims[0], H/4, W/4)``. + """ + for blk in self.uformer_list: + z0, z1, z2, z3 = blk(z0, z1, z2, z3) + return z0.contiguous() + + +# ── S2M2 top-level model ────────────────────────────────────────────────────── + + +class S2M2(nn.Module): + """Stereo Matching Model with Multi-scale transformer (S2M2). + + Full forward pass: CNN backbone → feature pyramid → stacked MRT transformer + → OT disparity init → global refinement → iterative local GRU refinement + → learned 4× and 1× upsampling. + + Instantiate via :func:`build_s2m2` rather than calling this constructor + directly, so architecture constants are applied consistently. + + Args: + feature_channels: Per-level channel count (variant-dependent). + dim_expansion: Q/K/V expansion factor (hardcoded to ``_DIM_EXPANSION``). + num_transformer: Number of StackedMRT stages (variant-dependent). + use_positivity: Clamp disparity ≥ 0 (hardcoded to ``_USE_POSITIVITY``). + output_upsample: Enable 2× output head (hardcoded to ``_OUTPUT_UPSAMPLE``). + refine_iter: Local GRU refinement steps (hardcoded to ``_REFINE_ITER``). + """ + + def __init__( + self, + feature_channels: int, + dim_expansion: int, + num_transformer: int, + use_positivity: bool = _USE_POSITIVITY, + output_upsample: bool = _OUTPUT_UPSAMPLE, + refine_iter: int = _REFINE_ITER, + ): + super().__init__() + ch = feature_channels + self.use_positivity = use_positivity + self.refine_iter = refine_iter + self.output_upsample = output_upsample + self.cnn_backbone = CNNEncoder(ch) + self.feat_pyramid = Unet( + [ch, ch, 2 * ch], + dim_expansion, + True, + n_attn=num_transformer * 2, + use_gate_fusion=True, + ) + self.transformer = StackedMRT(num_transformer, [ch, ch, 2 * ch], 1, dim_expansion, True) + self.disp_init = DispInit(ch, _OT_ITER, use_positivity) + self.upsample_mask_1x = UpsampleMask1x(ch) + self.upsample_mask_4x = UpsampleMask4x(ch) + self.global_refiner = GlobalRefiner(ch) + self.feat_fusion_layer = FeatureFusion(ch, 3, True) + self.refiner = LocalRefiner(ch, dim_expansion, 4, True) + self.ctx_feat = nn.Sequential(nn.Conv2d(ch, ch, 1), nn.GELU(), nn.Conv2d(ch, ch, 1)) + + def my_load_state_dict(self, state_dict: dict): + """Load *state_dict* with shape-mismatch tolerance. + + Silently skips keys whose tensor shapes differ from the current model + (e.g. when fine-tuning a pretrained model on a different image size). + + Args: + state_dict: Dict of parameter tensors, typically from + ``torch.load(...).get("state_dict", ckpt)``. + """ + own = self.state_dict() + for k in state_dict: + if k in own and state_dict[k].shape != own[k].shape: + state_dict[k] = own[k] + self.load_state_dict(state_dict, strict=False) + + @staticmethod + def _normalize(img0: Tensor, img1: Tensor): + """Normalise uint8 pixel values to ``[-1, 1]``.""" + return (img0 / 255.0 - 0.5) * 2, (img1 / 255.0 - 0.5) * 2 + + def _upsample4x(self, x: Tensor, up_weights: Tensor) -> Tensor: + """Apply the learned 4× convex-combination upsampler.""" + b, c, h, w = x.shape + x_unfold = custom_unfold(x.reshape(b, c, h, w), 3, 1) + x_unfold = F.interpolate(x_unfold, (h * 4, w * 4), mode="nearest").reshape( + b, 9, h * 4, w * 4 + ) + return (x_unfold * up_weights.softmax(dim=1)).sum(1, keepdim=True) + + def _upsample1x(self, disp: Tensor, filter_weights: Tensor) -> Tensor: + """Apply the guided 1× (or 2× when ``output_upsample``) disparity upsampler.""" + disp_unfold = custom_unfold(disp, 3, 1) + if self.output_upsample: + disp_unfold = F.interpolate(disp_unfold, scale_factor=2, mode="nearest") + filter_weights = F.interpolate( + filter_weights, scale_factor=2, mode="bilinear", align_corners=False + ) + return (disp_unfold * filter_weights.softmax(dim=1).to(disp.dtype)).sum(1, keepdim=True) + + def forward(self, img0: Tensor, img1: Tensor): + """Run full stereo forward pass. + + Args: + img0: Left image ``(B, 3, H, W)`` with pixel values in ``[0, 255]``. + img1: Right image ``(B, 3, H, W)`` with pixel values in ``[0, 255]``. + + Returns: + ``(disp_up, occ_up, conf_up)`` — disparity (in pixels at input resolution), + occlusion, and confidence maps, all ``(B, 1, H, W)``. + """ + img0_nor, img1_nor = self._normalize(img0, img1) + feature_4x, feature_2x = self.cnn_backbone(torch.cat([img0_nor, img1_nor], dim=0)) + feature0_2x, _ = feature_2x.chunk(2, dim=0) + feature_py_4x, feature_py_8x, feature_py_16x, feature_py_32x = self.feat_pyramid(feature_4x) + feature_tr_4x = self.transformer( + feature_py_4x, feature_py_8x, feature_py_16x, feature_py_32x + ) + disp, conf, occ, cv = self.disp_init(feature_tr_4x) + feature0_tr_4x, _ = feature_tr_4x.chunk(2, dim=0) + feature0_py_4x, _ = feature_py_4x.chunk(2, dim=0) + disp = self.global_refiner(feature0_tr_4x.contiguous(), disp.detach(), conf.detach()) + if self.use_positivity: + disp = disp.clamp(min=0) + feature0_fusion_4x = self.feat_fusion_layer(feature0_tr_4x, feature0_py_4x) + ctx0 = self.ctx_feat(feature0_fusion_4x) + hidden = torch.tanh(ctx0) + b, c, h, w = feature0_fusion_4x.shape + coords_4x = torch.arange( + w, device=feature0_fusion_4x.device, dtype=feature0_fusion_4x.dtype + ) + cv_fn = CostVolume(cv, coords_4x.reshape(1, 1, w, 1).repeat(b, h, 1, 1), radius=4) + for _ in range(self.refine_iter): + hidden, disp, conf, occ = self.refiner(hidden, ctx0, disp, conf, occ, cv_fn) + if self.use_positivity: + disp = disp.clamp(min=0) + occ = occ * torch.ge(coords_4x.reshape(1, 1, 1, -1) - disp, 0) + upsample_mask = self.upsample_mask_4x(hidden, feature0_2x) + disp_up = self._upsample4x(disp * 4, upsample_mask) + occ_up = self._upsample4x(occ, upsample_mask) + conf_up = self._upsample4x(conf, upsample_mask) + filter_weights = self.upsample_mask_1x(disp_up, img0_nor, feature0_2x) + disp_up = self._upsample1x(disp_up, filter_weights) + occ_up = self._upsample1x(occ_up, filter_weights) + conf_up = self._upsample1x(conf_up, filter_weights) + if self.output_upsample: + disp_up = 2 * disp_up + return disp_up, occ_up, conf_up + + +def build_s2m2(model_path: str) -> S2M2: + """Construct an :class:`S2M2` instance with the correct architecture for *model_path*. + + Architecture constants (_DIM_EXPANSION, _USE_POSITIVITY, _REFINE_ITER, + _OUTPUT_UPSAMPLE) are applied automatically. The variant (S/M/L/XL) is + detected from the checkpoint filename. + + Args: + model_path: Path to the ``.pth`` checkpoint. + + Returns: + Un-loaded :class:`S2M2` in eval-ready configuration (call + ``my_load_state_dict`` + ``.eval()`` to complete initialisation). + """ + cfg = s2m2_config_from_checkpoint(model_path) + return S2M2( + feature_channels=cfg["feature_channels"], + dim_expansion=_DIM_EXPANSION, + num_transformer=cfg["num_transformer"], + use_positivity=_USE_POSITIVITY, + output_upsample=_OUTPUT_UPSAMPLE, + refine_iter=_REFINE_ITER, + ) diff --git a/vizion3d/stereo/arch/utils.py b/vizion3d/stereo/arch/utils.py new file mode 100644 index 0000000..4ac9763 --- /dev/null +++ b/vizion3d/stereo/arch/utils.py @@ -0,0 +1,190 @@ +""" +Image manipulation, positional encoding, and numeric utilities for S2M2. + +This module contains pure tensor operations shared across the model architecture: +padding/cropping helpers, sinc-based positional encoding, cost-volume sampling, +and numerically stable log-sum-exp. +""" + +import math + +import torch +import torch.nn.functional as F +from torch import Tensor + + +def image_pad(img: Tensor, factor: int = 32) -> Tensor: + """Pad spatial dims to the next multiple of *factor* with content-aware fill. + + The border pixels are filled by downsampling then upsampling the padded region + so edge artefacts do not bleed into the network. The original image content is + then written back into the centre of the padded result. + + Args: + img: Float tensor of shape ``(B, C, H, W)``. + factor: Divisibility factor (default 32). + + Returns: + Padded tensor of shape ``(B, C, H', W')`` where ``H'`` and ``W'`` are the + nearest multiples of *factor* ≥ ``H`` and ``W``. + """ + H, W = img.shape[-2:] + H_new = math.ceil(H / factor) * factor + W_new = math.ceil(W / factor) * factor + pad_h = H_new - H + pad_w = W_new - W + img_pad = F.pad(img, (pad_w // 2, pad_w - pad_w // 2, 0, 0), "constant", 0) + img_pad = F.pad(img_pad, (0, 0, pad_h // 2, pad_h - pad_h // 2), "constant", 0) + img_pad_down = F.adaptive_avg_pool2d(img_pad.float(), output_size=[H // factor, W // factor]) + img_pad = F.interpolate(img_pad_down, size=[H_new, W_new], mode="bilinear") + h_s, h_e = pad_h // 2, pad_h - pad_h // 2 + w_s, w_e = pad_w // 2, pad_w - pad_w // 2 + if h_e == 0 and w_e == 0: + img_pad[:, :, h_s:, w_s:] = img + elif h_e == 0: + img_pad[:, :, h_s:, w_s:-w_e] = img + elif w_e == 0: + img_pad[:, :, h_s:-h_e, w_s:] = img + else: + img_pad[:, :, h_s:-h_e, w_s:-w_e] = img + return img_pad + + +def image_crop(img: Tensor, img_shape: tuple) -> Tensor: + """Remove symmetric padding added by :func:`image_pad`. + + Args: + img: Padded tensor of shape ``(B, C, H_pad, W_pad)``. + img_shape: ``(H_orig, W_orig)`` — the target crop size. + + Returns: + Cropped tensor of shape ``(B, C, H_orig, W_orig)``. + """ + H, W = img.shape[-2:] + H_new, W_new = img_shape + crop_h = H - H_new + if crop_h > 0: + s, e = crop_h // 2, crop_h - crop_h // 2 + img = img[:, :, s:-e] + crop_w = W - W_new + if crop_w > 0: + s, e = crop_w // 2, crop_w - crop_w // 2 + img = img[:, :, :, s:-e] + return img + + +def custom_sinc(x: Tensor) -> Tensor: + """Numerically stable sinc: ``sin(π·x)/(π·x)`` with the correct limit 1 at x=0. + + Args: + x: Input tensor (any shape, any dtype). + + Returns: + Sinc values, same shape and dtype as *x*. + """ + return torch.where( + torch.abs(x) < 1e-6, + torch.ones_like(x), + (torch.sin(3.1415 * x) / (3.1415 * x)).to(x.dtype), + ) + + +def custom_unfold(x: Tensor, kernel_size: int = 3, padding: int = 1) -> Tensor: + """Expand a feature map into per-pixel 3×3 neighbourhood patches. + + The output channels contain the ``kernel_size²`` shifted copies of the input, + concatenated along the channel dimension — equivalent to ``nn.Unfold`` but + implemented as a stack of slices to avoid large intermediate tensors. + + Args: + x: Input tensor ``(B, C, H, W)``. + kernel_size: Neighbourhood size (default 3). + padding: Replicate padding applied before extraction (default 1). + + Returns: + Tensor ``(B, C·kernel_size², H, W)``. + """ + B, C, H, W = x.shape + x_pad = F.pad(x, (padding, padding, padding, padding), "replicate") + parts = [] + for i in range(kernel_size): + for j in range(kernel_size): + parts.append(x_pad[:, :, i : i + H, j : j + W]) + return torch.cat(parts, dim=1) + + +def get_pe(h: int, w: int, pe_dim: int, dtype, device) -> Tensor: + """Compute 2D relative positional encoding via sinc interpolation. + + Produces a ``(h·w, h·w, pe_dim)`` tensor whose ``[i, j, :]`` entry encodes + the relative 2D offset between spatial positions *i* and *j*. + + Args: + h: Feature-map height (at the scale where PE is applied). + w: Feature-map width. + pe_dim: Encoding dimensionality (split equally between x and y). + dtype: Torch dtype for the output. + device: Torch device for the output. + + Returns: + Positional encoding tensor ``(h·w, h·w, pe_dim)``. + """ + with torch.no_grad(): + grid_y, grid_x = torch.meshgrid( + torch.linspace(0, h - 1, h, device=device, dtype=dtype), + torch.linspace(0, w - 1, w, device=device, dtype=dtype), + indexing="ij", + ) + rel_x = (grid_x.reshape(-1, 1) - grid_x.reshape(1, -1)).long() + rel_y = (grid_y.reshape(-1, 1) - grid_y.reshape(1, -1)).long() + + sig = 5 / pe_dim + x_pos = torch.linspace(-3, 3, 2 * w + 1, device=device, dtype=dtype).tanh() + dim_t = torch.linspace(-1, 1, pe_dim // 2, device=device, dtype=dtype) + pe_x = F.normalize(custom_sinc((dim_t[None] - x_pos[:, None]) / sig), p=2, dim=-1) + rel_pe_x = pe_x[rel_x + w - 1].reshape(h * w, h * w, pe_dim // 2) + + y_pos = torch.linspace(-3, 3, 2 * h + 1, device=device, dtype=dtype).tanh() + pe_y = F.normalize(custom_sinc((dim_t[None] - y_pos[:, None]) / sig), p=2, dim=-1) + rel_pe_y = pe_y[rel_y + h - 1].reshape(h * w, h * w, pe_dim // 2) + + pe = 0.5 * torch.cat([rel_pe_x, rel_pe_y], dim=2) + return pe.clone() + + +def bilinear_sampler(img: Tensor, coords: Tensor, mode: str = "bilinear") -> Tensor: + """Sample *img* at fractional pixel coordinates using bilinear interpolation. + + Args: + img: Feature map ``(B·H, 1, H_feat, W_feat)`` (squeezed batch format used + by :class:`CostVolume`). + coords: Sampling coordinates ``(B·H, W, K, 2)`` where the last dim is + ``(x, y)`` in pixel units (not normalised). + mode: Interpolation mode forwarded to ``F.grid_sample``. + + Returns: + Sampled values, same shape prefix as *coords* with the channel squeezed. + """ + W = torch.tensor(img.shape[-1], dtype=img.dtype, device=img.device) + H = torch.tensor(img.shape[-2], dtype=img.dtype, device=img.device) + xgrid, ygrid = coords.split([1, 1], dim=-1) + grid = torch.cat([2 * xgrid / (W - 1) - 1, 2 * ygrid / (H - 1) - 1], dim=-1) + return F.grid_sample(img, grid, mode=mode, align_corners=True) + + +def logsumexp_stable(x: Tensor, dim: int, keepdim: bool = False, eps: float = 1e-30) -> Tensor: + """Numerically stable log-sum-exp along *dim* using the max-subtraction trick. + + Args: + x: Input tensor. + dim: Reduction dimension. + keepdim: Whether to keep the reduced dimension. + eps: Floor clamped onto the inner sum to avoid ``log(0)``. + + Returns: + Log-sum-exp values, with *dim* reduced (or kept as size 1 if *keepdim*). + """ + m, _ = x.max(dim=dim, keepdim=True) + y = (x - m).exp().sum(dim=dim, keepdim=True) + y = m + torch.log(torch.clamp(y, min=eps)) + return y if keepdim else y.squeeze(dim) diff --git a/vizion3d/stereo/commands.py b/vizion3d/stereo/commands.py new file mode 100644 index 0000000..b78dd94 --- /dev/null +++ b/vizion3d/stereo/commands.py @@ -0,0 +1,59 @@ +""" +CQRS command payload for the Stereo Depth task. +""" + +from dataclasses import dataclass, field + +from vizion3d.core.cqrs import Command + +from .defaults import DEFAULT_STEREO_MODEL_URL +from .models import StereoDepthAdvancedConfig, StereoDepthResult + + +@dataclass +class StereoDepthCommand(Command[StereoDepthResult]): + """Command payload to trigger a stereo depth inference task. + + Stereo depth produces **real metric depth** (in metres) by matching + corresponding pixels across a rectified left/right image pair and applying + the stereo geometry formula: + + depth_m = baseline_mm × focal_length_px / disparity_px / 1000 + + Attributes: + left_image: The left-camera image. Pass a file-path string or raw + image bytes. The handler auto-detects which form is supplied. + right_image: The right-camera image (same resolution as *left_image*, + from a horizontally-offset camera). Same path/bytes convention. + model_backend: S2M2 checkpoint to use for inference. + + - Default value is the vizion3D release checkpoint URL + (``stereo-depth-s2m2-L.pth``), downloaded on first use and cached + under ``~/.cache/vizion3d/models/``. + Set ``VIZION3D_MODEL_CACHE`` to override the cache directory. + - A local ``.pth`` or ``.pt`` path is loaded directly. + - Any HTTPS URL is downloaded to the cache directory and loaded. + + return_depth_image: When ``True``, the result includes a 16-bit grayscale + ``open3d.geometry.Image`` (dtype ``uint16``) where the full 0–65535 + range maps linearly to ``[min_depth, max_depth]`` in metres. + return_point_cloud: When ``True``, the result includes an + ``open3d.geometry.PointCloud`` unprojected using the stereo camera + intrinsics in ``advanced_config``. Point coordinates are in metres. + return_mesh: When ``True``, the result includes an + ``open3d.geometry.TriangleMesh`` reconstructed from the point cloud + via ball-pivoting. Includes vertex colours. + advanced_config: Camera intrinsics and inference settings. Override any + field to match your stereo rig — e.g. + ``advanced_config=StereoDepthAdvancedConfig(focal_length=1733.74, + cx=792.27, cy=541.89, baseline=536.62)``. + Unspecified fields keep their defaults (1280×720 @ 100 mm baseline). + """ + + left_image: str | bytes + right_image: str | bytes + model_backend: str = DEFAULT_STEREO_MODEL_URL + return_depth_image: bool = False + return_point_cloud: bool = False + return_mesh: bool = False + advanced_config: StereoDepthAdvancedConfig = field(default_factory=StereoDepthAdvancedConfig) diff --git a/vizion3d/stereo/defaults.py b/vizion3d/stereo/defaults.py new file mode 100644 index 0000000..9ebbe46 --- /dev/null +++ b/vizion3d/stereo/defaults.py @@ -0,0 +1,36 @@ +""" +Default model configuration and download utilities for Stereo Depth. + +Model download logic is shared with the lifting module — URLs are resolved, +downloaded on first use, and cached under ``~/.cache/vizion3d/models/`` (or the +directory specified by the ``VIZION3D_MODEL_CACHE`` environment variable). +""" + +from __future__ import annotations + +from pathlib import Path + +from vizion3d.lifting.defaults import resolve_model_backend + +DEFAULT_STEREO_MODEL_URL = ( + "https://github.com/OlafenwaMoses/vizion3D/releases/download/" + "essentials-v1/stereo-depth-s2m2-L.pth" +) +DEFAULT_STEREO_MODEL_FILENAME = "stereo-depth-s2m2-L.pth" + + +def resolve_stereo_model_backend(model_backend: str, cache_dir: Path | None = None) -> str: + """Resolve *model_backend* to a local file path, downloading if needed. + + Delegates to :func:`vizion3d.lifting.defaults.resolve_model_backend` so the + stereo and depth-estimation modules share the same download/cache logic. + + Args: + model_backend: A URL or local file path to the S2M2 checkpoint. + cache_dir: Override for the model cache directory. ``None`` uses + ``default_model_cache_dir()`` from the lifting defaults. + + Returns: + Absolute local path to the resolved checkpoint file. + """ + return resolve_model_backend(model_backend, cache_dir=cache_dir) diff --git a/vizion3d/stereo/handlers.py b/vizion3d/stereo/handlers.py new file mode 100644 index 0000000..66cd2a5 --- /dev/null +++ b/vizion3d/stereo/handlers.py @@ -0,0 +1,372 @@ +""" +CQRS command handler for the Stereo Depth task. + +Orchestrates model loading, image pre-processing, stereo inference, depth +computation, and optional point-cloud / mesh generation. +""" + +from __future__ import annotations + +import contextlib +import io +import threading + +import numpy as np +from PIL import Image + +from vizion3d.core.cqrs import CommandHandler +from vizion3d.lifting.handlers import ( + OPEN3D_CAMERA_TO_IMAGE_VIEW_TRANSFORM, + DepthEstimationHandler, +) + +from .arch import build_s2m2 +from .arch.utils import image_crop, image_pad +from .commands import StereoDepthCommand +from .defaults import resolve_stereo_model_backend +from .models import StereoDepthResult + +# Flip-Y transform reused from the depth estimation handler so point-cloud +# orientation matches the image-plane convention (Y increases downward). +_CAMERA_TO_IMAGE_VIEW = OPEN3D_CAMERA_TO_IMAGE_VIEW_TRANSFORM + + +class StereoDepthHandler(CommandHandler[StereoDepthCommand, StereoDepthResult]): + """Handles :class:`~vizion3d.stereo.commands.StereoDepthCommand` inference requests. + + Models are loaded on first use and cached in a class-level dict so subsequent + calls within the same process reuse weights without re-reading disk. A lock + ensures thread-safe initialisation. + + The inference pipeline: + 1. Resolve / download the checkpoint. + 2. Load and cache the S2M2 model (thread-safe). + 3. Pre-process the image pair (resize, pad to 32-px multiple). + 4. Run stereo inference with mixed-precision autocast where available. + 5. Crop output back to original resolution; optionally upsample if scaled. + 6. Convert disparity → metric depth using the pinhole stereo formula. + 7. Unproject depth + colour into an Open3D PointCloud and optionally mesh. + """ + + _stereo_models: dict = {} + _model_lock = threading.Lock() + + def handle(self, command: StereoDepthCommand) -> StereoDepthResult: + """Execute the full stereo depth inference pipeline. + + Args: + command: Fully populated :class:`StereoDepthCommand`. + + Returns: + :class:`StereoDepthResult` with depth map, disparity map, and + optional depth image, point cloud, and mesh. + """ + model_id = resolve_stereo_model_backend(command.model_backend) + + def _load_image(src: str | bytes) -> Image.Image: + if isinstance(src, str): + return Image.open(src).convert("RGB") + return Image.open(io.BytesIO(src)).convert("RGB") + + left_pil = _load_image(command.left_image) + right_pil = _load_image(command.right_image) + + inference = self._run_s2m2(model_id, left_pil, right_pil, command.advanced_config) + if isinstance(inference, tuple): + disp_np, occ_np, conf_np = inference + else: + # Backwards-compatible path for tests or custom subclasses that patch + # _run_s2m2 to return only a disparity map. + disp_np, occ_np, conf_np = inference, None, None + + cfg = command.advanced_config + + # Disparity → metric depth (millimetre formula → metres) + with np.errstate(divide="ignore", invalid="ignore"): + depth_mm = cfg.baseline * cfg.focal_length / (disp_np + cfg.doffs) + depth_mm[disp_np <= 0] = 0.0 + depth_m = depth_mm / 1000.0 + + min_depth = float(np.min(depth_m)) + max_depth = float(np.max(depth_m)) + depth_map: list[list[float]] = depth_m.astype(np.float32).tolist() + disparity_map: list[list[float]] = disp_np.astype(np.float32).tolist() + + depth_image = None + if command.return_depth_image: + try: + import open3d as o3d + except ImportError: + raise ImportError( + "open3d is required for depth image output. Pin to Python 3.12 and run: uv sync" + ) + depth_range = max_depth - min_depth + if depth_range > 0: + normalized = (depth_m - min_depth) / depth_range + else: + normalized = np.zeros_like(depth_m) + depth_16bit = (normalized * 65535).astype(np.uint16) + depth_image = o3d.geometry.Image(depth_16bit) + + point_cloud = None + mesh = None + if command.return_point_cloud or command.return_mesh: + try: + import open3d as o3d + except ImportError: + raise ImportError( + "open3d is required for point cloud / mesh output. " + "Pin to Python 3.12 and run: uv sync" + ) + point_cloud, mesh = self._unproject( + left_pil, + depth_m, + disp_np, + cfg, + o3d, + occ_np=occ_np, + conf_np=conf_np, + want_cloud=command.return_point_cloud, + want_mesh=command.return_mesh, + ) + + return StereoDepthResult( + depth_map=depth_map, + disparity_map=disparity_map, + min_depth=min_depth, + max_depth=max_depth, + backend_used=model_id, + depth_image=depth_image, + point_cloud=point_cloud, + mesh=mesh, + point_cloud_scale=1.0, + ) + + # ── Model loading ───────────────────────────────────────────────────────── + + def _load_s2m2(self, model_path: str): + """Load and cache an S2M2 checkpoint (thread-safe double-checked locking). + + Args: + model_path: Resolved local path to the ``.pth`` checkpoint. + + Returns: + ``(model, torch, device)`` tuple ready for inference. + """ + if model_path in self._stereo_models: + return self._stereo_models[model_path] + + with self._model_lock: + if model_path in self._stereo_models: + return self._stereo_models[model_path] + + try: + import torch + except ImportError as exc: + raise ImportError("S2M2 stereo checkpoints require torch. Run: uv sync") from exc + + device = self._torch_device(torch) + model = build_s2m2(model_path) + + try: + ckpt = torch.load(model_path, map_location="cpu", weights_only=True) + except TypeError: + ckpt = torch.load(model_path, map_location="cpu") + + state_dict = ckpt.get("state_dict", ckpt) + model.my_load_state_dict(state_dict) + model = model.to(device).eval() + + self._stereo_models[model_path] = (model, torch, device) + return self._stereo_models[model_path] + + # ── Inference ───────────────────────────────────────────────────────────── + + def _run_s2m2( + self, + model_path: str, + left: Image.Image, + right: Image.Image, + cfg, + ) -> np.ndarray: + """Run one stereo forward pass and return the disparity map (in pixels). + + Handles input scaling, padding, autocast, output cropping, and upsampling. + + Args: + model_path: Resolved local path to the checkpoint. + left: Left PIL image (RGB). + right: Right PIL image (RGB). + cfg: :class:`~vizion3d.stereo.models.StereoDepthAdvancedConfig`. + + Returns: + Tuple of ``(disparity, occlusion, confidence)`` float32 numpy arrays, + each shape ``(H, W)`` at the original image resolution. + """ + model, torch, device = self._load_s2m2(model_path) + + import torch.nn.functional as F + + img_h, img_w = left.height, left.width + + left_t = torch.from_numpy(np.asarray(left).copy()).permute(2, 0, 1).unsqueeze(0).float() + right_t = torch.from_numpy(np.asarray(right).copy()).permute(2, 0, 1).unsqueeze(0).float() + + sf = cfg.scale_factor + if sf != 1.0: + scaled_h = round(img_h * sf / 32) * 32 + scaled_w = round(img_w * sf / 32) * 32 + left_t_inp = F.interpolate( + left_t, + (scaled_h, scaled_w), + mode="bilinear", + align_corners=False, + ) + right_t_inp = F.interpolate( + right_t, + (scaled_h, scaled_w), + mode="bilinear", + align_corners=False, + ) + else: + 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) + + device_type = device if isinstance(device, str) else device.type + if device_type == "cuda": + autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=True) + elif device_type == "mps": + autocast_ctx = torch.amp.autocast(device_type="mps", dtype=torch.float16, enabled=True) + else: + autocast_ctx = contextlib.nullcontext() + + with torch.inference_mode(), autocast_ctx: + pred_disp, pred_occ, pred_conf = model(left_pad, right_pad) + + if sf != 1.0: + scaled_h = round(img_h * sf / 32) * 32 + scaled_w = round(img_w * sf / 32) * 32 + pred_disp = image_crop(pred_disp, (scaled_h, scaled_w)) + pred_disp = ( + F.interpolate( + pred_disp, + (img_h, img_w), + mode="bilinear", + align_corners=False, + ) + / sf + ) + pred_occ = image_crop(pred_occ, (scaled_h, scaled_w)) + pred_occ = F.interpolate( + pred_occ, + (img_h, img_w), + mode="bilinear", + align_corners=False, + ) + pred_conf = image_crop(pred_conf, (scaled_h, scaled_w)) + pred_conf = F.interpolate( + pred_conf, + (img_h, img_w), + mode="bilinear", + align_corners=False, + ) + else: + pred_disp = image_crop(pred_disp, (img_h, img_w)) + pred_occ = image_crop(pred_occ, (img_h, img_w)) + pred_conf = image_crop(pred_conf, (img_h, img_w)) + + if device_type == "mps": + torch.mps.empty_cache() + + return ( + pred_disp.squeeze().float().cpu().numpy(), + pred_occ.squeeze().float().cpu().numpy(), + pred_conf.squeeze().float().cpu().numpy(), + ) + + # ── Point cloud / mesh ──────────────────────────────────────────────────── + + def _unproject( + self, + left_pil, + depth_m, + disp_np, + cfg, + o3d, + *, + occ_np=None, + conf_np=None, + want_cloud, + want_mesh, + ): + """Unproject the left image into a coloured 3D point cloud and optional mesh. + + Args: + left_pil: Left PIL image (RGB) — used for vertex colours. + depth_m: Metric depth map, shape ``(H, W)``. + disp_np: Raw disparity map, shape ``(H, W)``. + cfg: :class:`~vizion3d.stereo.models.StereoDepthAdvancedConfig`. + o3d: Imported open3d module. + occ_np: Optional per-pixel occlusion score, shape ``(H, W)``. + conf_np: Optional per-pixel confidence score, shape ``(H, W)``. + want_cloud: Whether to return the point cloud. + want_mesh: Whether to return the mesh. + + Returns: + ``(point_cloud, mesh)`` — each is the Open3D object or ``None`` if not requested. + """ + H, W = depth_m.shape + left_np = np.asarray(left_pil) + + uu, vv = np.meshgrid(np.arange(W, dtype=np.float32), np.arange(H, dtype=np.float32)) + x = (uu - cfg.cx) * depth_m / cfg.focal_length + y = (vv - cfg.cy) * depth_m / cfg.focal_length + z = depth_m + + valid = (z > 0) & (z < cfg.z_far) + if conf_np is not None: + valid &= conf_np >= cfg.conf_threshold + if occ_np is not None: + valid &= occ_np >= cfg.occ_threshold + pts = np.stack([x, y, z], axis=-1)[valid] + cols = left_np[valid].astype(np.float64) / 255.0 + + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pts) + pcd.colors = o3d.utility.Vector3dVector(cols) + + # Flip Y so Y increases downward in the viewer, matching image-plane orientation. + pcd.transform(_CAMERA_TO_IMAGE_VIEW) + + mesh = None + if want_mesh: + mesh = DepthEstimationHandler._mesh_from_point_cloud(pcd, o3d) + + return (pcd if want_cloud else None), mesh + + # ── Pre-loading ─────────────────────────────────────────────────────────── + + @classmethod + def preload(cls, model_path: str) -> None: + """Resolve *model_path* (downloading if a URL) and load it into the class-level cache. + + Call this at server startup to ensure the model is in memory before the first request. + """ + from .defaults import resolve_stereo_model_backend + + resolved = resolve_stereo_model_backend(model_path) + cls()._load_s2m2(resolved) + + # ── Device selection ────────────────────────────────────────────────────── + + @staticmethod + def _torch_device(torch_module) -> str: + """Return the best available device string: ``'cuda'``, ``'mps'``, or ``'cpu'``.""" + if torch_module.cuda.is_available(): + return "cuda" + if hasattr(torch_module.backends, "mps") and torch_module.backends.mps.is_available(): + return "mps" + return "cpu" diff --git a/vizion3d/stereo/models.py b/vizion3d/stereo/models.py new file mode 100644 index 0000000..41c91e9 --- /dev/null +++ b/vizion3d/stereo/models.py @@ -0,0 +1,96 @@ +""" +Data models for the Stereo Depth task. + +Defines the camera-configuration Pydantic model and the result payload. +""" + +from open3d.geometry import Image as O3dImage # type: ignore[import-untyped] +from open3d.geometry import PointCloud as O3dPointCloud # type: ignore[import-untyped] +from open3d.geometry import TriangleMesh as O3dTriangleMesh # type: ignore[import-untyped] +from pydantic import BaseModel, ConfigDict + + +class StereoDepthAdvancedConfig(BaseModel): + """Camera intrinsics and inference settings for stereo depth estimation. + + All fields are optional overrides — unspecified fields retain sensible defaults + for a 1280×720 commodity stereo camera with a 100 mm baseline. + + Attributes: + focal_length: Focal length in pixels (assumes square pixels, i.e. fx = fy). + Larger values mean a narrower field of view and more perspective + compression. Override with your camera's actual calibrated value. + cx: Principal point x — the pixel column of the optical axis, typically + near the horizontal image centre (``image_width / 2 - 0.5``). + cy: Principal point y — the pixel row of the optical axis, typically near + the vertical image centre (``image_height / 2 - 0.5``). + baseline: Stereo baseline in **millimetres** — the physical distance between + the two camera optical centres. Real metric depth is proportional to + this value, so an incorrect baseline scales all depth values uniformly. + doffs: Disparity offset in pixels. Non-zero for Middlebury-style calibration + where the principal points are not aligned across the two views. + Set to ``0.0`` for standard rectified pairs. + z_far: Maximum depth in metres. Points beyond this distance are excluded + from the point cloud and mesh to reduce noise and file size. + conf_threshold: Minimum per-pixel confidence score (in ``[0, 1]``) for a + point to be included in the point cloud. Lower values include more + uncertain points; higher values give sparser but more reliable clouds. + occ_threshold: Minimum occlusion score (in ``[0, 1]``) for a point to be + included. Points with low occlusion scores are likely partially occluded + in one view and produce less reliable depth estimates. + scale_factor: Input image downscale factor before inference. ``1.0`` means + full resolution (highest quality, slowest). ``0.5`` halves both spatial + dimensions (~3–4× faster at some quality cost). + """ + + focal_length: float = 1000.0 + cx: float = 640.0 + cy: float = 360.0 + baseline: float = 100.0 + doffs: float = 0.0 + z_far: float = 10.0 + conf_threshold: float = 0.1 + occ_threshold: float = 0.5 + scale_factor: float = 1.0 + + +class StereoDepthResult(BaseModel): + """Result payload returned after a stereo depth inference task. + + Attributes: + depth_map: Metric depth in **metres**, shape ``[H][W]``. Unlike monocular + depth estimation, these are real-world distances (assuming correct camera + calibration) — not relative or fictitious values. + disparity_map: Raw disparity map in pixels, shape ``[H][W]``. Disparity + is the horizontal pixel offset between matched features across the + left and right images. Depth = baseline × focal_length / disparity. + min_depth: Minimum value in ``depth_map`` (metres). + max_depth: Maximum value in ``depth_map`` (metres). Guaranteed + ``max_depth >= min_depth``. + backend_used: Resolved local file path of the checkpoint used. + depth_image: 16-bit grayscale ``open3d.geometry.Image`` (dtype ``uint16``) + where the full 0–65535 range maps linearly to ``[min_depth, max_depth]``. + Present when ``return_depth_image=True`` was set on the command. + point_cloud: Coloured ``open3d.geometry.PointCloud`` unprojected from the + RGB-D image using the camera intrinsics in ``advanced_config``. + Coordinates are in metres. Present when ``return_point_cloud=True``. + mesh: ``open3d.geometry.TriangleMesh`` reconstructed via ball-pivoting + from the point cloud. Includes vertex colours. + Present when ``return_mesh=True``. + point_cloud_scale: Scale factor: multiply any distance measured between + two points in the returned point cloud by this value to get the + equivalent distance in metres. Always ``1.0`` for stereo depth — + coordinates are already in real metric units. + """ + + depth_map: list[list[float]] + disparity_map: list[list[float]] + min_depth: float + max_depth: float + backend_used: str + depth_image: O3dImage | None = None + point_cloud: O3dPointCloud | None = None + mesh: O3dTriangleMesh | None = None + point_cloud_scale: float = 1.0 + + model_config = ConfigDict(arbitrary_types_allowed=True)