From dec5412316c8ef32c2bcfc363b8453a7cee7c039 Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Mon, 31 Aug 2026 12:02:16 -0400 Subject: [PATCH 1/9] Add a foundation-stereo option for warping detections between stereo cameras The existing client-side transfer matches each point by NCC template matching along its epipolar curve, which needs the source patch to be photometrically matchable in the other view. Where that fails -- obstructed viewpoints, repetitive substrate, low contrast -- it mismatches or declines. Adds a second correspondence method that runs a Fast-FoundationStereo ONNX export instead: rectify the pair once per frame, run a dense disparity network, and read each point's shift out of the disparity map. Both methods satisfy a shared StereoMatcher interface, so box/line/polygon warping, measurement and bulk transfer are unchanged either way, and a Correspondence method dropdown under Track Settings -> Stereo Settings picks between them. The foundation model is not bundled -- the exports run ~100 MB -- so it is supplied by URL and the method no-ops with a warning when none is served. NCC remains the default. --- .../components/TrackSettingsPanel.vue | 39 ++++ client/dive-common/store/settings.ts | 6 + client/dive-common/use/stereo/README.md | 88 +++++++- .../use/stereo/StereoFoundationMatcher.ts | 201 ++++++++++++++++++ client/dive-common/use/stereo/index.ts | 10 + client/dive-common/use/stereo/rectify.ts | 166 +++++++++++++++ .../dive-common/use/stereo/stereoMatcher.ts | 32 +++ .../use/stereo/tests/rectify.spec.ts | 122 +++++++++++ .../use/stereo/useStereoOnnxTransfer.ts | 11 +- .../platform/web-girder/useStereoOnnxWeb.ts | 53 +++-- 10 files changed, 710 insertions(+), 18 deletions(-) create mode 100644 client/dive-common/use/stereo/StereoFoundationMatcher.ts create mode 100644 client/dive-common/use/stereo/rectify.ts create mode 100644 client/dive-common/use/stereo/stereoMatcher.ts create mode 100644 client/dive-common/use/stereo/tests/rectify.spec.ts diff --git a/client/dive-common/components/TrackSettingsPanel.vue b/client/dive-common/components/TrackSettingsPanel.vue index 8144c3aea..a1ac15d86 100644 --- a/client/dive-common/components/TrackSettingsPanel.vue +++ b/client/dive-common/components/TrackSettingsPanel.vue @@ -7,6 +7,7 @@ import { computed, } from 'vue'; import { clientSettings } from 'dive-common/store/settings'; +import { STEREO_MATCH_METHODS } from 'dive-common/use/stereo/stereoMatcher'; import isDesktopRuntime from 'dive-common/isDesktopRuntime'; export default defineComponent({ @@ -40,6 +41,7 @@ export default defineComponent({ showMultiCamToolbar: 'Show multi-camera tools in the top toolbar when a track is selected', stereoUpdateLengths: 'When a line annotation is modified on a detection that is linked across both cameras, recompute its stereo measurement (length, midpoint, range, RMS) automatically.', stereoAutoCompute: 'When an annotation is drawn on one camera and the other camera has no detection for it yet, automatically warp it to the other camera using stereo disparity.', + stereoMatchMethod: 'How a point is matched on the other camera. Template matching correlates the source patch along the epipolar line. Foundation stereo runs a dense disparity network over the pair once, which is steadier where the patch is hard to match but needs its model to be configured.', }); const modes = ref(['Track', 'Detection']); // Add unknown as the default type to the typeList @@ -52,6 +54,7 @@ export default defineComponent({ help, modes, typeList, + stereoMatchMethods: STEREO_MATCH_METHODS, }; }, }); @@ -445,6 +448,42 @@ export default defineComponent({ + + + + + + + + {{ help.stereoMatchMethod }} + + + diff --git a/client/dive-common/store/settings.ts b/client/dive-common/store/settings.ts index c5d145b42..fdffc69a8 100644 --- a/client/dive-common/store/settings.ts +++ b/client/dive-common/store/settings.ts @@ -2,6 +2,8 @@ import { Ref, watch, reactive } from 'vue'; import { cloneDeep, merge } from 'lodash'; import { AnnotatorPreferences } from 'vue-media-annotator/types'; import isDesktopRuntime from 'dive-common/isDesktopRuntime'; +import { DEFAULT_STEREO_MATCH_METHOD } from 'dive-common/use/stereo/stereoMatcher'; +import type { StereoMatchMethod } from 'dive-common/use/stereo/stereoMatcher'; interface ColumnVisibilitySettings { type: boolean; @@ -94,6 +96,9 @@ interface AnnotationSettings { // Warp an annotation drawn on one camera to the other camera when that // camera has no detection for it yet. autoComputeOtherCamera: boolean; + // Which correspondence method the warp uses: 'ncc' template matching or + // 'foundation' dense disparity. + matchMethod: StereoMatchMethod; loading: boolean; loadingMessage: string; }; @@ -191,6 +196,7 @@ const defaultSettings: AnnotationSettings = { clearLengthOnCameraFileLoad: true, updateLengthsOnModify: true, autoComputeOtherCamera: false, + matchMethod: DEFAULT_STEREO_MATCH_METHOD, loading: false, loadingMessage: '', }, diff --git a/client/dive-common/use/stereo/README.md b/client/dive-common/use/stereo/README.md index 84a866588..387a8489a 100644 --- a/client/dive-common/use/stereo/README.md +++ b/client/dive-common/use/stereo/README.md @@ -1,9 +1,20 @@ # Client-side stereo transfer and measurement (ONNX) Warp a detection annotated on one camera onto the other camera and measure its -length, entirely in the browser / Electron renderer — no backend — using VIAME's -epipolar template-matching model (stereo measurement "method 1") exported to -ONNX and run with `onnxruntime-web`. +length, entirely in the browser / Electron renderer — no backend — running the +correspondence model with `onnxruntime-web`. + +Two correspondence methods are available, chosen from **Track Settings → Stereo +Settings → Correspondence method**: + +| Method | Model | How it matches | +| --- | --- | --- | +| **Template matching (NCC)** — default | VIAME's epipolar template-matching model (stereo measurement "method 1"), bundled | Per point: generate epipolar candidates, NCC the source patch along that curve | +| **Foundation stereo (disparity)** | A Fast-FoundationStereo ONNX export, **not bundled** | Once per frame: rectify the pair, run a dense disparity network, read each point's shift out of the map | + +They are interchangeable behind the `StereoMatcher` interface, so everything +downstream — box/line/polygon warping, measurement, bulk transfer — is identical +either way. This is the client counterpart to the desktop backend stereo service: the desktop `ViewerLoader` warps and measures via native IPC (`stereoTransferLine` / @@ -15,6 +26,9 @@ work client-side so it also works on the web. | File | Role | | --- | --- | | `StereoOnnxMatcher.ts` | Loads the `match` ONNX model and warps source points → target points via NCC along the epipolar curve. | +| `StereoFoundationMatcher.ts` | Loads a Fast-FoundationStereo ONNX export, rectifies the pair into the network's input resolution, and reads each point's correspondence from the dense disparity map. | +| `stereoMatcher.ts` | The `StereoMatcher` contract both matchers satisfy, the `StereoMatchMethod` union, and the dropdown's labels. | +| `rectify.ts` | Stereo rectification ported from OpenCV `cvStereoRectify` (Rodrigues, rectifying rotations, point rectify/unrectify, and the inverse map used to sample a rectified image). Only the foundation method needs it. | | `calibration.ts` | `StereoRig` + loaders (`rigFromNpz`, `rigFromJson`) mirroring VIAME's `read_stereo_rig`; `invertRig` to swap the source/target camera. | | `npz.ts` | Minimal `.npz`/`.npy` reader (calibration files are NumPy archives). | | `image.ts` | RGBA → BT.601 grayscale (matches OpenCV `BGR2GRAY` used by the C++ NCC). | @@ -107,3 +121,71 @@ The disparity range is scene-dependent — VIAME's batch measurement pipes ship binding, calibration download, and the GeoJS frame-pixel read in `frameSource.geoViewerToImageElement`) is type-checked and lint-clean but has not been exercised in a running web viewer with a real stereo dataset. + + +## Foundation stereo method + +### Why a second method + +The NCC matcher needs the source patch to be photometrically matchable in the +other view. Where that fails — obstructed viewpoints, repetitive substrate, low +contrast — it either mismatches or declines. A dense disparity network does not +depend on patch correlation, and it costs one network pass per frame no matter +how many points are warped, so bulk-warping a whole camera amortises well. + +Its trade is setup: the model is large and must be supplied. + +### Supplying the model + +Unlike the NCC graph (small, committed at `client/public/models/stereo_match.onnx`), +Fast-FoundationStereo exports run ~100 MB and are **not** committed. Obtain an +export from the Fast-FoundationStereo release, serve it, and point the web glue +at it: + +```ts +useStereoOnnxWeb({ + ..., + foundationModelUrl: '/models/stereo_foundation.onnx', + foundationModelSpec: { height: 576, width: 960 }, // the export's sidecar image_size +}); +``` + +The default URL is `/models/stereo_foundation.onnx` and the default spec is +576×960. `foundationModelSpec` **must** match the export: the graph fixes its +input resolution, and the sidecar `.yaml` shipped beside each export gives it as +`image_size: [H, W]`. With no model served, selecting the method reports that it +could not load and the warp no-ops — the same way a missing calibration does. + +### How it works + +1. Solve the rectifying rotations for the rig once per calibration + (`computeRectification`), sized to the network's input resolution. +2. Build the rectified pair by inverse-mapping each output pixel back to its + source pixel and bilinear-sampling. Rectify and resize are fused, so the cost + is the network's resolution rather than the frame's. +3. Run the network to get dense disparity in rectified pixels. +4. Per point: rectify it, pool the disparities in a small window by median, + shift `x` by that disparity, and unrectify into the target image. + +Step 4 pools rather than sampling the single pixel deliberately. A head or tail +tip is a couple of pixels wide at the network's working resolution, so the +disparity exactly at the tip is frequently the background's; the median over a +small window rejects that without dragging the estimate off the animal. + +The network emits no confidence channel, so the reported `score` is the fraction +of the pooled window carrying a finite positive disparity, and a match is +accepted when that clears `DEFAULT_MIN_VALID_FRACTION` **and** the implied +disparity falls inside the configured search range — the same range that bounds +the NCC search. + +### Testing status + +- **Tested** (`tests/rectify.spec.ts`): Rodrigues round-trip, orthonormality of + the rectifying rotations, the defining rectification property (a 3D point + lands on the same row in both rectified views), disparity positive and + decreasing with range, and pixel round-trip through rectify/unrectify with and + without distortion. +- **Not tested**: `StereoFoundationMatcher` end-to-end, which needs a ~100 MB + model the repo does not carry. The geometry it depends on is covered above; + the network call, disparity pooling and the settings dropdown have not been + exercised against a real export in a running viewer. diff --git a/client/dive-common/use/stereo/StereoFoundationMatcher.ts b/client/dive-common/use/stereo/StereoFoundationMatcher.ts new file mode 100644 index 000000000..a0250d697 --- /dev/null +++ b/client/dive-common/use/stereo/StereoFoundationMatcher.ts @@ -0,0 +1,201 @@ +/** + * Client-side wrapper around a Fast-FoundationStereo ONNX export (NVIDIA), as + * the second stereo correspondence method alongside {@link StereoOnnxMatcher}. + * + * Where the NCC matcher searches the epipolar curve per point, this one runs a + * dense disparity network over the whole pair once and reads each point's + * correspondence out of the disparity map. That costs one network pass per + * frame regardless of how many points are warped, and it does not depend on the + * source patch being photometrically matchable — which is what makes it hold up + * on footage where template correlation struggles (obstructed views, repetitive + * substrate, low contrast). + * + * The model is NOT bundled: the exports are ~100 MB, far past what belongs in + * the repo. Point {@link StereoFoundationMatcher.create} at a served or + * user-supplied model. Exports are published with the Fast-FoundationStereo + * release as `_iters__res_x.onnx` plus a sidecar `.yaml` giving + * `image_size`; the graph takes `left_image`/`right_image` as [1,3,H,W] RGB in + * [0,1] and returns `disparity` as [1,1,H,W] in rectified pixels. + */ + +import * as ort from 'onnxruntime-web'; + +import { GrayImage } from './image'; +import { StereoRig } from './calibration'; +import type { WarpOptions, WarpResult } from './StereoOnnxMatcher'; +import { + Rectification, computeRectification, rectifyPoint, rectifyMapper, unrectifyPoint, +} from './rectify'; + +/** + * Half-width of the window whose disparities are pooled for one point. + * + * A head or tail tip is a couple of pixels wide at the network's working + * resolution, so the disparity sampled exactly at the tip is often the + * background's. Pooling a small neighbourhood by median rejects that without + * dragging the estimate off the animal. + */ +export const DEFAULT_SAMPLE_RADIUS = 3; + +/** + * Fraction of the pooled window that must carry a finite positive disparity for + * the match to be accepted. The network emits a dense map with no confidence + * channel, so validity density is the available proxy. + */ +export const DEFAULT_MIN_VALID_FRACTION = 0.34; + +export interface FoundationModelSpec { + /** Network input size, from the export's sidecar yaml `image_size: [H, W]`. */ + height: number; + width: number; +} + +/** Bilinear sample of a single-channel image, NaN outside. */ +function sampleBilinear(data: Float32Array, width: number, height: number, x: number, y: number): number { + if (!(x >= 0 && y >= 0 && x <= width - 1 && y <= height - 1)) return NaN; + const x0 = Math.floor(x); + const y0 = Math.floor(y); + const x1 = Math.min(x0 + 1, width - 1); + const y1 = Math.min(y0 + 1, height - 1); + const fx = x - x0; + const fy = y - y0; + const a = data[y0 * width + x0]; + const b = data[y0 * width + x1]; + const c = data[y1 * width + x0]; + const d = data[y1 * width + x1]; + return a * (1 - fx) * (1 - fy) + b * fx * (1 - fy) + c * (1 - fx) * fy + d * fx * fy; +} + +/** Remap a grayscale frame through an inverse map into an RGB [1,3,H,W] tensor. */ +function remapToRgbTensor(src: GrayImage, mapX: Float32Array, mapY: Float32Array, width: number, height: number): ort.Tensor { + const plane = width * height; + const out = new Float32Array(plane * 3); + for (let i = 0; i < plane; i += 1) { + const v = sampleBilinear(src.data, src.width, src.height, mapX[i], mapY[i]); + const g = Number.isNaN(v) ? 0 : v; + out[i] = g; + out[plane + i] = g; + out[2 * plane + i] = g; + } + return new ort.Tensor('float32', out, [1, 3, height, width]); +} + +export class StereoFoundationMatcher { + private session: ort.InferenceSession; + + private spec: FoundationModelSpec; + + /** Rectification + inverse maps, rebuilt only when the rig or size changes. */ + private cache: { + key: string; rect: Rectification; + src: { mapX: Float32Array; mapY: Float32Array }; + tgt: { mapX: Float32Array; mapY: Float32Array }; + } | null = null; + + private constructor(session: ort.InferenceSession, spec: FoundationModelSpec) { + this.session = session; + this.spec = spec; + } + + /** + * Create a matcher from a model URL or model bytes. `spec` is the export's + * input resolution (its sidecar yaml `image_size`), which the graph fixes. + */ + static async create( + model: string | ArrayBuffer | Uint8Array, + spec: FoundationModelSpec, + opts: { threads?: number } = {}, + ): Promise { + ort.env.wasm.numThreads = opts.threads ?? 1; + ort.env.wasm.proxy = false; + const session = await ort.InferenceSession.create(model as string, { + executionProviders: ['wasm'], + graphOptimizationLevel: 'all', + }); + return new StereoFoundationMatcher(session, spec); + } + + /** Rectification and inverse maps for this rig at the model's resolution. */ + private geometry(rig: StereoRig) { + const key = `${rig.Kl.join(',')}|${rig.R.join(',')}|${rig.T.join(',')}`; + if (this.cache && this.cache.key === key) return this.cache; + const rect = computeRectification(rig, this.spec.width, this.spec.height); + this.cache = { + key, + rect, + src: rectifyMapper(rig, rect, false), + tgt: rectifyMapper(rig, rect, true), + }; + return this.cache; + } + + /** + * Warp source-image points onto the target image, matching + * {@link StereoOnnxMatcher.warpPoints} so the two are interchangeable. + * + * `opts.range` bounds the accepted disparity exactly as it bounds the NCC + * search: a correspondence outside it is rejected rather than trusted. + */ + async warpPoints( + points: [number, number][], + source: GrayImage, + target: GrayImage, + rig: StereoRig, + opts: WarpOptions, + ): Promise { + const { rect, src, tgt } = this.geometry(rig); + const { width, height } = this.spec; + + const feeds: Record = { + left_image: remapToRgbTensor(source, src.mapX, src.mapY, width, height), + right_image: remapToRgbTensor(target, tgt.mapX, tgt.mapY, width, height), + }; + const out = await this.session.run(feeds); + const disparity = out.disparity.data as Float32Array; + + const radius = DEFAULT_SAMPLE_RADIUS; + const minValid = DEFAULT_MIN_VALID_FRACTION; + const [minDisp, maxDisp] = 'minDisparity' in opts.range + ? [opts.range.minDisparity, opts.range.maxDisparity] + : [0, Number.POSITIVE_INFINITY]; + // The search range is expressed in source-image pixels; the network works + // at its own resolution, so carry the bound across in the same ratio. + const dispScale = width / source.width; + + return points.map(([px, py]) => { + const [rx, ry] = rectifyPoint(px, py, rig, rect, false); + const fail: WarpResult = { + x: NaN, y: NaN, score: 0, secondScore: 0, accepted: false, + }; + if (!Number.isFinite(rx) || !Number.isFinite(ry)) return fail; + + const samples: number[] = []; + let considered = 0; + for (let dy = -radius; dy <= radius; dy += 1) { + for (let dx = -radius; dx <= radius; dx += 1) { + considered += 1; + const v = sampleBilinear(disparity, width, height, rx + dx, ry + dy); + if (Number.isFinite(v) && v > 0) samples.push(v); + } + } + if (!samples.length) return fail; + + samples.sort((a, b) => a - b); + const d = samples[Math.floor(samples.length / 2)]; + const validFraction = samples.length / considered; + + const dSource = d / dispScale; + const inRange = dSource >= minDisp && dSource <= maxDisp; + const [ox, oy] = unrectifyPoint(rx - d, ry, rig, rect, true); + if (!Number.isFinite(ox) || !Number.isFinite(oy)) return fail; + + return { + x: ox, + y: oy, + score: validFraction, + secondScore: 0, + accepted: validFraction >= minValid && inRange, + }; + }); + } +} diff --git a/client/dive-common/use/stereo/index.ts b/client/dive-common/use/stereo/index.ts index cc5ad56bd..01db38c26 100644 --- a/client/dive-common/use/stereo/index.ts +++ b/client/dive-common/use/stereo/index.ts @@ -1,4 +1,14 @@ export { StereoOnnxMatcher } from './StereoOnnxMatcher'; +export { StereoFoundationMatcher } from './StereoFoundationMatcher'; +export type { FoundationModelSpec } from './StereoFoundationMatcher'; +export { + DEFAULT_STEREO_MATCH_METHOD, STEREO_MATCH_METHODS, +} from './stereoMatcher'; +export type { StereoMatcher, StereoMatchMethod } from './stereoMatcher'; +export { + computeRectification, rectifyPoint, unrectifyPoint, rectifyMapper, +} from './rectify'; +export type { Rectification } from './rectify'; export type { WarpOptions, WarpResult, SearchRange } from './StereoOnnxMatcher'; export { rigFromNpz, rigFromNpzArrays, rigFromJson, baseline, diff --git a/client/dive-common/use/stereo/rectify.ts b/client/dive-common/use/stereo/rectify.ts new file mode 100644 index 000000000..79cad859d --- /dev/null +++ b/client/dive-common/use/stereo/rectify.ts @@ -0,0 +1,166 @@ +/** + * Stereo rectification, ported from OpenCV's `cvStereoRectify` (Bouguet). + * + * The NCC matcher searches the epipolar curve directly and needs none of this. + * A disparity network does: it consumes a rectified pair, where corresponding + * points share a row and the correspondence is a pure horizontal shift. + * + * Nothing here builds a full-resolution rectified image. `rectifyMapper` + * returns the inverse map (rectified pixel -> source pixel) so a caller can + * sample straight into the network's input resolution, fusing rectify+resize + * into one bilinear read per output pixel. + */ + +import { StereoRig } from './calibration'; +import { mapPoint, unmap } from './triangulate'; + +export type Mat3 = Float32Array; +type Vec3 = [number, number, number]; + +export interface Rectification { + /** Rectifying rotations for the source and target cameras. */ + R1: Mat3; + R2: Mat3; + /** Shared focal length and principal point of the rectified pair. */ + f: number; + cx: number; + cy: number; + /** Rectified image size these were solved for. */ + width: number; + height: number; +} + +function matMul(a: ArrayLike, b: ArrayLike): Mat3 { + const m = new Float32Array(9); + for (let r = 0; r < 3; r += 1) { + for (let c = 0; c < 3; c += 1) { + m[r * 3 + c] = a[r * 3] * b[c] + a[r * 3 + 1] * b[3 + c] + a[r * 3 + 2] * b[6 + c]; + } + } + return m; +} + +function transpose(a: ArrayLike): Mat3 { + return Float32Array.from([a[0], a[3], a[6], a[1], a[4], a[7], a[2], a[5], a[8]]); +} + +function matVec(a: ArrayLike, v: ArrayLike): Vec3 { + return [ + a[0] * v[0] + a[1] * v[1] + a[2] * v[2], + a[3] * v[0] + a[4] * v[1] + a[5] * v[2], + a[6] * v[0] + a[7] * v[1] + a[8] * v[2], + ]; +} + +/** Rotation matrix -> rotation vector (axis * angle). */ +export function rodriguesInv(R: ArrayLike): Vec3 { + const trace = R[0] + R[4] + R[8]; + const cos = Math.min(1, Math.max(-1, (trace - 1) / 2)); + const angle = Math.acos(cos); + if (angle < 1e-9) return [0, 0, 0]; + const s = angle / (2 * Math.sin(angle)); + return [s * (R[7] - R[5]), s * (R[2] - R[6]), s * (R[3] - R[1])]; +} + +/** Rotation vector -> rotation matrix. */ +export function rodrigues(v: ArrayLike): Mat3 { + const theta = Math.sqrt(v[0] * v[0] + v[1] * v[1] + v[2] * v[2]); + if (theta < 1e-9) return Float32Array.from([1, 0, 0, 0, 1, 0, 0, 0, 1]); + const [x, y, z] = [v[0] / theta, v[1] / theta, v[2] / theta]; + const c = Math.cos(theta); + const s = Math.sin(theta); + const t = 1 - c; + return Float32Array.from([ + t * x * x + c, t * x * y - s * z, t * x * z + s * y, + t * x * y + s * z, t * y * y + c, t * y * z - s * x, + t * x * z - s * y, t * y * z + s * x, t * z * z + c, + ]); +} + +/** + * Solve the rectifying rotations for a rig. + * + * The focal length is the smaller of the two cameras' so the rectified frustum + * stays inside both, and the principal point is centred on the output. This is + * OpenCV's `alpha = 0`-free behaviour: no zoom-to-valid-region crop, which on a + * rig whose baseline sits far from horizontal demands an extreme zoom and can + * push the whole scene off canvas. + */ +export function computeRectification(rig: StereoRig, width: number, height: number): Rectification { + // Half-rotate both cameras toward each other: r = R^(-1/2). + const om = rodriguesInv(rig.R); + const r = rodrigues([-om[0] / 2, -om[1] / 2, -om[2] / 2]); + const t = matVec(r, rig.T); + + // New x axis along the (half-rotated) baseline. + const nt = Math.hypot(t[0], t[1], t[2]) || 1; + const horizontal = Math.abs(t[0]) > Math.abs(t[1]); + const idx = horizontal ? 0 : 1; + const uu: Vec3 = [0, 0, 0]; + uu[idx] = t[idx] > 0 ? 1 : -1; + + // Rotate about the axis that carries the baseline onto uu. + const ww: Vec3 = [ + t[1] * uu[2] - t[2] * uu[1], + t[2] * uu[0] - t[0] * uu[2], + t[0] * uu[1] - t[1] * uu[0], + ]; + const nw = Math.hypot(ww[0], ww[1], ww[2]); + let wR: Mat3; + if (nw < 1e-12) { + wR = Float32Array.from([1, 0, 0, 0, 1, 0, 0, 0, 1]); + } else { + const scale = Math.acos(Math.abs(t[idx]) / nt) / nw; + wR = rodrigues([ww[0] * scale, ww[1] * scale, ww[2] * scale]); + } + + return { + R1: matMul(wR, transpose(r)), + R2: matMul(wR, r), + f: Math.min(rig.Kl[0], rig.Kr[0]), + cx: (width - 1) / 2, + cy: (height - 1) / 2, + width, + height, + }; +} + +/** Source pixel -> rectified pixel, for the source (R1) or target (R2) camera. */ +export function rectifyPoint(px: number, py: number, rig: StereoRig, rect: Rectification, target: boolean): [number, number] { + const K = target ? rig.Kr : rig.Kl; + const d = target ? rig.distr : rig.distl; + const R = target ? rect.R2 : rect.R1; + const [nx, ny] = unmap(px, py, K, d); + const p = matVec(R, [nx, ny, 1]); + if (p[2] === 0) return [NaN, NaN]; + return [rect.f * (p[0] / p[2]) + rect.cx, rect.f * (p[1] / p[2]) + rect.cy]; +} + +/** Rectified pixel -> source pixel (the inverse of {@link rectifyPoint}). */ +export function unrectifyPoint(rx: number, ry: number, rig: StereoRig, rect: Rectification, target: boolean): [number, number] { + const K = target ? rig.Kr : rig.Kl; + const d = target ? rig.distr : rig.distl; + const R = target ? rect.R2 : rect.R1; + const p = matVec(transpose(R), [(rx - rect.cx) / rect.f, (ry - rect.cy) / rect.f, 1]); + if (p[2] === 0) return [NaN, NaN]; + return mapPoint(p[0] / p[2], p[1] / p[2], K, d); +} + +/** + * Inverse map for building a rectified image: for each rectified pixel, the + * source pixel to sample. Returned as flat x/y arrays of length width*height so + * the caller can bilinear-sample without recomputing the projection per frame. + */ +export function rectifyMapper(rig: StereoRig, rect: Rectification, target: boolean): { mapX: Float32Array; mapY: Float32Array } { + const { width, height } = rect; + const mapX = new Float32Array(width * height); + const mapY = new Float32Array(width * height); + for (let y = 0; y < height; y += 1) { + for (let x = 0; x < width; x += 1) { + const [sx, sy] = unrectifyPoint(x, y, rig, rect, target); + mapX[y * width + x] = sx; + mapY[y * width + x] = sy; + } + } + return { mapX, mapY }; +} diff --git a/client/dive-common/use/stereo/stereoMatcher.ts b/client/dive-common/use/stereo/stereoMatcher.ts new file mode 100644 index 000000000..be22ac7d7 --- /dev/null +++ b/client/dive-common/use/stereo/stereoMatcher.ts @@ -0,0 +1,32 @@ +/** + * The contract both correspondence methods satisfy, so the transfer composable + * and its callers never branch on which one is selected. + */ + +import { GrayImage } from './image'; +import { StereoRig } from './calibration'; +import type { WarpOptions, WarpResult } from './StereoOnnxMatcher'; + +/** + * `ncc` — epipolar candidates + NCC template matching (VIAME method 1). + * `foundation` — dense Fast-FoundationStereo disparity, read per point. + */ +export type StereoMatchMethod = 'ncc' | 'foundation'; + +export const DEFAULT_STEREO_MATCH_METHOD: StereoMatchMethod = 'ncc'; + +export interface StereoMatcher { + warpPoints( + points: [number, number][], + source: GrayImage, + target: GrayImage, + rig: StereoRig, + opts: WarpOptions, + ): Promise; +} + +/** Labels for the method selector. */ +export const STEREO_MATCH_METHODS: { value: StereoMatchMethod; text: string }[] = [ + { value: 'ncc', text: 'Template matching (NCC)' }, + { value: 'foundation', text: 'Foundation stereo (disparity)' }, +]; diff --git a/client/dive-common/use/stereo/tests/rectify.spec.ts b/client/dive-common/use/stereo/tests/rectify.spec.ts new file mode 100644 index 000000000..3daded55c --- /dev/null +++ b/client/dive-common/use/stereo/tests/rectify.spec.ts @@ -0,0 +1,122 @@ +/** + * Rectification is the only new geometry the foundation matcher adds, and it is + * the part that silently produces plausible-but-wrong warps if it is off. These + * check it against properties that must hold for any correct rectification, + * rather than against a golden matrix. + */ + +import { describe, it, expect } from 'vitest'; + +import { + computeRectification, rectifyPoint, unrectifyPoint, rodrigues, rodriguesInv, +} from '../rectify'; +import { StereoRig } from '../calibration'; +import { project } from '../triangulate'; + +const I3 = Float32Array.from([1, 0, 0, 0, 1, 0, 0, 0, 1]); +const Z3 = Float32Array.from([0, 0, 0]); + +/** Project a world point (left-camera frame) into one of the rig's cameras. */ +function projectInto(p: [number, number, number], rig: StereoRig, target: boolean): [number, number] { + return target + ? project(p, rig.Kr, rig.distr, rig.R, rig.T) + : project(p, rig.Kl, rig.distl, I3, Z3); +} + +const W = 960; +const H = 576; + +/** A rig with a mostly-horizontal baseline and a small relative rotation. */ +function makeRig(rotation: [number, number, number] = [0.01, -0.02, 0.004]): StereoRig { + const K = Float32Array.from([1000, 0, 640, 0, 1000, 400, 0, 0, 1]); + return { + Kl: K, + Kr: Float32Array.from(K), + distl: new Float32Array(8), + distr: new Float32Array(8), + R: rodrigues(rotation), + T: Float32Array.from([-200, -30, -5]), + }; +} + +describe('rodrigues', () => { + it('round-trips a rotation vector through the matrix form', () => { + const v: [number, number, number] = [0.11, -0.24, 0.07]; + const back = rodriguesInv(rodrigues(v)); + back.forEach((c, i) => expect(c).toBeCloseTo(v[i], 6)); + }); + + it('returns identity for a zero rotation', () => { + const R = rodrigues([0, 0, 0]); + expect(Array.from(R)).toEqual([1, 0, 0, 0, 1, 0, 0, 0, 1]); + }); +}); + +describe('computeRectification', () => { + it('produces orthonormal rectifying rotations', () => { + const { R1, R2 } = computeRectification(makeRig(), W, H); + [R1, R2].forEach((R) => { + for (let i = 0; i < 3; i += 1) { + for (let j = 0; j < 3; j += 1) { + const dot = R[i * 3] * R[j * 3] + R[i * 3 + 1] * R[j * 3 + 1] + R[i * 3 + 2] * R[j * 3 + 2]; + expect(dot).toBeCloseTo(i === j ? 1 : 0, 5); + } + } + }); + }); + + it('puts corresponding points on the same rectified row', () => { + // The defining property of rectification: a 3D point seen by both cameras + // must land on one row, so disparity is a pure horizontal shift. + const rig = makeRig(); + const rect = computeRectification(rig, W, H); + const worldPoints: [number, number, number][] = [ + [0, 0, 3000], [400, -200, 2500], [-350, 250, 4000], [120, 90, 1800], + ]; + worldPoints.forEach((p) => { + const left = projectInto(p, rig, false); + const right = projectInto(p, rig, true); + const [, ly] = rectifyPoint(left[0], left[1], rig, rect, false); + const [, ry] = rectifyPoint(right[0], right[1], rig, rect, true); + expect(ry).toBeCloseTo(ly, 2); + }); + }); + + it('gives a positive disparity that shrinks with range', () => { + const rig = makeRig(); + const rect = computeRectification(rig, W, H); + const disparityAt = (z: number) => { + const p: [number, number, number] = [0, 0, z]; + const [lx] = rectifyPoint(...projectInto(p, rig, false), rig, rect, false); + const [rx] = rectifyPoint(...projectInto(p, rig, true), rig, rect, true); + return lx - rx; + }; + const near = disparityAt(1500); + const far = disparityAt(6000); + expect(near).toBeGreaterThan(0); + expect(far).toBeGreaterThan(0); + expect(near).toBeGreaterThan(far); + }); +}); + +describe('rectifyPoint / unrectifyPoint', () => { + it('round-trips a pixel on both cameras, with and without distortion', () => { + const plain = makeRig(); + const distorted: StereoRig = { + ...plain, + distl: Float32Array.from([-0.16, 0.10, -0.001, 0.002, 0, 0, 0, 0]), + distr: Float32Array.from([-0.15, 0.09, -0.001, 0.002, 0, 0, 0, 0]), + }; + [plain, distorted].forEach((rig) => { + const rect = computeRectification(rig, W, H); + [false, true].forEach((target) => { + [[640, 400], [300, 180], [900, 550]].forEach(([px, py]) => { + const [rx, ry] = rectifyPoint(px, py, rig, rect, target); + const [bx, by] = unrectifyPoint(rx, ry, rig, rect, target); + expect(bx).toBeCloseTo(px, 2); + expect(by).toBeCloseTo(py, 2); + }); + }); + }); + }); +}); diff --git a/client/dive-common/use/stereo/useStereoOnnxTransfer.ts b/client/dive-common/use/stereo/useStereoOnnxTransfer.ts index 0b53ab744..57f062c51 100644 --- a/client/dive-common/use/stereo/useStereoOnnxTransfer.ts +++ b/client/dive-common/use/stereo/useStereoOnnxTransfer.ts @@ -19,7 +19,8 @@ import Track from 'vue-media-annotator/track'; import { RectBounds } from 'vue-media-annotator/utils'; import { HeadPointKey, TailPointKey, HeadTailLineKey } from 'dive-common/recipes/headtail'; import type { StereoAnnotationCompleteParams } from '../useModeManager'; -import { StereoOnnxMatcher, SearchRange } from './StereoOnnxMatcher'; +import type { SearchRange } from './StereoOnnxMatcher'; +import type { StereoMatcher } from './stereoMatcher'; import { StereoRig, invertRig } from './calibration'; import { rgbaToGray, RgbaImage } from './image'; import { measureLine, aggregateLengths, StereoMeasurement } from './triangulate'; @@ -32,8 +33,12 @@ export interface StereoOnnxTransferConfig { getLeftCameraName: () => string; /** Stereo calibration, or null if unavailable (transfer is then skipped). */ getRig: () => Promise; - /** The (lazily created / cached) ONNX matcher, or null if unavailable. */ - getMatcher: () => Promise; + /** + * The (lazily created / cached) matcher for the selected method, or null if + * unavailable. Either correspondence method satisfies {@link StereoMatcher}, + * so nothing downstream branches on which one is in use. + */ + getMatcher: () => Promise; /** Full-resolution RGBA pixels for a camera at a frame, or null. */ getFrame: (cameraName: string, frameNum: number) => Promise; /** Disparity- or depth-based search range for the correspondence search. */ diff --git a/client/platform/web-girder/useStereoOnnxWeb.ts b/client/platform/web-girder/useStereoOnnxWeb.ts index a27ed8fc2..64305ac48 100644 --- a/client/platform/web-girder/useStereoOnnxWeb.ts +++ b/client/platform/web-girder/useStereoOnnxWeb.ts @@ -17,6 +17,10 @@ import { clientSettings } from 'dive-common/store/settings'; import useStereoOnnxTransfer from 'dive-common/use/stereo/useStereoOnnxTransfer'; import { StereoOnnxMatcher } from 'dive-common/use/stereo/StereoOnnxMatcher'; +import { StereoFoundationMatcher } from 'dive-common/use/stereo/StereoFoundationMatcher'; +import type { FoundationModelSpec } from 'dive-common/use/stereo/StereoFoundationMatcher'; +import { DEFAULT_STEREO_MATCH_METHOD } from 'dive-common/use/stereo/stereoMatcher'; +import type { StereoMatcher, StereoMatchMethod } from 'dive-common/use/stereo/stereoMatcher'; import type { SearchRange } from 'dive-common/use/stereo/StereoOnnxMatcher'; import { rigFromNpz, rigFromJson, StereoRig, @@ -27,6 +31,16 @@ import type { StereoMeasurement } from 'dive-common/use/stereo/triangulate'; import { getCalibrationFile, getLastCalibration } from './multicamFileRegistry'; const DEFAULT_MODEL_URL = '/models/stereo_match.onnx'; +/** + * Fast-FoundationStereo is opt-in and unbundled: the exports run ~100 MB, so + * unlike the NCC graph this one is not committed. Serve an export here (or pass + * `foundationModelUrl`) and give its sidecar `image_size` as + * `foundationModelSpec`; with no model served the dropdown's foundation option + * reports that it could not load and the warp no-ops, exactly as a missing + * calibration does. + */ +const DEFAULT_FOUNDATION_MODEL_URL = '/models/stereo_foundation.onnx'; +const DEFAULT_FOUNDATION_SPEC: FoundationModelSpec = { height: 576, width: 960 }; // Mirrors epipolar_min_disparity / epipolar_max_disparity in VIAME's // configs/pipelines/interactive_stereo_template.conf, which is what the desktop // interactive stereo service loads. Scene-dependent, and hidden config there @@ -41,6 +55,11 @@ export interface StereoOnnxWebOptions { /** Dataset (folder) id used to look up the stored calibration. */ getDatasetId: () => string; modelUrl?: string; + foundationModelUrl?: string; + /** Input resolution of the foundation export (its sidecar yaml `image_size`). */ + foundationModelSpec?: FoundationModelSpec; + /** Overrides the user's dropdown choice; mainly for tests. */ + getMatchMethod?: () => StereoMatchMethod; range?: SearchRange; onStatus?: (message: string | null) => void; onError?: (message: string) => void; @@ -80,22 +99,32 @@ async function urlToRgba(url: string): Promise { export default function useStereoOnnxWeb(opts: StereoOnnxWebOptions) { const modelUrl = opts.modelUrl ?? DEFAULT_MODEL_URL; - let matcher: StereoOnnxMatcher | null = null; - let matcherTried = false; + const foundationModelUrl = opts.foundationModelUrl ?? DEFAULT_FOUNDATION_MODEL_URL; + const foundationSpec = opts.foundationModelSpec ?? DEFAULT_FOUNDATION_SPEC; + // Cached per method: switching the dropdown must not reload the other model, + // and a method that failed to load must not be retried on every warp. + const matchers: Partial> = {}; let rig: StereoRig | null = null; let rigKey: string | null = null; - async function getMatcher(): Promise { - if (!matcher && !matcherTried) { - matcherTried = true; - try { - matcher = await StereoOnnxMatcher.create(modelUrl); - } catch (err) { - console.warn('[StereoOnnx] failed to load model', modelUrl, err); - matcher = null; - } + function currentMethod(): StereoMatchMethod { + if (opts.getMatchMethod) return opts.getMatchMethod(); + return clientSettings.stereoSettings.matchMethod ?? DEFAULT_STEREO_MATCH_METHOD; + } + + async function getMatcher(): Promise { + const method = currentMethod(); + if (method in matchers) return matchers[method] ?? null; + const url = method === 'foundation' ? foundationModelUrl : modelUrl; + try { + matchers[method] = method === 'foundation' + ? await StereoFoundationMatcher.create(url, foundationSpec) + : await StereoOnnxMatcher.create(url); + } catch (err) { + console.warn('[StereoOnnx] failed to load model', method, url, err); + matchers[method] = null; } - return matcher; + return matchers[method] ?? null; } function parseRig(name: string, buffer: ArrayBuffer): Promise { From 7c90912034b3eba0d05b714574d0286be5ed2224 Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Mon, 14 Sep 2026 12:52:29 -0400 Subject: [PATCH 2/9] Serve the foundation stereo model from VIAME's add-on list and precompute its disparity per frame Girder resolves FAST-FDN-STEREO from download_viame_addons.csv, caches it, and serves it; the client caches by md5. Fixes ImageNet/RGB preprocessing and rectify-with-resize; LRU disparity cache filled on frame change. Settings labels: "Lower Accuracy, Higher Speed" / "Higher Accuracy, Lower Speed" (web only). Claude-Session: https://claude.ai/code/session_013oFD6dP2EQY3R9UdXzx7dL --- .../components/TrackSettingsPanel.vue | 5 +- client/dive-common/store/settings.ts | 5 +- client/dive-common/use/stereo/README.md | 127 +++++--- .../use/stereo/StereoFoundationMatcher.ts | 294 +++++++++++++++--- .../use/stereo/StereoOnnxMatcher.ts | 17 +- client/dive-common/use/stereo/image.ts | 8 + client/dive-common/use/stereo/index.ts | 10 +- client/dive-common/use/stereo/rectify.ts | 74 ++++- .../dive-common/use/stereo/stereoMatcher.ts | 22 +- .../use/stereo/tests/rectify.spec.ts | 34 +- .../use/stereo/tests/stereoFoundation.spec.ts | 262 ++++++++++++++++ .../use/stereo/useStereoOnnxTransfer.ts | 48 ++- .../web-girder/api/configuration.service.ts | 22 ++ .../platform/web-girder/useStereoOnnxWeb.ts | 180 +++++++++-- docker-compose.yml | 5 + server/dive_server/views_configuration.py | 51 ++- server/dive_utils/stereo_models.py | 191 ++++++++++++ server/tests/test_stereo_models.py | 105 +++++++ 18 files changed, 1304 insertions(+), 156 deletions(-) create mode 100644 client/dive-common/use/stereo/tests/stereoFoundation.spec.ts create mode 100644 server/dive_utils/stereo_models.py create mode 100644 server/tests/test_stereo_models.py diff --git a/client/dive-common/components/TrackSettingsPanel.vue b/client/dive-common/components/TrackSettingsPanel.vue index a1ac15d86..db9cb8983 100644 --- a/client/dive-common/components/TrackSettingsPanel.vue +++ b/client/dive-common/components/TrackSettingsPanel.vue @@ -41,7 +41,7 @@ export default defineComponent({ showMultiCamToolbar: 'Show multi-camera tools in the top toolbar when a track is selected', stereoUpdateLengths: 'When a line annotation is modified on a detection that is linked across both cameras, recompute its stereo measurement (length, midpoint, range, RMS) automatically.', stereoAutoCompute: 'When an annotation is drawn on one camera and the other camera has no detection for it yet, automatically warp it to the other camera using stereo disparity.', - stereoMatchMethod: 'How a point is matched on the other camera. Template matching correlates the source patch along the epipolar line. Foundation stereo runs a dense disparity network over the pair once, which is steadier where the patch is hard to match but needs its model to be configured.', + stereoMatchMethod: 'How points are located on the other camera. "Lower Accuracy, Higher Speed" template-matches each point along its epipolar line. "Higher Accuracy, Lower Speed" runs the Fast Foundation Stereo model over the whole image pair (downloaded once, about 100 MB) and reads every point from its disparity map, which is computed ahead of time whenever you change frames.', }); const modes = ref(['Track', 'Detection']); // Add unknown as the default type to the typeList @@ -449,6 +449,7 @@ export default defineComponent({ @@ -459,7 +460,7 @@ export default defineComponent({ class="my-0 ml-1 pt-0" dense hide-details - label="Correspondence method" + label="Point matching" /> _iters__res_x.onnx` plus a sidecar `.yaml` giving - * `image_size`; the graph takes `left_image`/`right_image` as [1,3,H,W] RGB in - * [0,1] and returns `disparity` as [1,1,H,W] in rectified pixels. + * The export is the one VIAME ships in its `FAST-FDN-STEREO` add-on + * (`fast_foundation_stereo_l.onnx` plus a sidecar `.yaml` giving `image_size`), + * the same file `plugins/onnx/fast_foundation_stereo.py` runs server-side. The + * runtime contract is shared with that plugin: `left_image`/`right_image` are + * [1,3,H,W] ImageNet-normalised RGB at the export's fixed resolution, and + * `disparity` is [1,1,H,W] in rectified pixels of that resolution. + * + * Disparity maps are cached per frame pair so a warp that follows a + * {@link StereoFoundationMatcher.prepare} for the same frame is immediate. */ import * as ort from 'onnxruntime-web'; -import { GrayImage } from './image'; +import { RgbaImage, GrayImage, isGrayImage } from './image'; import { StereoRig } from './calibration'; import type { WarpOptions, WarpResult } from './StereoOnnxMatcher'; +import type { StereoMatcher } from './stereoMatcher'; import { Rectification, computeRectification, rectifyPoint, rectifyMapper, unrectifyPoint, } from './rectify'; @@ -44,12 +48,51 @@ export const DEFAULT_SAMPLE_RADIUS = 3; */ export const DEFAULT_MIN_VALID_FRACTION = 0.34; +/** Disparity maps kept per matcher: two directions per frame, so ~4 frames. */ +export const DEFAULT_DISPARITY_CACHE_SIZE = 8; + +const IMAGENET_MEAN = [0.485, 0.456, 0.406]; +const IMAGENET_STD = [0.229, 0.224, 0.225]; + export interface FoundationModelSpec { /** Network input size, from the export's sidecar yaml `image_size: [H, W]`. */ height: number; width: number; } +/** The slice of an ONNX session the matcher uses; tests substitute a fake. */ +export interface DisparitySession { + run(feeds: Record): Promise>; +} + +export interface FoundationMatcherOptions { + /** + * Defaults to WebGPU, the only provider that can run this export in a + * browser: the CPU (wasm) path needs several GB of activations, past the + * 4 GB a wasm heap can hold, so it is not offered as a fallback. + */ + executionProviders?: string[]; + cacheSize?: number; +} + +export const WEBGPU_REQUIRED_MESSAGE = 'The higher-accuracy stereo model needs WebGPU, which this browser does not provide. ' + + 'Use a current Chrome or Edge, or switch the point matching setting to the faster method.'; + +function defaultExecutionProviders(): string[] { + const hasWebGpu = typeof navigator !== 'undefined' && 'gpu' in navigator && !!navigator.gpu; + if (!hasWebGpu) throw new Error(WEBGPU_REQUIRED_MESSAGE); + return ['webgpu']; +} + +function rigKey(rig: StereoRig): string { + return `${rig.Kl.join(',')}|${rig.Kr.join(',')}|${rig.R.join(',')}|${rig.T.join(',')}`; +} + +function cacheKey(frameKey: string | undefined, rig: StereoRig, source: { width: number; height: number }): string | null { + if (frameKey === undefined) return null; + return `${frameKey}|${source.width}x${source.height}|${rigKey(rig)}`; +} + /** Bilinear sample of a single-channel image, NaN outside. */ function sampleBilinear(data: Float32Array, width: number, height: number, x: number, y: number): number { if (!(x >= 0 && y >= 0 && x <= width - 1 && y <= height - 1)) return NaN; @@ -66,35 +109,99 @@ function sampleBilinear(data: Float32Array, width: number, height: number, x: nu return a * (1 - fx) * (1 - fy) + b * fx * (1 - fy) + c * (1 - fx) * fy + d * fx * fy; } -/** Remap a grayscale frame through an inverse map into an RGB [1,3,H,W] tensor. */ -function remapToRgbTensor(src: GrayImage, mapX: Float32Array, mapY: Float32Array, width: number, height: number): ort.Tensor { +/** + * Remap a frame through an inverse map into an ImageNet-normalised RGB + * [1,3,H,W] tensor. Pixels that map outside the source are filled black, as + * OpenCV's remap would. + */ +export function remapToInputTensor( + src: RgbaImage | GrayImage, + mapX: Float32Array, + mapY: Float32Array, + width: number, + height: number, +): ort.Tensor { const plane = width * height; const out = new Float32Array(plane * 3); + const { data, width: sw, height: sh } = src; + const gray = isGrayImage(src); + const stride = gray ? 1 : 4; + const rgb = [0, 0, 0]; for (let i = 0; i < plane; i += 1) { - const v = sampleBilinear(src.data, src.width, src.height, mapX[i], mapY[i]); - const g = Number.isNaN(v) ? 0 : v; - out[i] = g; - out[plane + i] = g; - out[2 * plane + i] = g; + const x = mapX[i]; + const y = mapY[i]; + rgb[0] = 0; + rgb[1] = 0; + rgb[2] = 0; + if (x >= 0 && y >= 0 && x <= sw - 1 && y <= sh - 1) { + const x0 = Math.floor(x); + const y0 = Math.floor(y); + const x1 = Math.min(x0 + 1, sw - 1); + const y1 = Math.min(y0 + 1, sh - 1); + const fx = x - x0; + const fy = y - y0; + const w00 = (1 - fx) * (1 - fy); + const w10 = fx * (1 - fy); + const w01 = (1 - fx) * fy; + const w11 = fx * fy; + const i00 = (y0 * sw + x0) * stride; + const i10 = (y0 * sw + x1) * stride; + const i01 = (y1 * sw + x0) * stride; + const i11 = (y1 * sw + x1) * stride; + for (let c = 0; c < 3; c += 1) { + const o = gray ? 0 : c; + rgb[c] = data[i00 + o] * w00 + data[i10 + o] * w10 + data[i01 + o] * w01 + data[i11 + o] * w11; + } + } + for (let c = 0; c < 3; c += 1) { + out[c * plane + i] = (rgb[c] / 255 - IMAGENET_MEAN[c]) / IMAGENET_STD[c]; + } } return new ort.Tensor('float32', out, [1, 3, height, width]); } -export class StereoFoundationMatcher { - private session: ort.InferenceSession; +interface Geometry { + key: string; + rect: Rectification; + src: { mapX: Float32Array; mapY: Float32Array }; + tgt: { mapX: Float32Array; mapY: Float32Array }; +} + +interface PendingDisparity { + promise: Promise; + /** Cleared when a warp needs the result, so the staleness check cannot skip it. */ + prefetch: boolean; +} + +export class StereoFoundationMatcher implements StereoMatcher { + private session: DisparitySession; private spec: FoundationModelSpec; + private cacheSize: number; + /** Rectification + inverse maps, rebuilt only when the rig or size changes. */ - private cache: { - key: string; rect: Rectification; - src: { mapX: Float32Array; mapY: Float32Array }; - tgt: { mapX: Float32Array; mapY: Float32Array }; - } | null = null; + private geometryCache: Geometry | null = null; + + /** Insertion-ordered so the first entry is the least recently used. */ + private disparities = new Map(); + + private pending = new Map(); + + /** The wasm/WebGPU session runs one inference at a time. */ + private queue: Promise = Promise.resolve(); + + /** + * An inference failure is a property of this model + runtime (unsupported + * operator, out of GPU memory), not of one frame, so it is remembered and + * rethrown instead of re-running a doomed pass on every frame change. + */ + private failure: Error | null = null; - private constructor(session: ort.InferenceSession, spec: FoundationModelSpec) { + constructor(session: DisparitySession, spec: FoundationModelSpec, cacheSize = DEFAULT_DISPARITY_CACHE_SIZE) { this.session = session; this.spec = spec; + this.cacheSize = cacheSize; } /** @@ -104,29 +211,131 @@ export class StereoFoundationMatcher { static async create( model: string | ArrayBuffer | Uint8Array, spec: FoundationModelSpec, - opts: { threads?: number } = {}, + opts: FoundationMatcherOptions = {}, ): Promise { - ort.env.wasm.numThreads = opts.threads ?? 1; + const executionProviders = opts.executionProviders ?? defaultExecutionProviders(); ort.env.wasm.proxy = false; const session = await ort.InferenceSession.create(model as string, { - executionProviders: ['wasm'], + executionProviders, graphOptimizationLevel: 'all', }); - return new StereoFoundationMatcher(session, spec); + return new StereoFoundationMatcher(session, spec, opts.cacheSize); + } + + get inputSize(): FoundationModelSpec { + return { ...this.spec }; } /** Rectification and inverse maps for this rig at the model's resolution. */ - private geometry(rig: StereoRig) { - const key = `${rig.Kl.join(',')}|${rig.R.join(',')}|${rig.T.join(',')}`; - if (this.cache && this.cache.key === key) return this.cache; - const rect = computeRectification(rig, this.spec.width, this.spec.height); - this.cache = { + private geometry(rig: StereoRig, source: { width: number; height: number }): Geometry { + const key = `${source.width}x${source.height}|${rigKey(rig)}`; + if (this.geometryCache && this.geometryCache.key === key) return this.geometryCache; + const rect = computeRectification(rig, source.width, source.height, this.spec.width, this.spec.height); + this.geometryCache = { key, rect, src: rectifyMapper(rig, rect, false), tgt: rectifyMapper(rig, rect, true), }; - return this.cache; + return this.geometryCache; + } + + private async infer(source: RgbaImage | GrayImage, target: RgbaImage | GrayImage, rig: StereoRig): Promise { + if (this.failure) throw this.failure; + const { src, tgt } = this.geometry(rig, source); + const { width, height } = this.spec; + let out: Record; + try { + out = await this.session.run({ + left_image: remapToInputTensor(source, src.mapX, src.mapY, width, height), + right_image: remapToInputTensor(target, tgt.mapX, tgt.mapY, width, height), + }); + } catch (err) { + this.failure = err instanceof Error ? err : new Error(String(err)); + throw this.failure; + } + const disparity = out.disparity.data as Float32Array; + for (let i = 0; i < disparity.length; i += 1) { + if (!(disparity[i] > 0)) disparity[i] = 0; + } + return disparity; + } + + private remember(key: string, disparity: Float32Array) { + this.disparities.delete(key); + this.disparities.set(key, disparity); + while (this.disparities.size > this.cacheSize) { + const oldest = this.disparities.keys().next().value as string; + this.disparities.delete(oldest); + } + } + + /** + * The disparity map for a pair, from the cache when present. Requests are + * serialised; a `prefetch` request whose `stillWanted` has turned false by + * the time it reaches the front of the queue resolves `null` without running. + */ + private disparityFor( + key: string | null, + source: RgbaImage | GrayImage, + target: RgbaImage | GrayImage, + rig: StereoRig, + prefetch: boolean, + stillWanted: () => boolean = () => true, + ): Promise { + if (key !== null) { + const hit = this.disparities.get(key); + if (hit) { + this.remember(key, hit); + return Promise.resolve(hit); + } + const inflight = this.pending.get(key); + if (inflight) { + if (!prefetch) inflight.prefetch = false; + return inflight.promise; + } + } + const entry: PendingDisparity = { prefetch, promise: Promise.resolve(null) }; + entry.promise = this.queue.then(async () => { + if (entry.prefetch && !stillWanted()) return null; + const hit = key !== null ? this.disparities.get(key) : undefined; + if (hit) return hit; + const disparity = await this.infer(source, target, rig); + if (key !== null) this.remember(key, disparity); + return disparity; + }); + this.queue = entry.promise.catch(() => undefined); + if (key !== null) { + this.pending.set(key, entry); + entry.promise.finally(() => this.pending.delete(key)).catch(() => undefined); + } + return entry.promise; + } + + /** The error that stopped this matcher, if an inference has failed. */ + get lastFailure(): Error | null { + return this.failure; + } + + /** Whether a disparity map is already cached for this frame pair. */ + isPrepared(frameKey: string, rig: StereoRig, source: { width: number; height: number }): boolean { + const key = cacheKey(frameKey, rig, source); + return key !== null && this.disparities.has(key); + } + + /** + * Compute and cache the disparity map for a frame pair ahead of any warp + * against it. Resolves once the map is cached or the request was dropped as + * stale. + */ + async prepare( + frameKey: string, + source: RgbaImage | GrayImage, + target: RgbaImage | GrayImage, + rig: StereoRig, + stillWanted: () => boolean = () => true, + ): Promise { + await this.disparityFor(cacheKey(frameKey, rig, source), source, target, rig, true, stillWanted); } /** @@ -138,20 +347,15 @@ export class StereoFoundationMatcher { */ async warpPoints( points: [number, number][], - source: GrayImage, - target: GrayImage, + source: RgbaImage | GrayImage, + target: RgbaImage | GrayImage, rig: StereoRig, opts: WarpOptions, ): Promise { - const { rect, src, tgt } = this.geometry(rig); + const { rect } = this.geometry(rig, source); const { width, height } = this.spec; - - const feeds: Record = { - left_image: remapToRgbTensor(source, src.mapX, src.mapY, width, height), - right_image: remapToRgbTensor(target, tgt.mapX, tgt.mapY, width, height), - }; - const out = await this.session.run(feeds); - const disparity = out.disparity.data as Float32Array; + const disparity = await this.disparityFor(cacheKey(opts.frameKey, rig, source), source, target, rig, false); + if (!disparity) throw new Error('The stereo disparity map could not be computed.'); const radius = DEFAULT_SAMPLE_RADIUS; const minValid = DEFAULT_MIN_VALID_FRACTION; @@ -174,8 +378,10 @@ export class StereoFoundationMatcher { for (let dy = -radius; dy <= radius; dy += 1) { for (let dx = -radius; dx <= radius; dx += 1) { considered += 1; - const v = sampleBilinear(disparity, width, height, rx + dx, ry + dy); - if (Number.isFinite(v) && v > 0) samples.push(v); + const sx = rx + dx; + const v = sampleBilinear(disparity, width, height, sx, ry + dy); + // A match that would land left of the target image is invisible. + if (Number.isFinite(v) && v > 0 && sx - v >= 0) samples.push(v); } } if (!samples.length) return fail; diff --git a/client/dive-common/use/stereo/StereoOnnxMatcher.ts b/client/dive-common/use/stereo/StereoOnnxMatcher.ts index 40f6fd274..8df33cad1 100644 --- a/client/dive-common/use/stereo/StereoOnnxMatcher.ts +++ b/client/dive-common/use/stereo/StereoOnnxMatcher.ts @@ -11,7 +11,7 @@ import * as ort from 'onnxruntime-web'; -import { GrayImage } from './image'; +import { GrayImage, RgbaImage, toGray } from './image'; import { StereoRig, baseline } from './calibration'; /** Search-range specification (disparity is unit-independent; depth needs calib units). */ @@ -25,6 +25,11 @@ export interface WarpOptions { threshold?: number; /** Reject if secondScore/score exceeds this (0 disables). Defaults to {@link DEFAULT_UNIQUENESS_RATIO}. */ uniquenessRatio?: number; + /** + * Identifies the frame pair the points belong to, so a matcher that computes + * per-frame state (a disparity map) can reuse it across calls. + */ + frameKey?: string; } /** @@ -94,16 +99,18 @@ export class StereoOnnxMatcher { /** * Warp a set of source-image points onto the target image. `source`/`target` - * are grayscale frames; `rig` is the stereo calibration with `source` as the - * left camera. Returns one {@link WarpResult} per input point. + * are RGBA or grayscale frames; `rig` is the stereo calibration with `source` + * as the left camera. Returns one {@link WarpResult} per input point. */ async warpPoints( points: [number, number][], - source: GrayImage, - target: GrayImage, + sourceImage: RgbaImage | GrayImage, + targetImage: RgbaImage | GrayImage, rig: StereoRig, opts: WarpOptions, ): Promise { + const source = toGray(sourceImage); + const target = toGray(targetImage); const [minDepth, maxDepth] = resolveDepthRange(rig, opts.range); const threshold = opts.threshold ?? DEFAULT_THRESHOLD; const uniqueness = opts.uniquenessRatio ?? DEFAULT_UNIQUENESS_RATIO; diff --git a/client/dive-common/use/stereo/image.ts b/client/dive-common/use/stereo/image.ts index 1944f6cb1..d1ba077ec 100644 --- a/client/dive-common/use/stereo/image.ts +++ b/client/dive-common/use/stereo/image.ts @@ -22,6 +22,10 @@ const R_W = 0.299; const G_W = 0.587; const B_W = 0.114; +export function isGrayImage(img: RgbaImage | GrayImage): img is GrayImage { + return img.data instanceof Float32Array; +} + /** Convert RGBA pixel data to BT.601 grayscale (matches cv2 BGR2GRAY). */ export function rgbaToGray(img: RgbaImage): GrayImage { const { data, width, height } = img; @@ -32,6 +36,10 @@ export function rgbaToGray(img: RgbaImage): GrayImage { return { data: out, width, height }; } +export function toGray(img: RgbaImage | GrayImage): GrayImage { + return isGrayImage(img) ? img : rgbaToGray(img); +} + /** * Draw an image source (HTMLImageElement / HTMLCanvasElement / * ImageBitmap) to an offscreen canvas and return BT.601 grayscale. Browser / diff --git a/client/dive-common/use/stereo/index.ts b/client/dive-common/use/stereo/index.ts index 01db38c26..10b9773a6 100644 --- a/client/dive-common/use/stereo/index.ts +++ b/client/dive-common/use/stereo/index.ts @@ -1,6 +1,8 @@ export { StereoOnnxMatcher } from './StereoOnnxMatcher'; -export { StereoFoundationMatcher } from './StereoFoundationMatcher'; -export type { FoundationModelSpec } from './StereoFoundationMatcher'; +export { StereoFoundationMatcher, remapToInputTensor } from './StereoFoundationMatcher'; +export type { + FoundationModelSpec, FoundationMatcherOptions, DisparitySession, +} from './StereoFoundationMatcher'; export { DEFAULT_STEREO_MATCH_METHOD, STEREO_MATCH_METHODS, } from './stereoMatcher'; @@ -16,7 +18,9 @@ export { export type { StereoRig } from './calibration'; export { parseNpz, parseNpy } from './npz'; export type { NpyArray } from './npz'; -export { rgbaToGray, drawableToGray } from './image'; +export { + rgbaToGray, drawableToGray, toGray, isGrayImage, +} from './image'; export type { GrayImage, RgbaImage } from './image'; export { triangulatePoint, measureLine, aggregateLengths, unmap, mapPoint, project, diff --git a/client/dive-common/use/stereo/rectify.ts b/client/dive-common/use/stereo/rectify.ts index 79cad859d..4267ac300 100644 --- a/client/dive-common/use/stereo/rectify.ts +++ b/client/dive-common/use/stereo/rectify.ts @@ -21,8 +21,13 @@ export interface Rectification { /** Rectifying rotations for the source and target cameras. */ R1: Mat3; R2: Mat3; - /** Shared focal length and principal point of the rectified pair. */ - f: number; + /** + * Shared focal lengths and principal point of the rectified pair, in + * rectified pixels. `fx` and `fy` differ when the rectified size has a + * different aspect ratio from the source. + */ + fx: number; + fy: number; cx: number; cy: number; /** Rectified image size these were solved for. */ @@ -78,15 +83,26 @@ export function rodrigues(v: ArrayLike): Mat3 { } /** - * Solve the rectifying rotations for a rig. + * Solve the rectifying rotations for a rig whose frames are `sourceWidth` x + * `sourceHeight`, producing rectified images of `width` x `height` (default: + * the source size). * - * The focal length is the smaller of the two cameras' so the rectified frustum - * stays inside both, and the principal point is centred on the output. This is - * OpenCV's `alpha = 0`-free behaviour: no zoom-to-valid-region crop, which on a - * rig whose baseline sits far from horizontal demands an extreme zoom and can - * push the whole scene off canvas. + * As in OpenCV, the focal length is the mean of the two cameras' along the axis + * perpendicular to the baseline, and the principal point is placed so the + * average of each camera's undistorted, rotated corners lands at the centre. + * Both are then scaled by the output/source size ratio per axis, which fuses + * rectification with a resize to the output resolution. This is the + * `alpha = -1` behaviour: no zoom-to-valid-region crop, which on a rig whose + * baseline sits far from horizontal demands an extreme zoom and can push the + * whole scene off canvas. */ -export function computeRectification(rig: StereoRig, width: number, height: number): Rectification { +export function computeRectification( + rig: StereoRig, + sourceWidth: number, + sourceHeight: number, + width = sourceWidth, + height = sourceHeight, +): Rectification { // Half-rotate both cameras toward each other: r = R^(-1/2). const om = rodriguesInv(rig.R); const r = rodrigues([-om[0] / 2, -om[1] / 2, -om[2] / 2]); @@ -113,13 +129,39 @@ export function computeRectification(rig: StereoRig, width: number, height: numb const scale = Math.acos(Math.abs(t[idx]) / nt) / nw; wR = rodrigues([ww[0] * scale, ww[1] * scale, ww[2] * scale]); } + const R1 = matMul(wR, transpose(r)); + const R2 = matMul(wR, r); + + const perp = horizontal ? 4 : 0; + const f = (rig.Kl[perp] + rig.Kr[perp]) / 2; + const fx = f * (width / sourceWidth); + const fy = f * (height / sourceHeight); + + const corners: [number, number][] = [ + [0, 0], [sourceWidth - 1, 0], [0, sourceHeight - 1], [sourceWidth - 1, sourceHeight - 1], + ]; + const centres = [false, true].map((target) => { + const K = target ? rig.Kr : rig.Kl; + const d = target ? rig.distr : rig.distl; + const R = target ? R2 : R1; + let sx = 0; + let sy = 0; + corners.forEach(([px, py]) => { + const [nx, ny] = unmap(px, py, K, d); + const p = matVec(R, [nx, ny, 1]); + sx += p[0] / p[2]; + sy += p[1] / p[2]; + }); + return [(width - 1) / 2 - fx * (sx / 4), (height - 1) / 2 - fy * (sy / 4)]; + }); return { - R1: matMul(wR, transpose(r)), - R2: matMul(wR, r), - f: Math.min(rig.Kl[0], rig.Kr[0]), - cx: (width - 1) / 2, - cy: (height - 1) / 2, + R1, + R2, + fx, + fy, + cx: (centres[0][0] + centres[1][0]) / 2, + cy: (centres[0][1] + centres[1][1]) / 2, width, height, }; @@ -133,7 +175,7 @@ export function rectifyPoint(px: number, py: number, rig: StereoRig, rect: Recti const [nx, ny] = unmap(px, py, K, d); const p = matVec(R, [nx, ny, 1]); if (p[2] === 0) return [NaN, NaN]; - return [rect.f * (p[0] / p[2]) + rect.cx, rect.f * (p[1] / p[2]) + rect.cy]; + return [rect.fx * (p[0] / p[2]) + rect.cx, rect.fy * (p[1] / p[2]) + rect.cy]; } /** Rectified pixel -> source pixel (the inverse of {@link rectifyPoint}). */ @@ -141,7 +183,7 @@ export function unrectifyPoint(rx: number, ry: number, rig: StereoRig, rect: Rec const K = target ? rig.Kr : rig.Kl; const d = target ? rig.distr : rig.distl; const R = target ? rect.R2 : rect.R1; - const p = matVec(transpose(R), [(rx - rect.cx) / rect.f, (ry - rect.cy) / rect.f, 1]); + const p = matVec(transpose(R), [(rx - rect.cx) / rect.fx, (ry - rect.cy) / rect.fy, 1]); if (p[2] === 0) return [NaN, NaN]; return mapPoint(p[0] / p[2], p[1] / p[2], K, d); } diff --git a/client/dive-common/use/stereo/stereoMatcher.ts b/client/dive-common/use/stereo/stereoMatcher.ts index be22ac7d7..2cb12fc67 100644 --- a/client/dive-common/use/stereo/stereoMatcher.ts +++ b/client/dive-common/use/stereo/stereoMatcher.ts @@ -3,7 +3,7 @@ * and its callers never branch on which one is selected. */ -import { GrayImage } from './image'; +import { GrayImage, RgbaImage } from './image'; import { StereoRig } from './calibration'; import type { WarpOptions, WarpResult } from './StereoOnnxMatcher'; @@ -18,15 +18,27 @@ export const DEFAULT_STEREO_MATCH_METHOD: StereoMatchMethod = 'ncc'; export interface StereoMatcher { warpPoints( points: [number, number][], - source: GrayImage, - target: GrayImage, + source: RgbaImage | GrayImage, + target: RgbaImage | GrayImage, rig: StereoRig, opts: WarpOptions, ): Promise; + /** + * Compute and cache whatever per-frame state a later `warpPoints` with the + * same `frameKey` would need, so the warp itself is quick. `stillWanted` is + * polled before the work starts so a stale request can be dropped. + */ + prepare?( + frameKey: string, + source: RgbaImage, + target: RgbaImage, + rig: StereoRig, + stillWanted?: () => boolean, + ): Promise; } /** Labels for the method selector. */ export const STEREO_MATCH_METHODS: { value: StereoMatchMethod; text: string }[] = [ - { value: 'ncc', text: 'Template matching (NCC)' }, - { value: 'foundation', text: 'Foundation stereo (disparity)' }, + { value: 'ncc', text: 'Lower Accuracy, Higher Speed' }, + { value: 'foundation', text: 'Higher Accuracy, Lower Speed' }, ]; diff --git a/client/dive-common/use/stereo/tests/rectify.spec.ts b/client/dive-common/use/stereo/tests/rectify.spec.ts index 3daded55c..7760c8c5a 100644 --- a/client/dive-common/use/stereo/tests/rectify.spec.ts +++ b/client/dive-common/use/stereo/tests/rectify.spec.ts @@ -82,6 +82,29 @@ describe('computeRectification', () => { }); }); + it('scales the rectified view to a different output size', () => { + // The network input is smaller than the frame, so rectification must also + // resize: the same source pixel lands at the same relative position, and a + // point at a source corner stays inside the output instead of being cropped. + const rig = makeRig(); + const full = computeRectification(rig, 1280, 800); + const small = computeRectification(rig, 1280, 800, W, H); + expect(small.fx).toBeCloseTo(full.fx * (W / 1280), 6); + expect(small.fy).toBeCloseTo(full.fy * (H / 800), 6); + // Pixel centres: position relative to the image centre scales with the size. + [[0, 0], [1279, 799], [640, 400]].forEach(([px, py]) => { + const [fx, fy] = rectifyPoint(px, py, rig, full, false); + const [sx, sy] = rectifyPoint(px, py, rig, small, false); + expect(sx - (W - 1) / 2).toBeCloseTo((fx - (1280 - 1) / 2) * (W / 1280), 3); + expect(sy - (H - 1) / 2).toBeCloseTo((fy - (800 - 1) / 2) * (H / 800), 3); + }); + const [cx, cy] = rectifyPoint(640, 400, rig, small, false); + expect(cx).toBeGreaterThan(W * 0.3); + expect(cx).toBeLessThan(W * 0.7); + expect(cy).toBeGreaterThan(H * 0.3); + expect(cy).toBeLessThan(H * 0.7); + }); + it('gives a positive disparity that shrinks with range', () => { const rig = makeRig(); const rect = computeRectification(rig, W, H); @@ -100,6 +123,15 @@ describe('computeRectification', () => { }); describe('rectifyPoint / unrectifyPoint', () => { + it('keeps corresponding points on one row when the output size differs', () => { + const rig = makeRig(); + const rect = computeRectification(rig, 1280, 800, W, H); + const p: [number, number, number] = [150, -80, 2200]; + const [, ly] = rectifyPoint(...projectInto(p, rig, false), rig, rect, false); + const [, ry] = rectifyPoint(...projectInto(p, rig, true), rig, rect, true); + expect(ry).toBeCloseTo(ly, 2); + }); + it('round-trips a pixel on both cameras, with and without distortion', () => { const plain = makeRig(); const distorted: StereoRig = { @@ -108,7 +140,7 @@ describe('rectifyPoint / unrectifyPoint', () => { distr: Float32Array.from([-0.15, 0.09, -0.001, 0.002, 0, 0, 0, 0]), }; [plain, distorted].forEach((rig) => { - const rect = computeRectification(rig, W, H); + const rect = computeRectification(rig, 1280, 800, W, H); [false, true].forEach((target) => { [[640, 400], [300, 180], [900, 550]].forEach(([px, py]) => { const [rx, ry] = rectifyPoint(px, py, rig, rect, target); diff --git a/client/dive-common/use/stereo/tests/stereoFoundation.spec.ts b/client/dive-common/use/stereo/tests/stereoFoundation.spec.ts new file mode 100644 index 000000000..6a0b6c046 --- /dev/null +++ b/client/dive-common/use/stereo/tests/stereoFoundation.spec.ts @@ -0,0 +1,262 @@ +/** + * The foundation matcher's bookkeeping — preprocessing, per-frame disparity + * caching, prefetch dropping, serialised inference — runs here against a fake + * session returning a constant disparity, so the geometry it implies is known. + * + * The real export (~100 MB) is not in the repo. Point + * DIVE_STEREO_FOUNDATION_MODEL at a `fast_foundation_stereo_l.onnx` (from + * VIAME's FAST-FDN-STEREO add-on) to also check its I/O contract. + */ + +import { existsSync, readFileSync } from 'fs'; +import { fileURLToPath } from 'url'; +import { PNG } from 'pngjs'; +import * as ort from 'onnxruntime-web'; +import { + describe, it, expect, +} from 'vitest'; + +import { + StereoFoundationMatcher, remapToInputTensor, DisparitySession, WEBGPU_REQUIRED_MESSAGE, +} from '../StereoFoundationMatcher'; +import { rigFromNpz, StereoRig } from '../calibration'; +import { rodrigues, computeRectification, rectifyPoint } from '../rectify'; +import { RgbaImage } from '../image'; + +const fixture = (name: string) => fileURLToPath(new URL(`./fixtures/${name}`, import.meta.url)); + +const SPEC = { width: 96, height: 64 }; +const SRC = { width: 320, height: 240 }; + +function makeRig(): StereoRig { + const K = Float32Array.from([300, 0, 160, 0, 300, 120, 0, 0, 1]); + return { + Kl: K, + Kr: Float32Array.from(K), + distl: new Float32Array(8), + distr: new Float32Array(8), + R: rodrigues([0, 0, 0]), + T: Float32Array.from([-100, 0, 0]), + }; +} + +function solidImage(width: number, height: number, rgb: [number, number, number]): RgbaImage { + const data = new Uint8ClampedArray(width * height * 4); + for (let i = 0; i < width * height; i += 1) { + data.set([...rgb, 255], i * 4); + } + return { data, width, height }; +} + +/** Returns a constant disparity and counts runs; `gate` can hold a run open. */ +function fakeSession(disparity: number): DisparitySession & { runs: number; feeds: Record[] } { + const session = { + runs: 0, + feeds: [] as Record[], + async run(feeds: Record) { + session.runs += 1; + session.feeds.push(feeds); + const plane = SPEC.width * SPEC.height; + return { + disparity: new ort.Tensor('float32', new Float32Array(plane).fill(disparity), [1, 1, SPEC.height, SPEC.width]), + }; + }, + }; + return session; +} + +const RANGE = { range: { minDisparity: 1, maxDisparity: 200 } }; + +describe('remapToInputTensor', () => { + it('produces ImageNet-normalised RGB planes and black outside the source', () => { + const img = solidImage(4, 4, [255, 128, 0]); + const mapX = Float32Array.from([1, 1, -5]); + const mapY = Float32Array.from([1, 2, 1]); + const t = remapToInputTensor(img, mapX, mapY, 3, 1); + expect(t.dims).toEqual([1, 3, 1, 3]); + const d = t.data as Float32Array; + expect(d[0]).toBeCloseTo((1 - 0.485) / 0.229, 4); + expect(d[3]).toBeCloseTo((128 / 255 - 0.456) / 0.224, 4); + expect(d[6]).toBeCloseTo((0 - 0.406) / 0.225, 4); + // Outside the source: black, not zero-after-normalisation. + expect(d[2]).toBeCloseTo((0 - 0.485) / 0.229, 4); + }); + + it('accepts a grayscale frame by replicating it across channels', () => { + const gray = { data: Float32Array.from([200, 200, 200, 200]), width: 2, height: 2 }; + const t = remapToInputTensor(gray, Float32Array.from([0.5]), Float32Array.from([0.5]), 1, 1); + const d = t.data as Float32Array; + expect(d[0]).toBeCloseTo((200 / 255 - 0.485) / 0.229, 4); + expect(d[1]).toBeCloseTo((200 / 255 - 0.456) / 0.224, 4); + }); +}); + +describe('StereoFoundationMatcher.create', () => { + it('refuses to run without WebGPU rather than falling back to wasm', async () => { + await expect(StereoFoundationMatcher.create(new Uint8Array(0), SPEC)).rejects.toThrow(WEBGPU_REQUIRED_MESSAGE); + }); +}); + +describe('StereoFoundationMatcher with a fake session', () => { + const rig = makeRig(); + const left = solidImage(SRC.width, SRC.height, [10, 20, 30]); + const right = solidImage(SRC.width, SRC.height, [30, 20, 10]); + const points: [number, number][] = [[160, 120], [200, 100]]; + + it('shifts each point by the disparity, scaled back to source pixels', async () => { + const disparity = 6; + const matcher = new StereoFoundationMatcher(fakeSession(disparity), SPEC); + const res = await matcher.warpPoints(points, left, right, rig, RANGE); + const rect = computeRectification(rig, SRC.width, SRC.height, SPEC.width, SPEC.height); + res.forEach((r, i) => { + const [px, py] = points[i]; + const [rx] = rectifyPoint(px, py, rig, rect, false); + expect(r.accepted).toBe(true); + expect(r.score).toBe(1); + // Aligned identical cameras: rectification is a pure scale, so the + // shift in source pixels is the disparity divided by the scale. + expect(r.x).toBeCloseTo(px - disparity / (SPEC.width / SRC.width), 3); + expect(r.y).toBeCloseTo(py, 3); + expect(rx).toBeGreaterThan(disparity); + }); + }); + + it('feeds the model at its own resolution', async () => { + const session = fakeSession(4); + const matcher = new StereoFoundationMatcher(session, SPEC); + await matcher.warpPoints(points, left, right, rig, RANGE); + expect(session.feeds[0].left_image.dims).toEqual([1, 3, SPEC.height, SPEC.width]); + expect(session.feeds[0].right_image.dims).toEqual([1, 3, SPEC.height, SPEC.width]); + }); + + it('rejects a disparity outside the configured range', async () => { + const matcher = new StereoFoundationMatcher(fakeSession(6), SPEC); + // 6 model px = 20 source px; a range that excludes it must reject. + const res = await matcher.warpPoints(points, left, right, rig, { range: { minDisparity: 30, maxDisparity: 200 } }); + expect(res.every((r) => !r.accepted)).toBe(true); + expect(res.every((r) => Number.isFinite(r.x))).toBe(true); + }); + + it('reuses the disparity map across warps with the same frame key', async () => { + const session = fakeSession(4); + const matcher = new StereoFoundationMatcher(session, SPEC); + await matcher.warpPoints(points, left, right, rig, { ...RANGE, frameKey: 'a>b@1' }); + await matcher.warpPoints(points.slice(0, 1), left, right, rig, { ...RANGE, frameKey: 'a>b@1' }); + expect(session.runs).toBe(1); + await matcher.warpPoints(points, left, right, rig, { ...RANGE, frameKey: 'a>b@2' }); + expect(session.runs).toBe(2); + // Without a key nothing can be reused. + await matcher.warpPoints(points, left, right, rig, RANGE); + await matcher.warpPoints(points, left, right, rig, RANGE); + expect(session.runs).toBe(4); + }); + + it('prepare makes the following warp free', async () => { + const session = fakeSession(4); + const matcher = new StereoFoundationMatcher(session, SPEC); + await matcher.prepare('a>b@7', left, right, rig); + expect(session.runs).toBe(1); + expect(matcher.isPrepared('a>b@7', rig, left)).toBe(true); + await matcher.warpPoints(points, left, right, rig, { ...RANGE, frameKey: 'a>b@7' }); + expect(session.runs).toBe(1); + }); + + it('drops a prefetch that is no longer wanted, but not one a warp is waiting on', async () => { + const session = fakeSession(4); + const matcher = new StereoFoundationMatcher(session, SPEC); + await matcher.prepare('stale', left, right, rig, () => false); + expect(session.runs).toBe(0); + expect(matcher.isPrepared('stale', rig, left)).toBe(false); + + // Queue a prefetch that would be stale, then ask for its result before it + // runs: the warp must upgrade it rather than get nothing. + const prefetch = matcher.prepare('upgraded', left, right, rig, () => false); + const warp = matcher.warpPoints(points, left, right, rig, { ...RANGE, frameKey: 'upgraded' }); + await Promise.all([prefetch, warp]); + expect(session.runs).toBe(1); + expect((await warp)[0].accepted).toBe(true); + }); + + it('runs one inference at a time and dedupes concurrent requests for a frame', async () => { + let active = 0; + let maxActive = 0; + const session: DisparitySession & { runs: number } = { + runs: 0, + async run() { + session.runs += 1; + active += 1; + maxActive = Math.max(maxActive, active); + await new Promise((resolve) => { setTimeout(resolve, 5); }); + active -= 1; + const plane = SPEC.width * SPEC.height; + return { disparity: new ort.Tensor('float32', new Float32Array(plane).fill(3), [1, 1, SPEC.height, SPEC.width]) }; + }, + }; + const matcher = new StereoFoundationMatcher(session, SPEC); + await Promise.all([ + matcher.prepare('f1', left, right, rig), + matcher.warpPoints(points, left, right, rig, { ...RANGE, frameKey: 'f1' }), + matcher.prepare('f2', left, right, rig), + matcher.warpPoints(points, left, right, rig, { ...RANGE, frameKey: 'f2' }), + ]); + expect(session.runs).toBe(2); + expect(maxActive).toBe(1); + }); + + it('evicts the least recently used map beyond the cache size', async () => { + const session = fakeSession(4); + const matcher = new StereoFoundationMatcher(session, SPEC, 2); + await matcher.prepare('k1', left, right, rig); + await matcher.prepare('k2', left, right, rig); + await matcher.warpPoints(points, left, right, rig, { ...RANGE, frameKey: 'k1' }); + await matcher.prepare('k3', left, right, rig); + expect(matcher.isPrepared('k1', rig, left)).toBe(true); + expect(matcher.isPrepared('k2', rig, left)).toBe(false); + expect(matcher.isPrepared('k3', rig, left)).toBe(true); + expect(session.runs).toBe(3); + }); + + it('remembers an inference failure instead of retrying every frame', async () => { + let runs = 0; + const session: DisparitySession = { + async run() { + runs += 1; + throw new Error('[WebGPU] Kernel failed'); + }, + }; + const matcher = new StereoFoundationMatcher(session, SPEC); + await expect(matcher.prepare('f1', left, right, rig)).rejects.toThrow('Kernel failed'); + await expect(matcher.warpPoints(points, left, right, rig, { ...RANGE, frameKey: 'f2' })).rejects.toThrow('Kernel failed'); + expect(runs).toBe(1); + expect(matcher.lastFailure?.message).toContain('Kernel failed'); + }); + + it('keys the cache by rig so a new calibration is not served an old map', async () => { + const session = fakeSession(4); + const matcher = new StereoFoundationMatcher(session, SPEC); + await matcher.prepare('f', left, right, rig); + const other = { ...rig, T: Float32Array.from([-120, 0, 0]) }; + expect(matcher.isPrepared('f', other, left)).toBe(false); + await matcher.warpPoints(points, left, right, other, { ...RANGE, frameKey: 'f' }); + expect(session.runs).toBe(2); + }); +}); + +const realModel = process.env.DIVE_STEREO_FOUNDATION_MODEL; + +describe.skipIf(!realModel || !existsSync(realModel))('the served export', () => { + // Running it needs WebGPU (the CPU path needs ~7 GB, past the wasm heap), so + // Node can only check the graph honours the contract the matcher assumes. + it('exposes left_image/right_image -> disparity at the sidecar resolution', async () => { + const session = await ort.InferenceSession.create(realModel as string, { + executionProviders: ['wasm'], + }); + expect(session.inputNames).toEqual(['left_image', 'right_image']); + expect(session.outputNames).toEqual(['disparity']); + const rig = await rigFromNpz(readFileSync(fixture('calibration.npz'))); + const png = PNG.sync.read(readFileSync(fixture('left.png'))); + const rect = computeRectification(rig, png.width, png.height, 960, 576); + expect(rect.width).toBe(960); + expect(rect.height).toBe(576); + }, 120000); +}); diff --git a/client/dive-common/use/stereo/useStereoOnnxTransfer.ts b/client/dive-common/use/stereo/useStereoOnnxTransfer.ts index 57f062c51..e53f5ebb7 100644 --- a/client/dive-common/use/stereo/useStereoOnnxTransfer.ts +++ b/client/dive-common/use/stereo/useStereoOnnxTransfer.ts @@ -1,8 +1,8 @@ /** * Client-side stereo transfer: when a detection is annotated on one camera, - * warp it onto the other camera using the VIAME "match" ONNX model - * ({@link StereoOnnxMatcher}) — no backend, so it works in both the web and - * desktop DIVE builds. + * warp it onto the other camera using the selected {@link StereoMatcher} — the + * VIAME "match" ONNX model or the Fast-FoundationStereo export — with no + * backend, so it works in both the web and desktop DIVE builds. * * This mirrors the desktop backend stereo handler (ViewerLoader's * `handleStereoAnnotationComplete`) but runs the correspondence search and the @@ -22,7 +22,7 @@ import type { StereoAnnotationCompleteParams } from '../useModeManager'; import type { SearchRange } from './StereoOnnxMatcher'; import type { StereoMatcher } from './stereoMatcher'; import { StereoRig, invertRig } from './calibration'; -import { rgbaToGray, RgbaImage } from './image'; +import { RgbaImage } from './image'; import { measureLine, aggregateLengths, StereoMeasurement } from './triangulate'; export interface StereoOnnxTransferConfig { @@ -261,26 +261,59 @@ export default function useStereoOnnxTransfer(config: StereoOnnxTransferConfig) return measurement; } + /** Identifies a (source camera -> other camera, frame) pair for matcher caches. */ + function frameKey(sourceCamera: string, otherCamera: string, frameNum: number) { + return `${sourceCamera}>${otherCamera}@${frameNum}`; + } + + /** Orient the rig so `sourceCamera` is the source ("left"). */ + function orientRig(rig0: StereoRig, sourceCamera: string) { + return sourceCamera === getLeftCameraName() ? rig0 : invertRig(rig0); + } + /** Run the correspondence search for one set of points, source camera -> other. */ async function warp(points: Point[], sourceCamera: string, otherCamera: string, frameNum: number) { const [rig0, matcher] = await Promise.all([getRig(), getMatcher()]); if (!rig0) throw new Error('No stereo calibration is available for this dataset.'); if (!matcher) throw new Error('The stereo matching model could not be loaded.'); - // Orient the rig so the annotated camera is the source ("left"). - const rig = sourceCamera === getLeftCameraName() ? rig0 : invertRig(rig0); + const rig = orientRig(rig0, sourceCamera); const [srcFrame, tgtFrame] = await Promise.all([ getFrame(sourceCamera, frameNum), getFrame(otherCamera, frameNum), ]); if (!srcFrame || !tgtFrame) throw new Error('Could not read the frame pixels for both cameras.'); - return matcher.warpPoints(points, rgbaToGray(srcFrame), rgbaToGray(tgtFrame), rig, { + return matcher.warpPoints(points, srcFrame, tgtFrame, rig, { range: getRange(), threshold: config.threshold, uniquenessRatio: config.uniquenessRatio, + frameKey: frameKey(sourceCamera, otherCamera, frameNum), }); } + /** + * Let a matcher that works per frame (the foundation method's dense + * disparity) compute its map for `frameNum` before the user draws anything + * there. Both warp directions are prepared, the calibration's left camera + * first. A no-op for matchers without per-frame state. `stillWanted` lets the + * host drop the work once the viewer has moved on to another frame. + */ + async function precomputeFrame(frameNum: number, stillWanted: () => boolean = () => true): Promise { + const cams = getMultiCamList(); + if (cams.length < 2) return; + const leftCamera = getLeftCameraName(); + const rightCamera = cams.find((c) => c !== leftCamera); + if (!rightCamera) return; + const [rig0, matcher] = await Promise.all([getRig(), getMatcher()]); + if (!rig0 || !matcher?.prepare || !stillWanted()) return; + const [leftFrame, rightFrame] = await Promise.all([ + getFrame(leftCamera, frameNum), getFrame(rightCamera, frameNum), + ]); + if (!leftFrame || !rightFrame || !stillWanted()) return; + await matcher.prepare(frameKey(leftCamera, rightCamera, frameNum), leftFrame, rightFrame, rig0, stillWanted); + await matcher.prepare(frameKey(rightCamera, leftCamera, frameNum), rightFrame, leftFrame, invertRig(rig0), stillWanted); + } + /** * Warp the just-completed annotation onto the other camera. * @@ -491,5 +524,6 @@ export default function useStereoOnnxTransfer(config: StereoOnnxTransferConfig) handleStereoTrackLinked, warpAllFromCamera, measureAtFrame, + precomputeFrame, }; } diff --git a/client/platform/web-girder/api/configuration.service.ts b/client/platform/web-girder/api/configuration.service.ts index 9be2d02f0..3dc0a0c65 100644 --- a/client/platform/web-girder/api/configuration.service.ts +++ b/client/platform/web-girder/api/configuration.service.ts @@ -11,6 +11,16 @@ export interface BrandData { trainingMessage?: string; } +/** The served Fast-FoundationStereo export, as `dive_configuration/stereo_foundation_model/spec` reports it. */ +export interface StereoFoundationModelSpec { + name: string; + url: string; + md5: string; + height: number; + width: number; + size: number; +} + export interface StatsResponse { table_stats: { datasets: number; // Total count of datasets @@ -89,6 +99,16 @@ function getAddons() { return girderRest.get('dive_configuration/addons'); } +function getStereoFoundationModelSpec() { + return girderRest.get('dive_configuration/stereo_foundation_model/spec'); +} + +function getStereoFoundationModel() { + return girderRest.get('dive_configuration/stereo_foundation_model', { + responseType: 'arraybuffer', + }); +} + function postAddons(urls: string[], forceDownload: boolean) { return girderRest.post(`dive_configuration/upgrade_pipelines?force=${forceDownload}`, urls); } @@ -116,6 +136,8 @@ export { getPipelineList, getTrainingConfigurations, getAddons, + getStereoFoundationModelSpec, + getStereoFoundationModel, postAddons, updateContainers, getStats, diff --git a/client/platform/web-girder/useStereoOnnxWeb.ts b/client/platform/web-girder/useStereoOnnxWeb.ts index 64305ac48..d238a3ac7 100644 --- a/client/platform/web-girder/useStereoOnnxWeb.ts +++ b/client/platform/web-girder/useStereoOnnxWeb.ts @@ -1,19 +1,28 @@ /** * Web wiring for client-side stereo transfer (warp a detection to the other - * camera and triangulate its length) using the VIAME "match" ONNX model. + * camera and triangulate its length) using the selected correspondence model. * Assembles the platform providers that {@link useStereoOnnxTransfer} needs: * - calibration, taken from the session's file stash when the user just * imported one and otherwise downloaded from the dataset's Girder folder, - * - the ONNX matcher (lazily created from a served model asset), + * - the matcher for the selected method (lazily created and cached), * - per-camera frame pixels, read from the GeoJS viewer for the frame on * screen and fetched from the frame's image URL for any other frame. * - * The exported model must be served as a static asset (default - * `/models/stereo_match.onnx`; produce it with - * `plugins/onnx/export_stereo_mapping.py --model match`). If no calibration or - * model is available the transfer reports the failure and no-ops. + * The NCC model is a static asset (`/models/stereo_match.onnx`, produced with + * `plugins/onnx/export_stereo_mapping.py --model match`). The foundation model + * is the export VIAME publishes in its FAST-FDN-STEREO add-on: the girder + * server resolves it from VIAME's add-on list and serves it, and the bytes are + * kept in the browser's Cache API keyed by the add-on's md5 so a page reload + * does not re-download ~100 MB. If no calibration or model is available the + * transfer reports the failure and no-ops. + * + * Because the foundation method costs one network pass per frame rather than + * per point, its disparity map is computed as soon as the viewer lands on a + * frame, so the warp itself is quick when the user draws. */ +import { watch } from 'vue'; + import { clientSettings } from 'dive-common/store/settings'; import useStereoOnnxTransfer from 'dive-common/use/stereo/useStereoOnnxTransfer'; import { StereoOnnxMatcher } from 'dive-common/use/stereo/StereoOnnxMatcher'; @@ -31,16 +40,8 @@ import type { StereoMeasurement } from 'dive-common/use/stereo/triangulate'; import { getCalibrationFile, getLastCalibration } from './multicamFileRegistry'; const DEFAULT_MODEL_URL = '/models/stereo_match.onnx'; -/** - * Fast-FoundationStereo is opt-in and unbundled: the exports run ~100 MB, so - * unlike the NCC graph this one is not committed. Serve an export here (or pass - * `foundationModelUrl`) and give its sidecar `image_size` as - * `foundationModelSpec`; with no model served the dropdown's foundation option - * reports that it could not load and the warp no-ops, exactly as a missing - * calibration does. - */ -const DEFAULT_FOUNDATION_MODEL_URL = '/models/stereo_foundation.onnx'; -const DEFAULT_FOUNDATION_SPEC: FoundationModelSpec = { height: 576, width: 960 }; +/** Browser cache holding the foundation model bytes across page loads. */ +const MODEL_CACHE_NAME = 'dive-stereo-models'; // Mirrors epipolar_min_disparity / epipolar_max_disparity in VIAME's // configs/pipelines/interactive_stereo_template.conf, which is what the desktop // interactive stereo service loads. Scene-dependent, and hidden config there @@ -55,8 +56,12 @@ export interface StereoOnnxWebOptions { /** Dataset (folder) id used to look up the stored calibration. */ getDatasetId: () => string; modelUrl?: string; + /** + * Serve the foundation model from a fixed URL instead of the girder + * endpoint; `foundationModelSpec` (the export's sidecar `image_size`) is then + * required. + */ foundationModelUrl?: string; - /** Input resolution of the foundation export (its sidecar yaml `image_size`). */ foundationModelSpec?: FoundationModelSpec; /** Overrides the user's dropdown choice; mainly for tests. */ getMatchMethod?: () => StereoMatchMethod; @@ -97,13 +102,53 @@ async function urlToRgba(url: string): Promise { } } +async function openModelCache(): Promise { + if (typeof caches === 'undefined') return null; + try { + return await caches.open(MODEL_CACHE_NAME); + } catch { + return null; + } +} + +/** + * The foundation model bytes for the export the server currently serves. The + * md5 comes from VIAME's add-on list, so a re-published export changes the + * cache key and the stale copy is dropped. + */ +async function fetchFoundationModel(): Promise<{ bytes: ArrayBuffer; spec: FoundationModelSpec }> { + // Imported lazily: the girder client touches `window` at load time, which + // breaks node-environment unit tests that import this file. + const { getStereoFoundationModelSpec, getStereoFoundationModel } = await import( + 'platform/web-girder/api/configuration.service' + ); + const { data: spec } = await getStereoFoundationModelSpec(); + const cacheUrl = `${window.location.origin}/dive-stereo-models/${spec.md5}/${spec.name}`; + const cache = await openModelCache(); + if (cache) { + const hit = await cache.match(cacheUrl); + if (hit) return { bytes: await hit.arrayBuffer(), spec }; + } + const { data: bytes } = await getStereoFoundationModel(); + if (cache) { + try { + const keys = await cache.keys(); + await Promise.all(keys.map((request) => cache.delete(request))); + await cache.put(cacheUrl, new Response(bytes, { + headers: { 'Content-Type': 'application/octet-stream' }, + })); + } catch (err) { + console.warn('[StereoOnnx] could not cache the foundation model', err); + } + } + return { bytes, spec }; +} + export default function useStereoOnnxWeb(opts: StereoOnnxWebOptions) { const modelUrl = opts.modelUrl ?? DEFAULT_MODEL_URL; - const foundationModelUrl = opts.foundationModelUrl ?? DEFAULT_FOUNDATION_MODEL_URL; - const foundationSpec = opts.foundationModelSpec ?? DEFAULT_FOUNDATION_SPEC; // Cached per method: switching the dropdown must not reload the other model, // and a method that failed to load must not be retried on every warp. - const matchers: Partial> = {}; + const matchers: Partial>> = {}; let rig: StereoRig | null = null; let rigKey: string | null = null; @@ -112,19 +157,38 @@ export default function useStereoOnnxWeb(opts: StereoOnnxWebOptions) { return clientSettings.stereoSettings.matchMethod ?? DEFAULT_STEREO_MATCH_METHOD; } - async function getMatcher(): Promise { - const method = currentMethod(); - if (method in matchers) return matchers[method] ?? null; - const url = method === 'foundation' ? foundationModelUrl : modelUrl; + async function createFoundationMatcher(): Promise { + if (opts.foundationModelUrl) { + if (!opts.foundationModelSpec) { + throw new Error('foundationModelSpec is required with foundationModelUrl'); + } + return StereoFoundationMatcher.create(opts.foundationModelUrl, opts.foundationModelSpec); + } + opts.onStatus?.('Loading the stereo model (about 100 MB on first use)...'); try { - matchers[method] = method === 'foundation' - ? await StereoFoundationMatcher.create(url, foundationSpec) - : await StereoOnnxMatcher.create(url); - } catch (err) { - console.warn('[StereoOnnx] failed to load model', method, url, err); - matchers[method] = null; + const { bytes, spec } = await fetchFoundationModel(); + return await StereoFoundationMatcher.create(new Uint8Array(bytes), { + height: spec.height, width: spec.width, + }); + } finally { + opts.onStatus?.(null); } - return matchers[method] ?? null; + } + + function getMatcher(): Promise { + const method = currentMethod(); + const existing = matchers[method]; + if (existing) return existing; + const created = (method === 'foundation' + ? createFoundationMatcher() + : StereoOnnxMatcher.create(modelUrl) + ).catch((err) => { + console.warn('[StereoOnnx] failed to load model', method, err); + opts.onError?.(`The stereo matching model could not be loaded. ${(err as Error).message ?? err}`); + return null; + }); + matchers[method] = created; + return created; } function parseRig(name: string, buffer: ArrayBuffer): Promise { @@ -208,6 +272,18 @@ export default function useStereoOnnxWeb(opts: StereoOnnxWebOptions) { return url ? urlToRgba(url) : null; } + /** The frame the viewer is on, once its media has loaded. */ + function currentFrame(): number | undefined { + const viewer = opts.getViewer(); + if (!viewer?.progress?.loaded) return undefined; + try { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + return fromViewer(viewer.aggregateController)?.frame?.value; + } catch { + return undefined; + } + } + // ViewerLoader is reused across /viewer/:id navigations while // remounts, so never close over a specific Viewer — rebuild when cameraStore // identity changes, and resolve multiCamList via getViewer() each call. @@ -240,6 +316,41 @@ export default function useStereoOnnxWeb(opts: StereoOnnxWebOptions) { return transfer; } + /** + * Only the foundation method has per-frame work to do ahead of time, and only + * when warps will actually happen (auto-compute on, two cameras). + */ + function precomputeWanted(): boolean { + return currentMethod() === 'foundation' + && clientSettings.stereoSettings.autoComputeOtherCamera + && (fromViewer(opts.getViewer()?.multiCamList) ?? []).length >= 2; + } + + // A precompute failure is the model failing on this browser/GPU, which the + // user should hear about once rather than on every frame change. + let precomputeErrorReported = false; + + /** Compute the current frame's disparity maps ahead of any warp there. */ + function precomputeCurrentFrame() { + const frame = currentFrame(); + if (frame === undefined || !precomputeWanted()) return; + getTransfer()?.precomputeFrame(frame, () => currentFrame() === frame && precomputeWanted()) + .catch((err) => { + console.warn('[StereoOnnx] disparity precompute failed', err); + if (!precomputeErrorReported) { + precomputeErrorReported = true; + opts.onError?.(`The higher-accuracy stereo model could not run in this browser. ${(err as Error).message ?? err}`); + } + }); + } + + watch( + () => [currentFrame(), precomputeWanted()] as const, + ([frame, wanted]) => { + if (frame !== undefined && wanted) precomputeCurrentFrame(); + }, + ); + type Transfer = ReturnType; async function handleStereoAnnotationComplete( @@ -256,5 +367,10 @@ export default function useStereoOnnxWeb(opts: StereoOnnxWebOptions) { return getTransfer()?.warpAllFromCamera(cameraName); } - return { handleStereoAnnotationComplete, handleStereoTrackLinked, warpAllFromCamera }; + return { + handleStereoAnnotationComplete, + handleStereoTrackLinked, + warpAllFromCamera, + precomputeCurrentFrame, + }; } diff --git a/docker-compose.yml b/docker-compose.yml index d959b283d..293b0c5a2 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -94,6 +94,9 @@ services: - memcached volumes: - girder_assetstore:/home/assetstore + # Models the girder service fetches on demand (e.g. the stereo model the + # web client runs), kept across container recreation. + - girder_model_cache:/tmp/dive_models # Bind mount nocopy public data - type: bind source: ${PUBLIC_DATA_PATH:-./docker/girder_data} @@ -108,6 +111,7 @@ services: - "GIRDER_SETTING_WORKER_API_URL=${GIRDER_SETTING_WORKER_API_URL:-http://girder:8080/api/v1}" - "GIRDER_STATIC_ROOT_DIR=/opt/dive/clients/girder" - "GIRDER_NOTIFICATION_REDIS_URL=redis://redis:6379" + - "DIVE_MODEL_CACHE_DIR=/tmp/dive_models" # Rabbitmq management variables - "RABBITMQ_MANAGEMENT_USERNAME=${RABBITMQ_MANAGEMENT_USERNAME:-guest}" @@ -233,4 +237,5 @@ volumes: addons: mongo_db: girder_assetstore: + girder_model_cache: redis_data: diff --git a/server/dive_server/views_configuration.py b/server/dive_server/views_configuration.py index 5001bab04..23de9a545 100644 --- a/server/dive_server/views_configuration.py +++ b/server/dive_server/views_configuration.py @@ -6,7 +6,7 @@ from girder.api import access from girder.api.describe import Description, autoDescribeRoute -from girder.api.rest import Resource +from girder.api.rest import Resource, setRawResponse, setResponseHeader from girder.constants import AccessType from girder.exceptions import RestException from girder.models.folder import Folder @@ -19,7 +19,7 @@ from dive_server import crud, crud_rpc, worker_capabilities from dive_tasks import tasks -from dive_utils import TRUTHY_META_VALUES, constants, models, types +from dive_utils import TRUTHY_META_VALUES, constants, models, stereo_models, types @setting_utilities.validator({constants.SETTINGS_CONST_JOBS_CONFIGS}) @@ -67,6 +67,10 @@ def __init__(self, resourceName): self.route("GET", ("brand_data",), self.get_brand_data) self.route("GET", ("pipelines",), self.get_pipelines) self.route("GET", ("training_configs",), self.get_training_configs) + self.route("GET", ("stereo_foundation_model",), self.get_stereo_foundation_model) + self.route( + "GET", ("stereo_foundation_model", "spec"), self.get_stereo_foundation_model_spec + ) self.route("PUT", ("brand_data",), self.update_brand_data) self.route("PUT", ("static_pipeline_configs",), self.update_static_pipeline_configs) @@ -118,6 +122,49 @@ def get_training_configs(self, params): } return training_configs + @staticmethod + def _stereo_foundation_model() -> stereo_models.FoundationModel: + try: + return stereo_models.ensure_stereo_foundation_model() + except stereo_models.ModelUnavailable as exc: + raise RestException(str(exc), code=502) + + @access.user + @autoDescribeRoute( + Description( + "Describe the Fast-FoundationStereo ONNX export served by " + "stereo_foundation_model, fetching it from the VIAME add-on list on first use" + ) + ) + def get_stereo_foundation_model_spec(self, params): + model = self._stereo_foundation_model() + return { + 'name': model.onnx_path.name, + 'url': model.url, + 'md5': model.md5, + 'height': model.height, + 'width': model.width, + 'size': model.onnx_path.stat().st_size, + } + + @access.user + @autoDescribeRoute(Description("Download the Fast-FoundationStereo ONNX export")) + def get_stereo_foundation_model(self, params): + model = self._stereo_foundation_model() + setResponseHeader('Content-Type', 'application/octet-stream') + setResponseHeader('Content-Length', str(model.onnx_path.stat().st_size)) + setResponseHeader('Content-Disposition', f'attachment; filename="{model.onnx_path.name}"') + setRawResponse() + + def stream(): + with open(model.onnx_path, 'rb') as handle: + chunk = handle.read(stereo_models.DOWNLOAD_CHUNK_BYTES) + while chunk: + yield chunk + chunk = handle.read(stereo_models.DOWNLOAD_CHUNK_BYTES) + + return stream() + @access.admin @autoDescribeRoute( Description("update brand data").jsonParam( diff --git a/server/dive_utils/stereo_models.py b/server/dive_utils/stereo_models.py new file mode 100644 index 000000000..d8a97c31b --- /dev/null +++ b/server/dive_utils/stereo_models.py @@ -0,0 +1,191 @@ +""" +Serve the Fast-FoundationStereo ONNX export to the web client. + +The export ships inside VIAME's ``FAST-FDN-STEREO`` add-on zip. Its download +URL and md5 are read from VIAME's ``download_viame_addons.csv`` (the same list +the add-on installer uses) rather than pinned here, so a re-published model is +picked up without a DIVE release. The zip is fetched once per md5 into a local +cache and only the model and its sidecar yaml are kept. +""" + +import csv +import fcntl +import hashlib +import os +from pathlib import Path +import re +import shutil +import tempfile +from typing import Callable, Iterable, List, NamedTuple, Optional +import zipfile + +import requests + +from dive_utils import constants + +STEREO_FOUNDATION_ADDON = 'FAST-FDN-STEREO' +MODEL_CACHE_DIR_ENV = 'DIVE_MODEL_CACHE_DIR' +DEFAULT_MODEL_CACHE_DIR = '/tmp/dive_models' +DOWNLOAD_CHUNK_BYTES = 1 << 20 +DOWNLOAD_TIMEOUT_SECONDS = 60 + + +class AddonSource(NamedTuple): + name: str + url: str + md5: str + + +class FoundationModel(NamedTuple): + onnx_path: Path + yaml_path: Path + url: str + md5: str + height: int + width: int + + +class ModelUnavailable(Exception): + """The model could not be resolved, downloaded or verified.""" + + +def parse_addon_rows(text: str) -> List[AddonSource]: + rows = [] + for item in csv.reader(text.splitlines(), delimiter=','): + if len(item) < 4: + continue + rows.append(AddonSource(item[0].strip(), item[1].strip(), item[3].strip().lower())) + return rows + + +def find_addon(rows: Iterable[AddonSource], name: str) -> Optional[AddonSource]: + return next((row for row in rows if row.name == name), None) + + +def resolve_addon(name: str = STEREO_FOUNDATION_ADDON) -> AddonSource: + try: + response = requests.get(constants.AddonsListURL, timeout=DOWNLOAD_TIMEOUT_SECONDS) + response.raise_for_status() + except requests.RequestException as exc: + raise ModelUnavailable(f'Could not read the VIAME add-on list: {exc}') from exc + addon = find_addon(parse_addon_rows(response.content.decode('utf-8')), name) + if addon is None or not addon.url: + raise ModelUnavailable(f'The VIAME add-on list has no {name} entry') + return addon + + +def parse_image_size(yaml_text: str) -> Optional[tuple]: + """``image_size: [H, W]`` from the export's sidecar, in block or flow form.""" + match = re.search(r'image_size:\s*\[\s*(\d+)\s*,\s*(\d+)\s*\]', yaml_text) or re.search( + r'image_size:\s*\n\s*-\s*(\d+)\s*\n\s*-\s*(\d+)', yaml_text + ) + if not match: + return None + return int(match.group(1)), int(match.group(2)) + + +def model_cache_dir() -> Path: + return Path(os.environ.get(MODEL_CACHE_DIR_ENV, DEFAULT_MODEL_CACHE_DIR)) + + +def _download(url: str, dest: Path) -> str: + digest = hashlib.md5() + with requests.get(url, stream=True, timeout=DOWNLOAD_TIMEOUT_SECONDS) as response: + response.raise_for_status() + with open(dest, 'wb') as out: + for chunk in response.iter_content(DOWNLOAD_CHUNK_BYTES): + out.write(chunk) + digest.update(chunk) + return digest.hexdigest() + + +def extract_model(zip_path: Path, dest_dir: Path) -> FoundationModel: + """Pull the single ``.onnx`` and its sidecar ``.yaml`` out of an add-on zip.""" + with zipfile.ZipFile(zip_path) as archive: + names = archive.namelist() + onnx_names = [n for n in names if n.lower().endswith('.onnx')] + if len(onnx_names) != 1: + raise ModelUnavailable(f'Expected one .onnx in the add-on, found {len(onnx_names)}') + onnx_name = onnx_names[0] + yaml_name = os.path.splitext(onnx_name)[0] + '.yaml' + if yaml_name not in names: + raise ModelUnavailable(f'The add-on has no sidecar {os.path.basename(yaml_name)}') + dest_dir.mkdir(parents=True, exist_ok=True) + targets = {} + for member in (onnx_name, yaml_name): + target = dest_dir / os.path.basename(member) + with archive.open(member) as src, open(target, 'wb') as out: + shutil.copyfileobj(src, out) + targets[member] = target + return _describe(targets[onnx_name], targets[yaml_name], url='', md5='') + + +def _describe(onnx_path: Path, yaml_path: Path, url: str, md5: str) -> FoundationModel: + size = parse_image_size(yaml_path.read_text()) + if size is None: + raise ModelUnavailable(f'{yaml_path.name} does not declare image_size') + return FoundationModel(onnx_path, yaml_path, url, md5, size[0], size[1]) + + +def _cached(addon: AddonSource, cache_dir: Path) -> Optional[FoundationModel]: + model_dir = cache_dir / addon.name / addon.md5 + if not model_dir.is_dir(): + return None + onnx_files = list(model_dir.glob('*.onnx')) + if len(onnx_files) != 1: + return None + yaml_path = onnx_files[0].with_suffix('.yaml') + if not yaml_path.is_file(): + return None + return _describe(onnx_files[0], yaml_path, addon.url, addon.md5) + + +def ensure_model( + addon: AddonSource, + cache_dir: Optional[Path] = None, + download: Callable[[str, Path], str] = _download, +) -> FoundationModel: + """ + The cached model for ``addon``, downloading and verifying it first when the + cache holds nothing for its md5. Older md5 directories are dropped once a + newer one is in place. Safe across processes sharing the cache directory. + """ + cache_dir = cache_dir or model_cache_dir() + addon_dir = cache_dir / addon.name + addon_dir.mkdir(parents=True, exist_ok=True) + with open(addon_dir / '.lock', 'w') as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + try: + cached = _cached(addon, cache_dir) + if cached is not None: + return cached + with tempfile.TemporaryDirectory(dir=addon_dir) as tmp: + zip_path = Path(tmp) / 'addon.zip' + try: + actual_md5 = download(addon.url, zip_path) + except requests.RequestException as exc: + raise ModelUnavailable(f'Could not download {addon.url}: {exc}') from exc + if addon.md5 and actual_md5 != addon.md5: + raise ModelUnavailable( + f'{addon.name} download did not match the add-on list md5 ' + f'({actual_md5} != {addon.md5})' + ) + staging = Path(tmp) / 'model' + extract_model(zip_path, staging) + final_dir = addon_dir / addon.md5 + if final_dir.exists(): + shutil.rmtree(final_dir) + shutil.move(str(staging), str(final_dir)) + for stale in addon_dir.iterdir(): + if stale.is_dir() and stale.name != addon.md5: + shutil.rmtree(stale, ignore_errors=True) + cached = _cached(addon, cache_dir) + if cached is None: + raise ModelUnavailable(f'{addon.name} was downloaded but could not be read back') + return cached + finally: + fcntl.flock(lock, fcntl.LOCK_UN) + + +def ensure_stereo_foundation_model() -> FoundationModel: + return ensure_model(resolve_addon(STEREO_FOUNDATION_ADDON)) diff --git a/server/tests/test_stereo_models.py b/server/tests/test_stereo_models.py new file mode 100644 index 000000000..0e84f01dc --- /dev/null +++ b/server/tests/test_stereo_models.py @@ -0,0 +1,105 @@ +import hashlib +import io +from pathlib import Path +import zipfile + +import pytest + +from dive_utils import stereo_models + +CSV = ( + 'DEFAULT-FISH, https://example.com/fish/download, Default fish, ' + '1b71862b7fa39def315c0c08c44048a7, ALL-PLATFORMS, "PYTORCH, ONNX", models/fish.zip\n' + 'FAST-FDN-STEREO, https://example.com/stereo/download, Fast foundation stereo, ' + '29D4CDE2E33500E74844D07C5AB3DEE0, ALL-PLATFORMS, "PYTORCH", \n' +) + +YAML = 'image_size:\n- 576\n- 960\nvalid_iters: 8\n' + + +def make_addon_zip(yaml_text=YAML, extra_onnx=False) -> bytes: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, 'w') as archive: + archive.writestr('configs/pipelines/interactive_stereo_default.conf', 'include x.conf\n') + archive.writestr('configs/pipelines/models/fast_foundation_stereo_l.onnx', b'onnx-bytes') + archive.writestr('configs/pipelines/models/fast_foundation_stereo_l.yaml', yaml_text) + if extra_onnx: + archive.writestr('configs/pipelines/models/other.onnx', b'x') + return buffer.getvalue() + + +def fake_downloader(payload: bytes): + def download(url: str, dest: Path) -> str: + dest.write_bytes(payload) + return hashlib.md5(payload).hexdigest() + + return download + + +def test_parse_addon_rows_strips_whitespace_and_lowercases_md5(): + rows = stereo_models.parse_addon_rows(CSV) + stereo = stereo_models.find_addon(rows, 'FAST-FDN-STEREO') + assert stereo == stereo_models.AddonSource( + 'FAST-FDN-STEREO', + 'https://example.com/stereo/download', + '29d4cde2e33500e74844d07c5ab3dee0', + ) + assert stereo_models.find_addon(rows, 'MISSING') is None + + +def test_parse_image_size_block_and_flow(): + assert stereo_models.parse_image_size(YAML) == (576, 960) + assert stereo_models.parse_image_size('image_size: [320, 736]\n') == (320, 736) + assert stereo_models.parse_image_size('valid_iters: 8\n') is None + + +def test_ensure_model_downloads_once_and_verifies_md5(tmp_path): + payload = make_addon_zip() + addon = stereo_models.AddonSource( + 'FAST-FDN-STEREO', 'https://example.com/stereo', hashlib.md5(payload).hexdigest() + ) + calls = [] + + def download(url, dest): + calls.append(url) + return fake_downloader(payload)(url, dest) + + model = stereo_models.ensure_model(addon, tmp_path, download) + assert model.onnx_path.read_bytes() == b'onnx-bytes' + assert (model.height, model.width) == (576, 960) + assert model.md5 == addon.md5 + assert model.onnx_path.parent == tmp_path / addon.name / addon.md5 + + again = stereo_models.ensure_model(addon, tmp_path, download) + assert again.onnx_path == model.onnx_path + assert calls == ['https://example.com/stereo'] + + +def test_ensure_model_rejects_md5_mismatch(tmp_path): + payload = make_addon_zip() + addon = stereo_models.AddonSource('FAST-FDN-STEREO', 'https://example.com/stereo', 'f' * 32) + with pytest.raises(stereo_models.ModelUnavailable): + stereo_models.ensure_model(addon, tmp_path, fake_downloader(payload)) + assert not (tmp_path / addon.name / addon.md5).exists() + + +def test_ensure_model_replaces_previous_md5(tmp_path): + old_payload = make_addon_zip(yaml_text='image_size: [320, 736]\n') + new_payload = make_addon_zip() + old = stereo_models.AddonSource( + 'FAST-FDN-STEREO', 'https://example.com/old', hashlib.md5(old_payload).hexdigest() + ) + new = stereo_models.AddonSource( + 'FAST-FDN-STEREO', 'https://example.com/new', hashlib.md5(new_payload).hexdigest() + ) + stereo_models.ensure_model(old, tmp_path, fake_downloader(old_payload)) + model = stereo_models.ensure_model(new, tmp_path, fake_downloader(new_payload)) + assert (model.height, model.width) == (576, 960) + assert not (tmp_path / old.name / old.md5).exists() + + +def test_extract_model_requires_exactly_one_onnx(tmp_path): + zip_path = tmp_path / 'addon.zip' + zip_path.write_bytes(make_addon_zip(extra_onnx=True)) + with pytest.raises(stereo_models.ModelUnavailable): + stereo_models.extract_model(zip_path, tmp_path / 'out') From 81e792cc781733e6f6048efe46cc8743736aba04 Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Mon, 14 Sep 2026 19:13:46 -0400 Subject: [PATCH 3/9] Run the foundation stereo model on the native WebGPU provider at basic optimisation JSEP returns an all-zero cost volume for this graph and the extended fusions corrupt the GRU gates; verified probe-by-probe against CPU onnxruntime: this configuration matches to 5e-5 px. Requires VIAME's web export of the model (export_fast_foundation_stereo_web.py). Claude-Session: https://claude.ai/code/session_013oFD6dP2EQY3R9UdXzx7dL --- client/dive-common/use/stereo/README.md | 63 ++++++++++--------- .../use/stereo/StereoFoundationMatcher.ts | 50 ++++++++++++--- .../use/stereo/tests/stereoFoundation.spec.ts | 4 +- client/vite.config.ts | 2 +- 4 files changed, 79 insertions(+), 40 deletions(-) diff --git a/client/dive-common/use/stereo/README.md b/client/dive-common/use/stereo/README.md index f71412c07..40c66b439 100644 --- a/client/dive-common/use/stereo/README.md +++ b/client/dive-common/use/stereo/README.md @@ -141,7 +141,8 @@ it needs a GPU. The export is the one VIAME publishes in its `FAST-FDN-STEREO` add-on (`fast_foundation_stereo_l.onnx` + sidecar `.yaml`), the same file -`plugins/onnx/fast_foundation_stereo.py` runs server-side. Nothing is pinned +`plugins/onnx/fast_foundation_stereo.py` runs server-side (see *Runtime +requirements* for why it has to be VIAME's web build of it). Nothing is pinned in DIVE: 1. The girder server reads the add-on's URL and md5 from VIAME's @@ -161,30 +162,35 @@ foundationModelSpec })` bypasses the server. ### Runtime requirements -The export runs only on onnxruntime-web's **WebGPU** provider: the CPU (wasm) -path needs ~7 GB of activations, past the 4 GB a wasm heap can address, so the -matcher refuses to start without `navigator.gpu` and says so. An inference -failure (unsupported operator, out of GPU memory) is remembered by the matcher -and surfaced once, rather than retried on every frame change. - -**As of onnxruntime-web 1.27–1.31 the add-on's current export does not run on -WebGPU either**, checked in headless Chrome against the stock export: - -| Provider | Stops at | -| --- | --- | -| WebGPU (JSEP, the default `onnxruntime-web` bundle) | `Conv` 3D with asymmetric padding (`0,1,1,0,1,1`); after rewriting those into `Pad` + `Conv`, `ConvTranspose` 3D ("only support 2-dimensional conv") | -| WebGPU (native EP, `onnxruntime-web/webgpu`) | a 48-input `Concat` ("Too many storage buffers in shader: 11, max 10") | - -The graph also carries several `[1, 8, 28, 48, 144, 240]` fp32 intermediates -(1.5 GB each) and peaks near 16 GB on the CUDA provider, so even with operator -coverage it needs a browser-exposed GPU with a very large buffer budget. A -web-viable export therefore needs to come from the VIAME side: a smaller -`image_size` / fewer refinement iterations to shrink the cost volume, and 3D -cost-aggregation ops expressed in a form onnxruntime-web supports (or a -`Pad`-rewrite plus `ConvTranspose` decomposition, and `Concat` split into ≤8 -inputs for the native EP). The client and server here already take whatever -the `FAST-FDN-STEREO` add-on ships, so a re-published export needs no DIVE -change beyond the CSV md5 it carries. +The export runs only on a GPU: the CPU (wasm) path needs several GB of +activations, past the 4 GB a wasm heap can address, so the matcher refuses to +start without `navigator.gpu` and says so. An inference failure (unsupported +operator, out of GPU memory) is remembered by the matcher and surfaced once, +rather than retried on every frame change. + +Two runtime choices are deliberate, both verified probe-by-probe against CPU +onnxruntime on the fixture pair (headless Chrome, SwiftShader adapter): + +- **The native WebGPU provider** (`onnxruntime-web/webgpu`, imported lazily), + not the default bundle's JSEP kernels: JSEP returns an all-zero cost volume + for this graph (a Reshape/Cast of the right-camera features reads wrong data) + on 1.27, 1.29 and 1.31-dev alike. +- **`graphOptimizationLevel: 'basic'`**: with `'all'`, one of the provider's + extended-level fusions corrupts the GRU gate convolutions and the disparity + drifts by ~2 px; at `'basic'` (and `'disabled'`) the browser output matches + CPU to 5e-5 px. + +The **export itself must be VIAME's web build** of the model +(`plugins/onnx/export_fast_foundation_stereo_web.py`, the file the +`FAST-FDN-STEREO` add-on ships): NVIDIA's stock single-file export materialises +1.5 GB correlation tensors and uses 3-D `ConvTranspose` and asymmetric 3-D +`Conv` padding that no onnxruntime-web provider runs. The web build is +numerically identical (~1e-4 px on CPU) and also cuts CPU peak memory from +16 GB to 3 GB. + +Real-GPU timing has not been measured (this machine's Chrome only exposes the +SwiftShader software adapter); `Desktop/Active/fast-fdn-stereo-web/webgpu_harness` +is a self-contained page that reports it. ### How it works @@ -237,8 +243,9 @@ has already left — unless a warp is waiting on it, which upgrades it. Python with this preprocessing puts the fixture's head/tail disparities within 1 px of the NCC reference; the padding rewrite leaves the output bit-identical. -- **Not runnable yet**: the WebGPU pass with the add-on's current export (see - *Runtime requirements*), so the end-to-end warp, the settings dropdown and - the frame watcher have not been exercised in a running viewer. Set +- **Verified in a browser (SwiftShader)**: the web export on the native WebGPU + provider at `basic` optimisation matches CPU onnxruntime to 5e-5 px on the + fixture pair. **Not yet exercised**: a real GPU adapter (timing), and the + end-to-end warp, settings dropdown and frame watcher in a running viewer. Set `DIVE_STEREO_FOUNDATION_MODEL` to an export to have the Node suite check its I/O contract. diff --git a/client/dive-common/use/stereo/StereoFoundationMatcher.ts b/client/dive-common/use/stereo/StereoFoundationMatcher.ts index 9937c016a..df34b232f 100644 --- a/client/dive-common/use/stereo/StereoFoundationMatcher.ts +++ b/client/dive-common/use/stereo/StereoFoundationMatcher.ts @@ -31,6 +31,17 @@ import { Rectification, computeRectification, rectifyPoint, rectifyMapper, unrectifyPoint, } from './rectify'; +/** + * The foundation model runs on onnxruntime-web's *native* WebGPU provider + * (`onnxruntime-web/webgpu`), not the default bundle's JSEP kernels: the JSEP + * provider returns an all-zero cost volume for this graph (verified probe by + * probe against CPU), while the native provider matches CPU through the + * refinement stage. The bundle is imported lazily so Node tests and pages that + * never select the method do not load it, and feeds are built with that + * bundle's own Tensor class because onnxruntime-common checks `instanceof`. + */ +type OrtModule = typeof ort; + /** * Half-width of the window whose disparities are pooled for one point. * @@ -65,6 +76,12 @@ export interface DisparitySession { run(feeds: Record): Promise>; } +/** A model input: NCHW float32 planes plus dims, wrapped into a Tensor by the matcher. */ +export interface InputPlanes { + data: Float32Array; + dims: [number, number, number, number]; +} + export interface FoundationMatcherOptions { /** * Defaults to WebGPU, the only provider that can run this export in a @@ -120,7 +137,7 @@ export function remapToInputTensor( mapY: Float32Array, width: number, height: number, -): ort.Tensor { +): InputPlanes { const plane = width * height; const out = new Float32Array(plane * 3); const { data, width: sw, height: sh } = src; @@ -157,7 +174,7 @@ export function remapToInputTensor( out[c * plane + i] = (rgb[c] / 255 - IMAGENET_MEAN[c]) / IMAGENET_STD[c]; } } - return new ort.Tensor('float32', out, [1, 3, height, width]); + return { data: out, dims: [1, 3, height, width] }; } interface Geometry { @@ -176,6 +193,9 @@ interface PendingDisparity { export class StereoFoundationMatcher implements StereoMatcher { private session: DisparitySession; + /** The runtime whose Tensor class the session accepts. */ + private ort: OrtModule; + private spec: FoundationModelSpec; private cacheSize: number; @@ -198,10 +218,16 @@ export class StereoFoundationMatcher implements StereoMatcher { */ private failure: Error | null = null; - constructor(session: DisparitySession, spec: FoundationModelSpec, cacheSize = DEFAULT_DISPARITY_CACHE_SIZE) { + constructor( + session: DisparitySession, + spec: FoundationModelSpec, + cacheSize = DEFAULT_DISPARITY_CACHE_SIZE, + runtime: OrtModule = ort, + ) { this.session = session; this.spec = spec; this.cacheSize = cacheSize; + this.ort = runtime; } /** @@ -214,12 +240,16 @@ export class StereoFoundationMatcher implements StereoMatcher { opts: FoundationMatcherOptions = {}, ): Promise { const executionProviders = opts.executionProviders ?? defaultExecutionProviders(); - ort.env.wasm.proxy = false; - const session = await ort.InferenceSession.create(model as string, { + // eslint-disable-next-line import/no-unresolved + const runtime = (await import('onnxruntime-web/webgpu')) as unknown as OrtModule; + runtime.env.wasm.proxy = false; + const session = await runtime.InferenceSession.create(model as string, { executionProviders, - graphOptimizationLevel: 'all', + // 'basic' (level 1) is exact against CPU; one of the extended-level + // fusions the WebGPU provider applies corrupts the GRU gates (~2 px). + graphOptimizationLevel: 'basic', }); - return new StereoFoundationMatcher(session, spec, opts.cacheSize); + return new StereoFoundationMatcher(session, spec, opts.cacheSize, runtime); } get inputSize(): FoundationModelSpec { @@ -244,11 +274,13 @@ export class StereoFoundationMatcher implements StereoMatcher { if (this.failure) throw this.failure; const { src, tgt } = this.geometry(rig, source); const { width, height } = this.spec; + const left = remapToInputTensor(source, src.mapX, src.mapY, width, height); + const right = remapToInputTensor(target, tgt.mapX, tgt.mapY, width, height); let out: Record; try { out = await this.session.run({ - left_image: remapToInputTensor(source, src.mapX, src.mapY, width, height), - right_image: remapToInputTensor(target, tgt.mapX, tgt.mapY, width, height), + left_image: new this.ort.Tensor('float32', left.data, left.dims), + right_image: new this.ort.Tensor('float32', right.data, right.dims), }); } catch (err) { this.failure = err instanceof Error ? err : new Error(String(err)); diff --git a/client/dive-common/use/stereo/tests/stereoFoundation.spec.ts b/client/dive-common/use/stereo/tests/stereoFoundation.spec.ts index 6a0b6c046..1c791f5d8 100644 --- a/client/dive-common/use/stereo/tests/stereoFoundation.spec.ts +++ b/client/dive-common/use/stereo/tests/stereoFoundation.spec.ts @@ -74,7 +74,7 @@ describe('remapToInputTensor', () => { const mapY = Float32Array.from([1, 2, 1]); const t = remapToInputTensor(img, mapX, mapY, 3, 1); expect(t.dims).toEqual([1, 3, 1, 3]); - const d = t.data as Float32Array; + const d = t.data; expect(d[0]).toBeCloseTo((1 - 0.485) / 0.229, 4); expect(d[3]).toBeCloseTo((128 / 255 - 0.456) / 0.224, 4); expect(d[6]).toBeCloseTo((0 - 0.406) / 0.225, 4); @@ -85,7 +85,7 @@ describe('remapToInputTensor', () => { it('accepts a grayscale frame by replicating it across channels', () => { const gray = { data: Float32Array.from([200, 200, 200, 200]), width: 2, height: 2 }; const t = remapToInputTensor(gray, Float32Array.from([0.5]), Float32Array.from([0.5]), 1, 1); - const d = t.data as Float32Array; + const d = t.data; expect(d[0]).toBeCloseTo((200 / 255 - 0.485) / 0.229, 4); expect(d[1]).toBeCloseTo((200 / 255 - 0.456) / 0.224, 4); }); diff --git a/client/vite.config.ts b/client/vite.config.ts index 50aa04fe2..4f11969c2 100644 --- a/client/vite.config.ts +++ b/client/vite.config.ts @@ -102,7 +102,7 @@ export default defineConfig(({ mode }) => { // answers with the SPA HTML fallback, so the runtime fails to instantiate. // Served unbundled, the .wasm sits next to the module and loads. The // production build already emits it as a hashed asset either way. - exclude: ['onnxruntime-web'], + exclude: ['onnxruntime-web', 'onnxruntime-web/webgpu'], }, build: { sourcemap: true, From 2cb9f9d53edce4800c5946d84a953b40a5bda2bb Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Mon, 14 Sep 2026 19:36:33 -0400 Subject: [PATCH 4/9] Record the measured browser performance of the foundation stereo export RTX 4090 through Chrome, native WebGPU provider at basic optimisation: 0.69 s per pair at 576x960, 0.26 s at 320x736, matching CPU onnxruntime to 1e-3 px. Claude-Session: https://claude.ai/code/session_013oFD6dP2EQY3R9UdXzx7dL --- client/dive-common/use/stereo/README.md | 38 ++++++++++++++++--------- 1 file changed, 24 insertions(+), 14 deletions(-) diff --git a/client/dive-common/use/stereo/README.md b/client/dive-common/use/stereo/README.md index 40c66b439..53307345e 100644 --- a/client/dive-common/use/stereo/README.md +++ b/client/dive-common/use/stereo/README.md @@ -169,28 +169,38 @@ operator, out of GPU memory) is remembered by the matcher and surfaced once, rather than retried on every frame change. Two runtime choices are deliberate, both verified probe-by-probe against CPU -onnxruntime on the fixture pair (headless Chrome, SwiftShader adapter): +onnxruntime on the fixture pair, in headless Chrome (SwiftShader) and then on +an RTX 4090 through windowed Chrome: - **The native WebGPU provider** (`onnxruntime-web/webgpu`, imported lazily), not the default bundle's JSEP kernels: JSEP returns an all-zero cost volume for this graph (a Reshape/Cast of the right-camera features reads wrong data) - on 1.27, 1.29 and 1.31-dev alike. + on 1.27, 1.29 and 1.31-dev alike, on real hardware too. - **`graphOptimizationLevel: 'basic'`**: with `'all'`, one of the provider's extended-level fusions corrupts the GRU gate convolutions and the disparity - drifts by ~2 px; at `'basic'` (and `'disabled'`) the browser output matches - CPU to 5e-5 px. + drifts by ~2 px (also on real hardware); at `'basic'` the browser output + matches CPU to 1e-3 px. The **export itself must be VIAME's web build** of the model (`plugins/onnx/export_fast_foundation_stereo_web.py`, the file the `FAST-FDN-STEREO` add-on ships): NVIDIA's stock single-file export materialises 1.5 GB correlation tensors and uses 3-D `ConvTranspose` and asymmetric 3-D -`Conv` padding that no onnxruntime-web provider runs. The web build is -numerically identical (~1e-4 px on CPU) and also cuts CPU peak memory from -16 GB to 3 GB. +`Conv` padding that no browser provider runs. The web build is numerically +identical (~1e-4 px), cuts CPU peak memory from 16 GB to ~3 GB, and expresses +every 3-D conv as 2-D convs (the providers' 3-D conv kernel is naive and was +90% of the runtime). -Real-GPU timing has not been measured (this machine's Chrome only exposes the -SwiftShader software adapter); `Desktop/Active/fast-fdn-stereo-web/webgpu_harness` -is a self-contained page that reports it. +Measured on an RTX 4090 Laptop GPU through Chrome, per stereo pair after the +first (warm-up) run: + +| Export | Browser, native WebGPU, `basic` | CPU onnxruntime | +| --- | --- | --- | +| 576×960, 8 iterations (the add-on) | 0.69 s | 5–8 s | +| 320×736, 4 iterations | 0.26 s | ~3 s | + +A hardware adapter needs Chrome in a normal window with +`--enable-features=Vulkan --enable-unsafe-webgpu` (or the matching +`chrome://flags`); headless Chrome only ever provides SwiftShader. ### How it works @@ -243,9 +253,9 @@ has already left — unless a warp is waiting on it, which upgrades it. Python with this preprocessing puts the fixture's head/tail disparities within 1 px of the NCC reference; the padding rewrite leaves the output bit-identical. -- **Verified in a browser (SwiftShader)**: the web export on the native WebGPU - provider at `basic` optimisation matches CPU onnxruntime to 5e-5 px on the - fixture pair. **Not yet exercised**: a real GPU adapter (timing), and the - end-to-end warp, settings dropdown and frame watcher in a running viewer. Set +- **Verified in a browser on an RTX 4090**: the web export on the native + WebGPU provider at `basic` optimisation matches CPU onnxruntime to 1e-3 px on + the fixture pair at 0.69 s per pair. **Not yet exercised**: the end-to-end + warp, settings dropdown and frame watcher in a running viewer. Set `DIVE_STEREO_FOUNDATION_MODEL` to an export to have the Node suite check its I/O contract. From 5f0aeadad3b73846cf11fc70d18966ff7ba2dcc0 Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Tue, 15 Sep 2026 09:20:34 -0400 Subject: [PATCH 5/9] Take the add-on's *_web.onnx export when it ships one alongside the desktop export The browser build rewrites 3-D convs, which TensorRT does not want; the add-on now carries both. Claude-Session: https://claude.ai/code/session_013oFD6dP2EQY3R9UdXzx7dL --- client/dive-common/use/stereo/README.md | 7 +++++-- server/dive_utils/stereo_models.py | 24 +++++++++++++++++++----- server/tests/test_stereo_models.py | 14 +++++++++++++- 3 files changed, 37 insertions(+), 8 deletions(-) diff --git a/client/dive-common/use/stereo/README.md b/client/dive-common/use/stereo/README.md index 53307345e..c07a421ab 100644 --- a/client/dive-common/use/stereo/README.md +++ b/client/dive-common/use/stereo/README.md @@ -149,8 +149,11 @@ in DIVE: `cmake/download_viame_addons.csv` (the list the add-on installer already uses), downloads the zip once into `DIVE_MODEL_CACHE_DIR` (default `/tmp/dive_models`, a named volume in `docker-compose.yml`) and - keeps just the model and yaml. A re-published add-on has a new md5, so it is - fetched and the old copy dropped. + keeps just the model and yaml. The add-on carries two exports of the same + weights: `fast_foundation_stereo_l.onnx` for VIAME's own CUDA/TensorRT + path and `fast_foundation_stereo_l_web.onnx` with the 3-D convolutions + rewritten for browsers; the server takes the `*_web.onnx` one. A + re-published add-on has a new md5, so it is fetched and the old copy dropped. 2. `GET dive_configuration/stereo_foundation_model/spec` reports the export's `image_size` and md5; `GET dive_configuration/stereo_foundation_model` streams the bytes. diff --git a/server/dive_utils/stereo_models.py b/server/dive_utils/stereo_models.py index d8a97c31b..05775a1a7 100644 --- a/server/dive_utils/stereo_models.py +++ b/server/dive_utils/stereo_models.py @@ -1,7 +1,8 @@ """ Serve the Fast-FoundationStereo ONNX export to the web client. -The export ships inside VIAME's ``FAST-FDN-STEREO`` add-on zip. Its download +The export ships inside VIAME's ``FAST-FDN-STEREO`` add-on zip (as the +``*_web.onnx`` build when the add-on carries one for each runtime). Its download URL and md5 are read from VIAME's ``download_viame_addons.csv`` (the same list the add-on installer uses) rather than pinned here, so a re-published model is picked up without a DIVE release. The zip is fetched once per md5 into a local @@ -99,14 +100,27 @@ def _download(url: str, dest: Path) -> str: return digest.hexdigest() +def select_web_onnx(onnx_names: List[str]) -> str: + """ + The add-on may ship two exports: the browser build (`*_web.onnx`, 3-D + convs rewritten for onnxruntime-web) next to the one VIAME's own CUDA / + TensorRT path uses. Prefer the web build; fall back to a single export. + """ + web = [n for n in onnx_names if os.path.basename(n).lower().endswith('_web.onnx')] + if len(web) == 1: + return web[0] + if len(onnx_names) == 1: + return onnx_names[0] + raise ModelUnavailable( + f'Expected one *_web.onnx or a single .onnx in the add-on, found {sorted(onnx_names)}' + ) + + def extract_model(zip_path: Path, dest_dir: Path) -> FoundationModel: """Pull the single ``.onnx`` and its sidecar ``.yaml`` out of an add-on zip.""" with zipfile.ZipFile(zip_path) as archive: names = archive.namelist() - onnx_names = [n for n in names if n.lower().endswith('.onnx')] - if len(onnx_names) != 1: - raise ModelUnavailable(f'Expected one .onnx in the add-on, found {len(onnx_names)}') - onnx_name = onnx_names[0] + onnx_name = select_web_onnx([n for n in names if n.lower().endswith('.onnx')]) yaml_name = os.path.splitext(onnx_name)[0] + '.yaml' if yaml_name not in names: raise ModelUnavailable(f'The add-on has no sidecar {os.path.basename(yaml_name)}') diff --git a/server/tests/test_stereo_models.py b/server/tests/test_stereo_models.py index 0e84f01dc..1d76ed300 100644 --- a/server/tests/test_stereo_models.py +++ b/server/tests/test_stereo_models.py @@ -17,7 +17,7 @@ YAML = 'image_size:\n- 576\n- 960\nvalid_iters: 8\n' -def make_addon_zip(yaml_text=YAML, extra_onnx=False) -> bytes: +def make_addon_zip(yaml_text=YAML, extra_onnx=False, web_onnx=False) -> bytes: buffer = io.BytesIO() with zipfile.ZipFile(buffer, 'w') as archive: archive.writestr('configs/pipelines/interactive_stereo_default.conf', 'include x.conf\n') @@ -25,6 +25,9 @@ def make_addon_zip(yaml_text=YAML, extra_onnx=False) -> bytes: archive.writestr('configs/pipelines/models/fast_foundation_stereo_l.yaml', yaml_text) if extra_onnx: archive.writestr('configs/pipelines/models/other.onnx', b'x') + if web_onnx: + archive.writestr('configs/pipelines/models/fast_foundation_stereo_l_web.onnx', b'web-bytes') + archive.writestr('configs/pipelines/models/fast_foundation_stereo_l_web.yaml', yaml_text) return buffer.getvalue() @@ -98,6 +101,15 @@ def test_ensure_model_replaces_previous_md5(tmp_path): assert not (tmp_path / old.name / old.md5).exists() +def test_extract_model_prefers_the_web_build(tmp_path): + zip_path = tmp_path / 'addon.zip' + zip_path.write_bytes(make_addon_zip(web_onnx=True)) + model = stereo_models.extract_model(zip_path, tmp_path / 'out') + assert model.onnx_path.name == 'fast_foundation_stereo_l_web.onnx' + assert model.onnx_path.read_bytes() == b'web-bytes' + assert (model.height, model.width) == (576, 960) + + def test_extract_model_requires_exactly_one_onnx(tmp_path): zip_path = tmp_path / 'addon.zip' zip_path.write_bytes(make_addon_zip(extra_onnx=True)) From 4a563f6b5e8e551645ee0c4c61eaab62186ae81c Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Tue, 15 Sep 2026 13:57:41 -0400 Subject: [PATCH 6/9] Fetch the web stereo model as a bare .onnx from the FAST-FDN-STEREO-WEB add-on row The desktop zip and the browser file are separate downloads; the input size comes from the graph, so the web file needs no sidecar. WEB-ONLY rows are hidden from the add-on installer. Claude-Session: https://claude.ai/code/session_013oFD6dP2EQY3R9UdXzx7dL --- client/dive-common/use/stereo/README.md | 43 +++++++------- .../use/stereo/StereoFoundationMatcher.ts | 21 +++++-- .../use/stereo/tests/stereoFoundation.spec.ts | 10 +++- .../web-girder/api/configuration.service.ts | 5 +- .../platform/web-girder/useStereoOnnxWeb.ts | 24 ++++---- server/dive_server/views_configuration.py | 3 + server/dive_utils/stereo_models.py | 58 +++++++++++-------- server/tests/test_stereo_models.py | 16 +++++ 8 files changed, 114 insertions(+), 66 deletions(-) diff --git a/client/dive-common/use/stereo/README.md b/client/dive-common/use/stereo/README.md index c07a421ab..0db2842db 100644 --- a/client/dive-common/use/stereo/README.md +++ b/client/dive-common/use/stereo/README.md @@ -139,29 +139,28 @@ it needs a GPU. ### Where the model comes from -The export is the one VIAME publishes in its `FAST-FDN-STEREO` add-on -(`fast_foundation_stereo_l.onnx` + sidecar `.yaml`), the same file -`plugins/onnx/fast_foundation_stereo.py` runs server-side (see *Runtime -requirements* for why it has to be VIAME's web build of it). Nothing is pinned -in DIVE: - -1. The girder server reads the add-on's URL and md5 from VIAME's - `cmake/download_viame_addons.csv` (the list the add-on installer already - uses), downloads the zip once into `DIVE_MODEL_CACHE_DIR` - (default `/tmp/dive_models`, a named volume in `docker-compose.yml`) and - keeps just the model and yaml. The add-on carries two exports of the same - weights: `fast_foundation_stereo_l.onnx` for VIAME's own CUDA/TensorRT - path and `fast_foundation_stereo_l_web.onnx` with the 3-D convolutions - rewritten for browsers; the server takes the `*_web.onnx` one. A - re-published add-on has a new md5, so it is fetched and the old copy dropped. -2. `GET dive_configuration/stereo_foundation_model/spec` reports the export's - `image_size` and md5; `GET dive_configuration/stereo_foundation_model` - streams the bytes. +The browser build is published as a **bare `.onnx`** under the +`FAST-FDN-STEREO-WEB` row (platform `WEB-ONLY`) of VIAME's +`cmake/download_viame_addons.csv`, separate from the desktop add-on zip +(`FAST-FDN-STEREO`), which carries the onnxruntime-CUDA export and a TensorRT +engine that browsers cannot use. Nothing is pinned in DIVE: + +1. The girder server reads the row's URL and md5 from the CSV, downloads the + file once into `DIVE_MODEL_CACHE_DIR` (default `/tmp/dive_models`, a named + volume in `docker-compose.yml`) and verifies the md5. A re-published model + has a new md5, so it is fetched and the old copy dropped. A zip is accepted + too (its `*_web.onnx`/single `.onnx` and yaml are extracted). +2. `GET dive_configuration/stereo_foundation_model/spec` reports the md5 and, + when a sidecar yaml exists, the input size; `GET + dive_configuration/stereo_foundation_model` streams the bytes. With a bare + `.onnx` the matcher reads the fixed input size from the graph itself + (`inputSizeOf`), so no sidecar is needed. 3. The client stores the bytes in the browser Cache API keyed by md5, so a page reload does not re-download. -For tests or a custom export, `useStereoOnnxWeb({ foundationModelUrl, -foundationModelSpec })` bypasses the server. +`WEB-ONLY` rows are skipped by VIAME's desktop add-on installer and hidden from +DIVE's add-on manager. For tests or a custom export, `useStereoOnnxWeb({ +foundationModelUrl, foundationModelSpec? })` bypasses the server. ### Runtime requirements @@ -185,8 +184,8 @@ an RTX 4090 through windowed Chrome: matches CPU to 1e-3 px. The **export itself must be VIAME's web build** of the model -(`plugins/onnx/export_fast_foundation_stereo_web.py`, the file the -`FAST-FDN-STEREO` add-on ships): NVIDIA's stock single-file export materialises +(`plugins/onnx/export_fast_foundation_stereo_web.py`, published as the +`FAST-FDN-STEREO-WEB` file): NVIDIA's stock single-file export materialises 1.5 GB correlation tensors and uses 3-D `ConvTranspose` and asymmetric 3-D `Conv` padding that no browser provider runs. The web build is numerically identical (~1e-4 px), cuts CPU peak memory from 16 GB to ~3 GB, and expresses diff --git a/client/dive-common/use/stereo/StereoFoundationMatcher.ts b/client/dive-common/use/stereo/StereoFoundationMatcher.ts index df34b232f..3493ec693 100644 --- a/client/dive-common/use/stereo/StereoFoundationMatcher.ts +++ b/client/dive-common/use/stereo/StereoFoundationMatcher.ts @@ -101,6 +101,18 @@ function defaultExecutionProviders(): string[] { return ['webgpu']; } +/** The fixed [1,3,H,W] input size of a loaded session, from its input metadata. */ +export function inputSizeOf(session: { inputMetadata?: readonly { name: string; shape?: ReadonlyArray }[] }): FoundationModelSpec { + const left = session.inputMetadata?.find((m) => m.name === 'left_image') ?? session.inputMetadata?.[0]; + const shape = left?.shape ?? []; + const height = Number(shape[2]); + const width = Number(shape[3]); + if (!(height > 0 && width > 0)) { + throw new Error('The stereo model does not declare a fixed input size; supply foundationModelSpec.'); + } + return { height, width }; +} + function rigKey(rig: StereoRig): string { return `${rig.Kl.join(',')}|${rig.Kr.join(',')}|${rig.R.join(',')}|${rig.T.join(',')}`; } @@ -231,12 +243,13 @@ export class StereoFoundationMatcher implements StereoMatcher { } /** - * Create a matcher from a model URL or model bytes. `spec` is the export's - * input resolution (its sidecar yaml `image_size`), which the graph fixes. + * Create a matcher from a model URL or model bytes. The graph fixes its + * input resolution; it is read from the session's input metadata unless + * `spec` (the sidecar yaml `image_size`) is given. */ static async create( model: string | ArrayBuffer | Uint8Array, - spec: FoundationModelSpec, + spec?: FoundationModelSpec, opts: FoundationMatcherOptions = {}, ): Promise { const executionProviders = opts.executionProviders ?? defaultExecutionProviders(); @@ -249,7 +262,7 @@ export class StereoFoundationMatcher implements StereoMatcher { // fusions the WebGPU provider applies corrupts the GRU gates (~2 px). graphOptimizationLevel: 'basic', }); - return new StereoFoundationMatcher(session, spec, opts.cacheSize, runtime); + return new StereoFoundationMatcher(session, spec ?? inputSizeOf(session), opts.cacheSize, runtime); } get inputSize(): FoundationModelSpec { diff --git a/client/dive-common/use/stereo/tests/stereoFoundation.spec.ts b/client/dive-common/use/stereo/tests/stereoFoundation.spec.ts index 1c791f5d8..fc0ef8314 100644 --- a/client/dive-common/use/stereo/tests/stereoFoundation.spec.ts +++ b/client/dive-common/use/stereo/tests/stereoFoundation.spec.ts @@ -17,7 +17,7 @@ import { } from 'vitest'; import { - StereoFoundationMatcher, remapToInputTensor, DisparitySession, WEBGPU_REQUIRED_MESSAGE, + StereoFoundationMatcher, remapToInputTensor, DisparitySession, WEBGPU_REQUIRED_MESSAGE, inputSizeOf, } from '../StereoFoundationMatcher'; import { rigFromNpz, StereoRig } from '../calibration'; import { rodrigues, computeRectification, rectifyPoint } from '../rectify'; @@ -91,6 +91,14 @@ describe('remapToInputTensor', () => { }); }); +describe('inputSizeOf', () => { + it('reads the fixed [1,3,H,W] input from the session metadata', () => { + expect(inputSizeOf({ inputMetadata: [{ name: 'right_image', shape: [1, 3, 8, 9] }, { name: 'left_image', shape: [1, 3, 576, 960] }] })) + .toEqual({ height: 576, width: 960 }); + expect(() => inputSizeOf({ inputMetadata: [{ name: 'left_image', shape: [1, 3, 'H', 'W'] }] })).toThrow('fixed input size'); + }); +}); + describe('StereoFoundationMatcher.create', () => { it('refuses to run without WebGPU rather than falling back to wasm', async () => { await expect(StereoFoundationMatcher.create(new Uint8Array(0), SPEC)).rejects.toThrow(WEBGPU_REQUIRED_MESSAGE); diff --git a/client/platform/web-girder/api/configuration.service.ts b/client/platform/web-girder/api/configuration.service.ts index 3dc0a0c65..3a589dc7c 100644 --- a/client/platform/web-girder/api/configuration.service.ts +++ b/client/platform/web-girder/api/configuration.service.ts @@ -16,8 +16,9 @@ export interface StereoFoundationModelSpec { name: string; url: string; md5: string; - height: number; - width: number; + /** Input size from a sidecar yaml; null for a bare .onnx (read from the graph instead). */ + height: number | null; + width: number | null; size: number; } diff --git a/client/platform/web-girder/useStereoOnnxWeb.ts b/client/platform/web-girder/useStereoOnnxWeb.ts index d238a3ac7..42a6c168a 100644 --- a/client/platform/web-girder/useStereoOnnxWeb.ts +++ b/client/platform/web-girder/useStereoOnnxWeb.ts @@ -10,10 +10,10 @@ * * The NCC model is a static asset (`/models/stereo_match.onnx`, produced with * `plugins/onnx/export_stereo_mapping.py --model match`). The foundation model - * is the export VIAME publishes in its FAST-FDN-STEREO add-on: the girder - * server resolves it from VIAME's add-on list and serves it, and the bytes are - * kept in the browser's Cache API keyed by the add-on's md5 so a page reload - * does not re-download ~100 MB. If no calibration or model is available the + * is the bare `.onnx` VIAME publishes under the FAST-FDN-STEREO-WEB row of its + * add-on list: the girder server resolves and serves it, and the bytes are + * kept in the browser's Cache API keyed by the row's md5 so a page reload does + * not re-download ~100 MB. If no calibration or model is available the * transfer reports the failure and no-ops. * * Because the foundation method costs one network pass per frame rather than @@ -28,6 +28,7 @@ import useStereoOnnxTransfer from 'dive-common/use/stereo/useStereoOnnxTransfer' import { StereoOnnxMatcher } from 'dive-common/use/stereo/StereoOnnxMatcher'; import { StereoFoundationMatcher } from 'dive-common/use/stereo/StereoFoundationMatcher'; import type { FoundationModelSpec } from 'dive-common/use/stereo/StereoFoundationMatcher'; +import type { StereoFoundationModelSpec } from 'platform/web-girder/api/configuration.service'; import { DEFAULT_STEREO_MATCH_METHOD } from 'dive-common/use/stereo/stereoMatcher'; import type { StereoMatcher, StereoMatchMethod } from 'dive-common/use/stereo/stereoMatcher'; import type { SearchRange } from 'dive-common/use/stereo/StereoOnnxMatcher'; @@ -58,8 +59,8 @@ export interface StereoOnnxWebOptions { modelUrl?: string; /** * Serve the foundation model from a fixed URL instead of the girder - * endpoint; `foundationModelSpec` (the export's sidecar `image_size`) is then - * required. + * endpoint. `foundationModelSpec` overrides the input size read from the + * graph. */ foundationModelUrl?: string; foundationModelSpec?: FoundationModelSpec; @@ -116,7 +117,7 @@ async function openModelCache(): Promise { * md5 comes from VIAME's add-on list, so a re-published export changes the * cache key and the stale copy is dropped. */ -async function fetchFoundationModel(): Promise<{ bytes: ArrayBuffer; spec: FoundationModelSpec }> { +async function fetchFoundationModel(): Promise<{ bytes: ArrayBuffer; spec: StereoFoundationModelSpec }> { // Imported lazily: the girder client touches `window` at load time, which // breaks node-environment unit tests that import this file. const { getStereoFoundationModelSpec, getStereoFoundationModel } = await import( @@ -159,17 +160,14 @@ export default function useStereoOnnxWeb(opts: StereoOnnxWebOptions) { async function createFoundationMatcher(): Promise { if (opts.foundationModelUrl) { - if (!opts.foundationModelSpec) { - throw new Error('foundationModelSpec is required with foundationModelUrl'); - } return StereoFoundationMatcher.create(opts.foundationModelUrl, opts.foundationModelSpec); } opts.onStatus?.('Loading the stereo model (about 100 MB on first use)...'); try { const { bytes, spec } = await fetchFoundationModel(); - return await StereoFoundationMatcher.create(new Uint8Array(bytes), { - height: spec.height, width: spec.width, - }); + // A bare .onnx has no sidecar; the matcher then reads the size from the graph. + const size = spec.height && spec.width ? { height: spec.height, width: spec.width } : undefined; + return await StereoFoundationMatcher.create(new Uint8Array(bytes), size); } finally { opts.onStatus?.(null); } diff --git a/server/dive_server/views_configuration.py b/server/dive_server/views_configuration.py index 23de9a545..b354a0bd6 100644 --- a/server/dive_server/views_configuration.py +++ b/server/dive_server/views_configuration.py @@ -231,6 +231,9 @@ def get_addons(self): decoded_content = download.content.decode('utf-8') cr = csv.reader(decoded_content.splitlines(), delimiter=',') my_list = list(cr) + # WEB-ONLY rows are bare model files the web client fetches itself + # (see dive_utils.stereo_models), not pipeline add-ons to install. + my_list = [item for item in my_list if len(item) < 5 or item[4].strip() != 'WEB-ONLY'] for item in my_list: addon = item[1] download_name = urlparse(addon).path.replace(os.path.sep, '_') diff --git a/server/dive_utils/stereo_models.py b/server/dive_utils/stereo_models.py index 05775a1a7..918d0e8fa 100644 --- a/server/dive_utils/stereo_models.py +++ b/server/dive_utils/stereo_models.py @@ -1,12 +1,14 @@ """ Serve the Fast-FoundationStereo ONNX export to the web client. -The export ships inside VIAME's ``FAST-FDN-STEREO`` add-on zip (as the -``*_web.onnx`` build when the add-on carries one for each runtime). Its download -URL and md5 are read from VIAME's ``download_viame_addons.csv`` (the same list -the add-on installer uses) rather than pinned here, so a re-published model is -picked up without a DIVE release. The zip is fetched once per md5 into a local -cache and only the model and its sidecar yaml are kept. +The browser build of the model is published as a bare ``.onnx`` file under the +``FAST-FDN-STEREO-WEB`` row (platform ``WEB-ONLY``) of VIAME's +``download_viame_addons.csv``, separate from the desktop add-on zip. Its URL +and md5 are read from that list at request time rather than pinned here, so a +re-published model is picked up without a DIVE release. The file is fetched +once per md5 into a local cache. A zip is accepted too (the model and sidecar +yaml are extracted); a bare file needs no sidecar, since onnxruntime-web reads +the input size from the graph. """ import csv @@ -24,7 +26,7 @@ from dive_utils import constants -STEREO_FOUNDATION_ADDON = 'FAST-FDN-STEREO' +STEREO_FOUNDATION_ADDON = 'FAST-FDN-STEREO-WEB' MODEL_CACHE_DIR_ENV = 'DIVE_MODEL_CACHE_DIR' DEFAULT_MODEL_CACHE_DIR = '/tmp/dive_models' DOWNLOAD_CHUNK_BYTES = 1 << 20 @@ -39,11 +41,12 @@ class AddonSource(NamedTuple): class FoundationModel(NamedTuple): onnx_path: Path - yaml_path: Path + yaml_path: Optional[Path] url: str md5: str - height: int - width: int + # From the sidecar yaml when there is one; None for a bare .onnx. + height: Optional[int] + width: Optional[int] class ModelUnavailable(Exception): @@ -116,29 +119,34 @@ def select_web_onnx(onnx_names: List[str]) -> str: ) +def is_zip(path: Path) -> bool: + with open(path, 'rb') as handle: + return handle.read(4) == b'PK\x03\x04' + + def extract_model(zip_path: Path, dest_dir: Path) -> FoundationModel: """Pull the single ``.onnx`` and its sidecar ``.yaml`` out of an add-on zip.""" with zipfile.ZipFile(zip_path) as archive: names = archive.namelist() onnx_name = select_web_onnx([n for n in names if n.lower().endswith('.onnx')]) yaml_name = os.path.splitext(onnx_name)[0] + '.yaml' - if yaml_name not in names: - raise ModelUnavailable(f'The add-on has no sidecar {os.path.basename(yaml_name)}') dest_dir.mkdir(parents=True, exist_ok=True) targets = {} for member in (onnx_name, yaml_name): + if member not in names: + continue target = dest_dir / os.path.basename(member) with archive.open(member) as src, open(target, 'wb') as out: shutil.copyfileobj(src, out) targets[member] = target - return _describe(targets[onnx_name], targets[yaml_name], url='', md5='') + return _describe(targets[onnx_name], targets.get(yaml_name), url='', md5='') -def _describe(onnx_path: Path, yaml_path: Path, url: str, md5: str) -> FoundationModel: - size = parse_image_size(yaml_path.read_text()) - if size is None: - raise ModelUnavailable(f'{yaml_path.name} does not declare image_size') - return FoundationModel(onnx_path, yaml_path, url, md5, size[0], size[1]) +def _describe(onnx_path: Path, yaml_path: Optional[Path], url: str, md5: str) -> FoundationModel: + size = parse_image_size(yaml_path.read_text()) if yaml_path else None + return FoundationModel( + onnx_path, yaml_path, url, md5, size[0] if size else None, size[1] if size else None + ) def _cached(addon: AddonSource, cache_dir: Path) -> Optional[FoundationModel]: @@ -149,9 +157,7 @@ def _cached(addon: AddonSource, cache_dir: Path) -> Optional[FoundationModel]: if len(onnx_files) != 1: return None yaml_path = onnx_files[0].with_suffix('.yaml') - if not yaml_path.is_file(): - return None - return _describe(onnx_files[0], yaml_path, addon.url, addon.md5) + return _describe(onnx_files[0], yaml_path if yaml_path.is_file() else None, addon.url, addon.md5) def ensure_model( @@ -174,9 +180,9 @@ def ensure_model( if cached is not None: return cached with tempfile.TemporaryDirectory(dir=addon_dir) as tmp: - zip_path = Path(tmp) / 'addon.zip' + download_path = Path(tmp) / 'download' try: - actual_md5 = download(addon.url, zip_path) + actual_md5 = download(addon.url, download_path) except requests.RequestException as exc: raise ModelUnavailable(f'Could not download {addon.url}: {exc}') from exc if addon.md5 and actual_md5 != addon.md5: @@ -185,7 +191,11 @@ def ensure_model( f'({actual_md5} != {addon.md5})' ) staging = Path(tmp) / 'model' - extract_model(zip_path, staging) + if is_zip(download_path): + extract_model(download_path, staging) + else: + staging.mkdir() + download_path.rename(staging / f'{addon.name.lower()}.onnx') final_dir = addon_dir / addon.md5 if final_dir.exists(): shutil.rmtree(final_dir) diff --git a/server/tests/test_stereo_models.py b/server/tests/test_stereo_models.py index 1d76ed300..f17a5e840 100644 --- a/server/tests/test_stereo_models.py +++ b/server/tests/test_stereo_models.py @@ -12,6 +12,8 @@ '1b71862b7fa39def315c0c08c44048a7, ALL-PLATFORMS, "PYTORCH, ONNX", models/fish.zip\n' 'FAST-FDN-STEREO, https://example.com/stereo/download, Fast foundation stereo, ' '29D4CDE2E33500E74844D07C5AB3DEE0, ALL-PLATFORMS, "PYTORCH", \n' + 'FAST-FDN-STEREO-WEB, https://example.com/stereo_web/download, Browser build, ' + 'ab' * 16 + ', WEB-ONLY, "", models/fast_foundation_stereo_l_web.onnx\n' ) YAML = 'image_size:\n- 576\n- 960\nvalid_iters: 8\n' @@ -41,6 +43,7 @@ def download(url: str, dest: Path) -> str: def test_parse_addon_rows_strips_whitespace_and_lowercases_md5(): rows = stereo_models.parse_addon_rows(CSV) + assert stereo_models.find_addon(rows, 'FAST-FDN-STEREO-WEB').url == 'https://example.com/stereo_web/download' stereo = stereo_models.find_addon(rows, 'FAST-FDN-STEREO') assert stereo == stereo_models.AddonSource( 'FAST-FDN-STEREO', @@ -78,6 +81,19 @@ def download(url, dest): assert calls == ['https://example.com/stereo'] +def test_ensure_model_accepts_a_bare_onnx_without_sidecar(tmp_path): + payload = b'raw-onnx-bytes' + addon = stereo_models.AddonSource( + 'FAST-FDN-STEREO-WEB', 'https://example.com/web', hashlib.md5(payload).hexdigest() + ) + model = stereo_models.ensure_model(addon, tmp_path, fake_downloader(payload)) + assert model.onnx_path.name == 'fast-fdn-stereo-web.onnx' + assert model.onnx_path.read_bytes() == payload + assert model.yaml_path is None + assert (model.height, model.width) == (None, None) + assert stereo_models.ensure_model(addon, tmp_path, fake_downloader(payload)).onnx_path == model.onnx_path + + def test_ensure_model_rejects_md5_mismatch(tmp_path): payload = make_addon_zip() addon = stereo_models.AddonSource('FAST-FDN-STEREO', 'https://example.com/stereo', 'f' * 32) From 28b31a83e126066f40c0ae5b270e21ab4ff880c8 Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Tue, 15 Sep 2026 20:55:48 -0400 Subject: [PATCH 7/9] Fall back to the published browser exports until the VIAME add-on list carries the web row The 448x768 and 576x960 files on viame.kitware.com are the defaults (DIVE_STEREO_WEB_MODEL picks one); a FAST-FDN-STEREO-WEB CSV row overrides them, and they also cover an unreachable list. Claude-Session: https://claude.ai/code/session_013oFD6dP2EQY3R9UdXzx7dL --- client/dive-common/use/stereo/README.md | 6 +++-- server/dive_utils/stereo_models.py | 33 ++++++++++++++++++++++--- server/tests/test_stereo_models.py | 28 +++++++++++++++++++++ 3 files changed, 61 insertions(+), 6 deletions(-) diff --git a/client/dive-common/use/stereo/README.md b/client/dive-common/use/stereo/README.md index 0db2842db..385182a7b 100644 --- a/client/dive-common/use/stereo/README.md +++ b/client/dive-common/use/stereo/README.md @@ -145,8 +145,10 @@ The browser build is published as a **bare `.onnx`** under the (`FAST-FDN-STEREO`), which carries the onnxruntime-CUDA export and a TensorRT engine that browsers cannot use. Nothing is pinned in DIVE: -1. The girder server reads the row's URL and md5 from the CSV, downloads the - file once into `DIVE_MODEL_CACHE_DIR` (default `/tmp/dive_models`, a named +1. The girder server reads the row's URL and md5 from the CSV (until that row + exists, or when the list cannot be fetched, it uses the builds published on + viame.kitware.com, the 448×768 one by default; `DIVE_STEREO_WEB_MODEL=576x960` + selects the full-size one), downloads the file once into `DIVE_MODEL_CACHE_DIR` (default `/tmp/dive_models`, a named volume in `docker-compose.yml`) and verifies the md5. A re-published model has a new md5, so it is fetched and the old copy dropped. A zip is accepted too (its `*_web.onnx`/single `.onnx` and yaml are extracted). diff --git a/server/dive_utils/stereo_models.py b/server/dive_utils/stereo_models.py index 918d0e8fa..a2ecf28e9 100644 --- a/server/dive_utils/stereo_models.py +++ b/server/dive_utils/stereo_models.py @@ -27,6 +27,17 @@ from dive_utils import constants STEREO_FOUNDATION_ADDON = 'FAST-FDN-STEREO-WEB' + +# Used when the VIAME add-on list has no row for the model yet (or cannot be +# fetched): the browser builds published on viame.kitware.com. The 448x768 +# export is the default for its 3.4 GB VRAM footprint; the CSV row overrides +# this, so a re-published model needs no DIVE change. +GIRDER_ITEM = 'https://viame.kitware.com/api/v1/item/{}/download' +DEFAULT_WEB_MODELS = { + '448x768': (GIRDER_ITEM.format('6aa9e846a723aa14eb79b1d4'), '0cfea82cc48435a4955a03223e1bbb02'), + '576x960': (GIRDER_ITEM.format('6aa9e850e4e84dbe3cb5b403'), 'da20c1e1837eb520d89f12bb337e38ad'), +} +DEFAULT_WEB_MODEL_ENV = 'DIVE_STEREO_WEB_MODEL' MODEL_CACHE_DIR_ENV = 'DIVE_MODEL_CACHE_DIR' DEFAULT_MODEL_CACHE_DIR = '/tmp/dive_models' DOWNLOAD_CHUNK_BYTES = 1 << 20 @@ -66,16 +77,30 @@ def find_addon(rows: Iterable[AddonSource], name: str) -> Optional[AddonSource]: return next((row for row in rows if row.name == name), None) +def default_web_model(name: str = STEREO_FOUNDATION_ADDON) -> Optional[AddonSource]: + variant = os.environ.get(DEFAULT_WEB_MODEL_ENV, '448x768') + if name != STEREO_FOUNDATION_ADDON or variant not in DEFAULT_WEB_MODELS: + return None + url, md5 = DEFAULT_WEB_MODELS[variant] + return AddonSource(name, url, md5) + + def resolve_addon(name: str = STEREO_FOUNDATION_ADDON) -> AddonSource: + """The add-on list's row for ``name``, else the built-in default for it.""" + fallback = default_web_model(name) try: response = requests.get(constants.AddonsListURL, timeout=DOWNLOAD_TIMEOUT_SECONDS) response.raise_for_status() + addon = find_addon(parse_addon_rows(response.content.decode('utf-8')), name) except requests.RequestException as exc: - raise ModelUnavailable(f'Could not read the VIAME add-on list: {exc}') from exc - addon = find_addon(parse_addon_rows(response.content.decode('utf-8')), name) - if addon is None or not addon.url: + if fallback is None: + raise ModelUnavailable(f'Could not read the VIAME add-on list: {exc}') from exc + addon = None + if addon is not None and addon.url: + return addon + if fallback is None: raise ModelUnavailable(f'The VIAME add-on list has no {name} entry') - return addon + return fallback def parse_image_size(yaml_text: str) -> Optional[tuple]: diff --git a/server/tests/test_stereo_models.py b/server/tests/test_stereo_models.py index f17a5e840..f33762e63 100644 --- a/server/tests/test_stereo_models.py +++ b/server/tests/test_stereo_models.py @@ -53,6 +53,34 @@ def test_parse_addon_rows_strips_whitespace_and_lowercases_md5(): assert stereo_models.find_addon(rows, 'MISSING') is None +def test_resolve_addon_prefers_the_csv_row_and_falls_back_to_the_published_item(monkeypatch): + class Response: + def __init__(self, text): + self.content = text.encode('utf-8') + + def raise_for_status(self): + pass + + monkeypatch.setattr(stereo_models.requests, 'get', lambda *a, **k: Response(CSV)) + assert stereo_models.resolve_addon().url == 'https://example.com/stereo_web/download' + + monkeypatch.setattr(stereo_models.requests, 'get', lambda *a, **k: Response('A, https://x, d, 1, ALL-PLATFORMS, "", \n')) + fallback = stereo_models.resolve_addon() + assert fallback.url == stereo_models.DEFAULT_WEB_MODELS['448x768'][0] + assert fallback.md5 == stereo_models.DEFAULT_WEB_MODELS['448x768'][1] + + monkeypatch.setenv(stereo_models.DEFAULT_WEB_MODEL_ENV, '576x960') + assert stereo_models.resolve_addon().url == stereo_models.DEFAULT_WEB_MODELS['576x960'][0] + + def offline(*a, **k): + raise stereo_models.requests.RequestException('offline') + + monkeypatch.setattr(stereo_models.requests, 'get', offline) + assert stereo_models.resolve_addon().url == stereo_models.DEFAULT_WEB_MODELS['576x960'][0] + with pytest.raises(stereo_models.ModelUnavailable): + stereo_models.resolve_addon('SOMETHING-ELSE') + + def test_parse_image_size_block_and_flow(): assert stereo_models.parse_image_size(YAML) == (576, 960) assert stereo_models.parse_image_size('image_size: [320, 736]\n') == (320, 736) From 15fc2852f74c962394a552ad6cb53e33362a943e Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Tue, 15 Sep 2026 20:58:28 -0400 Subject: [PATCH 8/9] Format the stereo model server code for the lint job Claude-Session: https://claude.ai/code/session_013oFD6dP2EQY3R9UdXzx7dL --- server/dive_server/views_configuration.py | 3 ++- server/dive_utils/stereo_models.py | 4 +++- server/tests/test_stereo_models.py | 24 ++++++++++++++++++----- 3 files changed, 24 insertions(+), 7 deletions(-) diff --git a/server/dive_server/views_configuration.py b/server/dive_server/views_configuration.py index 6182f9f7f..26dcbe9ea 100644 --- a/server/dive_server/views_configuration.py +++ b/server/dive_server/views_configuration.py @@ -234,7 +234,8 @@ def get_addons(self): # files the web client fetches itself (see dive_utils.stereo_models), # not pipeline add-ons to install. my_list = [ - item for item in cr + item + for item in cr if len(item) >= 5 and item[4].strip() not in ('ALL-EXCEPT-DIVE', 'WEB-ONLY') ] for item in my_list: diff --git a/server/dive_utils/stereo_models.py b/server/dive_utils/stereo_models.py index a2ecf28e9..3314c7aca 100644 --- a/server/dive_utils/stereo_models.py +++ b/server/dive_utils/stereo_models.py @@ -182,7 +182,9 @@ def _cached(addon: AddonSource, cache_dir: Path) -> Optional[FoundationModel]: if len(onnx_files) != 1: return None yaml_path = onnx_files[0].with_suffix('.yaml') - return _describe(onnx_files[0], yaml_path if yaml_path.is_file() else None, addon.url, addon.md5) + return _describe( + onnx_files[0], yaml_path if yaml_path.is_file() else None, addon.url, addon.md5 + ) def ensure_model( diff --git a/server/tests/test_stereo_models.py b/server/tests/test_stereo_models.py index f33762e63..1aacacc4f 100644 --- a/server/tests/test_stereo_models.py +++ b/server/tests/test_stereo_models.py @@ -28,8 +28,12 @@ def make_addon_zip(yaml_text=YAML, extra_onnx=False, web_onnx=False) -> bytes: if extra_onnx: archive.writestr('configs/pipelines/models/other.onnx', b'x') if web_onnx: - archive.writestr('configs/pipelines/models/fast_foundation_stereo_l_web.onnx', b'web-bytes') - archive.writestr('configs/pipelines/models/fast_foundation_stereo_l_web.yaml', yaml_text) + archive.writestr( + 'configs/pipelines/models/fast_foundation_stereo_l_web.onnx', b'web-bytes' + ) + archive.writestr( + 'configs/pipelines/models/fast_foundation_stereo_l_web.yaml', yaml_text + ) return buffer.getvalue() @@ -43,7 +47,10 @@ def download(url: str, dest: Path) -> str: def test_parse_addon_rows_strips_whitespace_and_lowercases_md5(): rows = stereo_models.parse_addon_rows(CSV) - assert stereo_models.find_addon(rows, 'FAST-FDN-STEREO-WEB').url == 'https://example.com/stereo_web/download' + assert ( + stereo_models.find_addon(rows, 'FAST-FDN-STEREO-WEB').url + == 'https://example.com/stereo_web/download' + ) stereo = stereo_models.find_addon(rows, 'FAST-FDN-STEREO') assert stereo == stereo_models.AddonSource( 'FAST-FDN-STEREO', @@ -64,7 +71,11 @@ def raise_for_status(self): monkeypatch.setattr(stereo_models.requests, 'get', lambda *a, **k: Response(CSV)) assert stereo_models.resolve_addon().url == 'https://example.com/stereo_web/download' - monkeypatch.setattr(stereo_models.requests, 'get', lambda *a, **k: Response('A, https://x, d, 1, ALL-PLATFORMS, "", \n')) + monkeypatch.setattr( + stereo_models.requests, + 'get', + lambda *a, **k: Response('A, https://x, d, 1, ALL-PLATFORMS, "", \n'), + ) fallback = stereo_models.resolve_addon() assert fallback.url == stereo_models.DEFAULT_WEB_MODELS['448x768'][0] assert fallback.md5 == stereo_models.DEFAULT_WEB_MODELS['448x768'][1] @@ -119,7 +130,10 @@ def test_ensure_model_accepts_a_bare_onnx_without_sidecar(tmp_path): assert model.onnx_path.read_bytes() == payload assert model.yaml_path is None assert (model.height, model.width) == (None, None) - assert stereo_models.ensure_model(addon, tmp_path, fake_downloader(payload)).onnx_path == model.onnx_path + assert ( + stereo_models.ensure_model(addon, tmp_path, fake_downloader(payload)).onnx_path + == model.onnx_path + ) def test_ensure_model_rejects_md5_mismatch(tmp_path): From 180b1f8cf418bbc4884f2438559451c92adb6abd Mon Sep 17 00:00:00 2001 From: Matt Dawkins Date: Tue, 15 Sep 2026 21:04:33 -0400 Subject: [PATCH 9/9] Name the stereo model exception with the Error suffix the lint job expects Claude-Session: https://claude.ai/code/session_013oFD6dP2EQY3R9UdXzx7dL --- server/dive_server/views_configuration.py | 2 +- server/dive_utils/stereo_models.py | 16 +++++++++------- server/tests/test_stereo_models.py | 6 +++--- 3 files changed, 13 insertions(+), 11 deletions(-) diff --git a/server/dive_server/views_configuration.py b/server/dive_server/views_configuration.py index 26dcbe9ea..ad43fe821 100644 --- a/server/dive_server/views_configuration.py +++ b/server/dive_server/views_configuration.py @@ -126,7 +126,7 @@ def get_training_configs(self, params): def _stereo_foundation_model() -> stereo_models.FoundationModel: try: return stereo_models.ensure_stereo_foundation_model() - except stereo_models.ModelUnavailable as exc: + except stereo_models.ModelUnavailableError as exc: raise RestException(str(exc), code=502) @access.user diff --git a/server/dive_utils/stereo_models.py b/server/dive_utils/stereo_models.py index 3314c7aca..b2e8ca313 100644 --- a/server/dive_utils/stereo_models.py +++ b/server/dive_utils/stereo_models.py @@ -60,7 +60,7 @@ class FoundationModel(NamedTuple): width: Optional[int] -class ModelUnavailable(Exception): +class ModelUnavailableError(Exception): """The model could not be resolved, downloaded or verified.""" @@ -94,12 +94,12 @@ def resolve_addon(name: str = STEREO_FOUNDATION_ADDON) -> AddonSource: addon = find_addon(parse_addon_rows(response.content.decode('utf-8')), name) except requests.RequestException as exc: if fallback is None: - raise ModelUnavailable(f'Could not read the VIAME add-on list: {exc}') from exc + raise ModelUnavailableError(f'Could not read the VIAME add-on list: {exc}') from exc addon = None if addon is not None and addon.url: return addon if fallback is None: - raise ModelUnavailable(f'The VIAME add-on list has no {name} entry') + raise ModelUnavailableError(f'The VIAME add-on list has no {name} entry') return fallback @@ -139,7 +139,7 @@ def select_web_onnx(onnx_names: List[str]) -> str: return web[0] if len(onnx_names) == 1: return onnx_names[0] - raise ModelUnavailable( + raise ModelUnavailableError( f'Expected one *_web.onnx or a single .onnx in the add-on, found {sorted(onnx_names)}' ) @@ -211,9 +211,9 @@ def ensure_model( try: actual_md5 = download(addon.url, download_path) except requests.RequestException as exc: - raise ModelUnavailable(f'Could not download {addon.url}: {exc}') from exc + raise ModelUnavailableError(f'Could not download {addon.url}: {exc}') from exc if addon.md5 and actual_md5 != addon.md5: - raise ModelUnavailable( + raise ModelUnavailableError( f'{addon.name} download did not match the add-on list md5 ' f'({actual_md5} != {addon.md5})' ) @@ -232,7 +232,9 @@ def ensure_model( shutil.rmtree(stale, ignore_errors=True) cached = _cached(addon, cache_dir) if cached is None: - raise ModelUnavailable(f'{addon.name} was downloaded but could not be read back') + raise ModelUnavailableError( + f'{addon.name} was downloaded but could not be read back' + ) return cached finally: fcntl.flock(lock, fcntl.LOCK_UN) diff --git a/server/tests/test_stereo_models.py b/server/tests/test_stereo_models.py index 1aacacc4f..e9b4b2641 100644 --- a/server/tests/test_stereo_models.py +++ b/server/tests/test_stereo_models.py @@ -88,7 +88,7 @@ def offline(*a, **k): monkeypatch.setattr(stereo_models.requests, 'get', offline) assert stereo_models.resolve_addon().url == stereo_models.DEFAULT_WEB_MODELS['576x960'][0] - with pytest.raises(stereo_models.ModelUnavailable): + with pytest.raises(stereo_models.ModelUnavailableError): stereo_models.resolve_addon('SOMETHING-ELSE') @@ -139,7 +139,7 @@ def test_ensure_model_accepts_a_bare_onnx_without_sidecar(tmp_path): def test_ensure_model_rejects_md5_mismatch(tmp_path): payload = make_addon_zip() addon = stereo_models.AddonSource('FAST-FDN-STEREO', 'https://example.com/stereo', 'f' * 32) - with pytest.raises(stereo_models.ModelUnavailable): + with pytest.raises(stereo_models.ModelUnavailableError): stereo_models.ensure_model(addon, tmp_path, fake_downloader(payload)) assert not (tmp_path / addon.name / addon.md5).exists() @@ -171,5 +171,5 @@ def test_extract_model_prefers_the_web_build(tmp_path): def test_extract_model_requires_exactly_one_onnx(tmp_path): zip_path = tmp_path / 'addon.zip' zip_path.write_bytes(make_addon_zip(extra_onnx=True)) - with pytest.raises(stereo_models.ModelUnavailable): + with pytest.raises(stereo_models.ModelUnavailableError): stereo_models.extract_model(zip_path, tmp_path / 'out')