Vectorising a cost volume in NumPy (and why it still wasn't enough)

The first post in this chapter walked through SSD2.py’s per-column loop as the implementation, the one that actually produced every number in the report. It was written that way for a reason, and the reason is a memory figure I never actually worked out at the time. This post works it out. There is an earlier, more direct version of the same idea sitting right next to it in the repo (SSD.py, not SSD2.py) and it computes the entire cost volume in one line. It is also the reason the loop exists: the direct version doesn’t run.

Everything below is about four files. Two of them, SSD.py and ZNCC.py, are the “just do it all at once” scratch versions. The other two, SSD2.py and ZNCC2.py, are what survived contact with a real image. The grading post covers what came out of the surviving versions; this one is about why they had to change shape at all.

The four-deep loop nobody wrote

Block matching, stated plainly, is four nested loops: for every pixel, for every candidate disparity, for every offset inside the window, for every colour channel, accumulate a difference. Report page 12 calls direct matching “computationally expensive” for exactly this reason: in the worst case every pixel is compared against every pixel, and each comparison is itself a small loop.

Nobody in this repo actually wrote that four-deep loop. NumPy gives you a shortcut for the inner two levels (the window and the channel) before you’ve written a single for, and SSD.py uses it immediately.

The gather: three loop levels become one index expression

SSD.py:9–21:

W = 5
actualW = 2 * W + 1
V = W + np.arange(img1.shape[0] - 2 * W)
U = W + np.arange(img1.shape[1] - 2 * W)

WindowVectorV = np.arange(-W, W + 1).repeat(2 * W + 1)
WindowVectorU = np.tile(np.arange(-W, W + 1), 2 * W + 1)

Wv = V.reshape(-1, 1) + WindowVectorV
Wu = U.reshape(-1, 1) + WindowVectorU

ordered_img1 = img1[Wv[:, None], Wu[None, :]]
ordered_img2 = img2[Wv[:, None], Wu[None, :]]

Read it from the inside out. arange(-W, W+1) is one window’s worth of offsets, -5..5 for the hard-coded W = 5. .repeat(actualW) turns it into [-5,-5,…,-5, -4,-4,…,-4, …, 5,…,5] , each value repeated actualW times in a row. .tile(actualW) turns the same range into [-5,-4,…,5, -5,-4,…,5, …], the whole sequence repeated actualW times. Zipped together, element i of WindowVectorV and element i of WindowVectorU is one (dv, du) pair, and running i from 0 to actualW² − 1 visits every offset in the window exactly once. repeat walks the rows; tile walks the columns inside each row. That is the whole trick: two one-dimensional array operations produce the flattened index set of a two-dimensional window.

Add the pixel’s own coordinates: V.reshape(-1, 1) + WindowVectorV broadcasts a column of pixel rows against a row of window offsets, and Wv, Wu become, for every valid pixel in the image at once, the row and column indices of every pixel in its window. Index img1 with both and NumPy’s fancy indexing does the rest: ordered_img1 comes out with shape (rows, cols, window, channels), one window, fully materialised, for every pixel in the image, without a single Python-level loop.

Nothing about this is expensive. For the whole image at the repo’s default W = 5 (actualW = 11), ordered_img1 and ordered_img2 together are 72.0 MB, one array, computed once, held for the rest of the function. The gather is the good part of SSD.py, and SSD2.py keeps every line of it unchanged.

The line that doesn’t run

SSD.py:23:

SSD = np.sum(np.square(ordered_img1[:, :, None] - ordered_img2[:, None]), axis=(3,4))

This is the same move one level up. ordered_img1[:, :, None] and ordered_img2[:, None] broadcast against each other over a new axis (not the window axis, the column axis), so every window in the left image is compared against every window in the right image in one expression. SSD comes out with shape (rows, cols, cols): for every pixel, a score against every column in the row. That is not “bounded disparity”, it is not disparity at all yet: it is the full two-dimensional correspondence problem that the epipolar constraint was supposed to collapse, minus only the row search. SSD.py never calls max_disp; there is no parameter for it.

Before that sum runs, NumPy has to build the array inside it. ordered_img1[:, :, None] has shape (rows', cols', 1, window, channels); ordered_img2[:, None] has shape (rows', 1, cols', window, channels); broadcasting them against each other materialises (rows', cols', cols', window, channels), every window, against every other window, in the same row, all at once. At the repo’s Tsukuba size, 384 × 288, and the hard-coded W = 5:

$$\text{rows}' \times \text{cols}' \times \text{cols}' \times (2W{+}1)^2 \times 3 = 278 \times 374 \times 374 \times 121 \times 3 \approx 1.41 \times 10^{10}$$

uint8 elements: 13.1 GB, for the difference array alone. np.square(...) then allocates a second array the same size to hold the result, because NumPy doesn’t square in place unless told to. For the moment both exist (and np.sum needs the squared one to still be there to read from). The two temporaries plus the 72 MB gather come to 26.4 GB. That’s report page 12’s “computationally expensive”, worked out in actual bytes rather than repeated as a phrase: not “slow”, but a single Python statement asking for more RAM than the laptop that wrote it has ever had installed.

A bug the memory story doesn’t explain

While checking these numbers I found something the report doesn’t mention and the plan for this post didn’t flag: ordered_img1 and ordered_img2 are uint8 (that’s what cv2.imread returns) and ordered_img1[:, :, None] - ordered_img2[:, None] is a subtraction of two uint8 arrays, which NumPy keeps as uint8. Any difference below zero wraps around modulo 256, and np.square of that wrapped value also happens in uint8, wrapping again if the true square exceeds 255. I checked this rather than asserted it:

>>> a = np.uint8(20) - np.uint8(0)   # true difference: 20
>>> np.square(a)
np.uint8(144)                        # true square is 400; 400 % 256 = 144

The retreat

SSD2.py:8–31 is the function that ran. The gather is untouched: same Wv, same Wu, same ordered_img1, same ordered_img2. What changes is everything after it:

disparityMap = np.zeros(ordered_img1.shape[:2], dtype=np.uint8)

for row in range(len(ordered_img1[0])):
    SSD = np.sum(np.square(ordered_img1[:, (row,), None] - ordered_img2[:, None, max(row-max_disp, 0):max(row-min_disp+1, 1)]), axis=(3,4))
    SSD.shape = SSD.shape[0], -1
    best_i = np.argmin(SSD, axis=1)
    disparityArray = row-np.arange(max(row-max_disp, 0), max(row-min_disp+1, 1))
    best_i_scaled = (disparityArray[best_i]*disp_scale).astype(np.uint8)
    disparityMap[:, row] = np.where(SSD[np.arange(SSD.shape[0]), best_i] < thresh, best_i_scaled.flat, 0)

Two changes, and it’s worth pulling them apart, because only one of them is the whole story.

Bound the disparity. The slice max(row-max_disp, 0):max(row-min_disp+1, 1) (the same bound the first post covers as physically motivated, not just convenient) replaces the full cols' axis with a search of max_disp - min_disp + 1 candidates. At the report’s own numbers for Tsukuba, W = 4 and max_disp = 31 (TestDisp.py:44), that alone would take the broadcast from 18.0 GB down to 1.57 GB, a real win, about 30×, and still nowhere close to fitting comfortably.

Loop over columns. The for row in range(...) (misleadingly named; it walks the columns of img2, one at a time) is what actually makes the function run. Inside the loop, only ordered_img1[:, (row,), None], one column’s worth of windows, is compared against the bounded slice of ordered_img2. The huge middle axis (every column against every column) is never materialised at all; it’s visited one column at a time and thrown away. At W = 4, max_disp = 31, that per-column array plus the persistent gather comes to 52.9 MB. Bounding bought a 30× reduction; the loop buys another ~11× on top of that, and the two together are the ~347× the difference between 18.0 GB and 52.9 MB actually is.

	Shape (Tsukuba, W = 4, max_disp = 31)	Peak memory
Explicit loops (not in the repo)	one running sum	no array
Gather + per-column loop, SSD2.py:22-28	gather + one column’s window stack	52.9 MB
Gather + full broadcast, bounded (hypothetical)	every column, bounded disparity	1.57 GB
Gather + full broadcast, unbounded, SSD.py:23	every column, every column	18.0 GB

Neither change alone gets you to 52.9 MB. Bounding the disparity is the optimisation report page 13 argues for on accuracy grounds and gets a computational discount for free; the loop is the one that makes the discount enough to actually run. That’s the retreat: not “give up on vectorising”, just don’t vectorise the one axis (every column against every other column) that was never bounded by anything physical to begin with.

Why ZNCC doesn’t need the loop’s help as much

ZNCC2.py:22–33 runs the identical gather, then does something the SSD versions can’t:

mean_img_1 = (ordered_img1 - np.mean(ordered_img1, axis=2)[:, :, None])/(np.std(ordered_img1, axis=2)[:, :, None]+1e-10)
mean_img_2 = (ordered_img2 - np.mean(ordered_img2, axis=2)[:, :, None])/(np.std(ordered_img2, axis=2)[:, :, None]+1e-10)

scales = 3*actualW * actualW

for row in range(len(mean_img_1[0])):
    NCC = np.einsum("ilm,iklm->ik", mean_img_1[:, row], mean_img_2[:, max(row-max_disp, 0):max(row-min_disp+1, 1)]) / scales

Subtracting the mean and dividing by the standard deviation happens once per window, outside the disparity loop entirely: mean_img_1 and mean_img_2 are computed a single time and reused for every candidate column. SSD can’t do this: the cost of an SSD window depends on which other window it’s compared against, so the difference has to be recomputed from raw pixels at every disparity. ZNCC’s cost, once normalised, is just einsum("ilm,iklm->ik", …), a dot product between two already-finished vectors. That’s the actual reason report page 18 can say ZNCC survives a lighting change and SSD can’t: not that correlation is inherently more robust arithmetic, but that normalising once and comparing many times means the per-comparison cost is nearly free, while SSD pays for the comparison in full every single time.

Where the loop still loses

The per-column loop solves the memory problem. It doesn’t solve the other one: for every column, NumPy re-reads the whole window stack and recomputes the difference from scratch, even though 8 of a 9-pixel window’s contribution didn’t change between one disparity and the next. Post 29 covers the fix: box-aggregation in the Rust crate turns that into a running sum, so a 19×19 window costs the same as a 3×3 one, and a full pass over Tsukuba drops from something NumPy takes whole seconds over to tens of milliseconds. The memory story explains why SSD2.py loops; it doesn’t explain why the loop is still slow, and that’s the gap Rust closes.

Try it: the gather, at any window size, on the real numbers

With JavaScript on, this becomes two tools. The first lets you pick a window half-width and a pixel on a small synthetic grid, and prints the actual WindowVectorV, WindowVectorU, Wv and Wu arrays it builds (the same repeat/tile arithmetic as above), alongside a live table of what four different ways of writing the same cost volume would cost in memory, computed for the real 384×288 Tsukuba size this whole chapter uses. The second runs the same 3×3 SSD pass as a naive JS loop and as post 29’s WASM kernel, timed in your browser, with a button to run it again.

The window-half-width slider caps at 6 (past that the highlighted window runs off the edge of a grid small enough to stay legible), but the byte counter has no such limit: at W = 6 the unbounded broadcast is already at 36.2 GB, and even the smallest possible window, W = 1, a 3×3, still needs 2.10 GB unbounded on this image size. There is no small-window escape from a broadcast that scales with the square of the image width; the fourth table row does not move when you drag the disparity slider, on purpose, because SSD.py never bounded that axis either.

The benchmark

The benchmark section of the widget above runs the identical 3×3-window SSD pass (same synthetic pair, same disparity range, min_disp = 0, max_disp = 31, Tsukuba’s own bound) two ways: a naive four-deep JavaScript loop, and stereo-wasm’s stereo_match (post 29’s crate, reused here without changes). On my machine, under Node.js, three timed runs each after a warm-up, median taken:

	median
naive JS loop	274.6 ms
stereo-wasm (stereo_match)	39.73 ms
ratio	6.9×

Notice that the WASM figure, 39.73 ms, is close to post 29’s own number for a full pass at its own 9×9 window and the same 32 disparities (about 40 ms) even though that window has nine times the area of this one. That’s box-aggregation doing exactly what it says: the kernel doesn’t pay more for a bigger window. A naive loop has no such property: every one of those nine times more window pixels is nine times more work, so the gap between the two columns above would only widen at the report’s own window size, not narrow. I haven’t run that larger naive-JS pass myself to put a number on it, but the reasoning follows directly from the loop shapes in both files, which is exactly the axis a NumPy broadcast can’t touch either. The widget times your own browser and reports the ratio it measures; it will not match mine, and it shouldn’t: different engines, different cores, different everything.

What this chapter adds up to

Bound the search with geometry, gather the window with two array operations instead of a loop, and you still can’t broadcast the disparity axis, because nothing bounded it the way the epipolar line bounded the row search. SSD2.py’s loop is not a failure to vectorise; it’s the one axis where vectorising and running out of memory turned out to be the same decision, made honestly, in a script sitting right next to the version that tried the other way first. And when even the loop’s per-column recomputation turns out to be the next bottleneck, the fix isn’t a cleverer NumPy expression: it’s a language where a running sum can replace a window sum without becoming an unreadable one-liner. That’s the whole argument of this chapter, and it took two failed scripts and one working one to make it honestly.