Mark/optimization

[markcamera]

Optimize Roma2 Inference
Gain roughly an 80x speedup in roma2 processing (assuming only the
coarse non-bidirectional match is needed)

• Introduces Docker container
• Add feature caching to romav2.py
• Add batch coarse matching inference (using cached features) to romav2.py
• Introduce batch_match_min script calculating the minimal coarse
  matching using the new romav2.py API
• Add profiling statements to romav2.py
• Add profiling statements to batch_match.py

[sentry]

Bug: The code attempts to index a Python list using a PyTorch tensor for the index idx, which will raise a TypeError and crash the script.
_{Severity: CRITICAL}

Suggested Fix

Convert the tensor index to an integer before using it to access the list element. Change icache[iset[idx]] = im to icache[iset[idx.item()]] = im.

Prompt for AI Agent

Review the code at the location below. A potential bug has been identified by an AI
agent.
Verify if this is a real issue. If it is, propose a fix; if not, explain why it's not
valid.

Location: scripts/batch_match_min.py#L152

Potential issue: In the `load_images` and `streaming_cache` functions, the code iterates
over a PyTorch `DataLoader`. The `ImageDataset` returns an integer index, but the
`DataLoader`'s default collate function converts this integer into a 1D tensor. When the
code attempts to use this tensor `idx` to index the Python list `iset` (i.e.,
`iset[idx]`), it will raise a `TypeError` because lists can only be indexed by integers
or slices. Since the `streaming_cache` function is called from `main`, this will cause
the script to crash during execution.

_{Did we get this right? 👍 / 👎 to inform future reviews.}

[sentry]

Bug: A cached image tensor has its batch dimension dropped, causing a shape mismatch error when re-batched later with torch.cat() instead of torch.stack().
_{Severity: CRITICAL}

Suggested Fix

To preserve the batch dimension when caching the image, change the indexing in romav2.py from ibatch[idx] to ibatch[idx:idx+1]. Alternatively, in batch_match_min.py, use torch.stack(img_A) instead of torch.cat(img_A) to create a correct batch dimension.

Prompt for AI Agent

Review the code at the location below. A potential bug has been identified by an AI
agent.
Verify if this is a real issue. If it is, propose a fix; if not, explain why it's not
valid.

Location: src/romav2/romav2.py#L353

Potential issue: The `cache_batch_features` method in `romav2.py` caches rescaled images
by extracting them from a batch tensor using single-item indexing (`ibatch[idx]`). This
removes the batch dimension, resulting in a 3D tensor of shape `(C, H, W)`. Later, in
`batch_match_min.py`, the `process_real_batch` function uses `torch.cat()` to combine
these 3D tensors, incorrectly creating a tensor of shape `(B*C, H, W)`. The downstream
matcher model expects a 4D tensor of shape `(B, C, H, W)`, causing a shape mismatch that
will lead to a runtime error or incorrect results.

_{Did we get this right? 👍 / 👎 to inform future reviews.}

[cdcseacave]

it seems to be working, however there is still a problem with the resolution, no matter what setting is used, the output resolutiion is always 160x160

[cdcseacave]

can we remove this or explain how to use it?

[cdcseacave]

I fixed the script compilation, I merged the latest main, I added support to export as NPZ and a way to disable profiling, and now it is close to be production ready.

However, there is still a problem with the resolution, no matter what setting is used, the output resolution is always 160x160

I did not find a way to fix this. @markcamera can u pls help?

[markcamera]

Thanks for the review. I'll take a look into the resolution problem. To confirm though you're talking about resolution variations with respect to this setting choices=['turbo', 'fast', 'base', 'precise'], , right?

I would recommend switching over to the torch save format though if it doesn't impair other workflows as the numpy conversions were a large part of the execution time. The 'best' version involved saving batches of results (or memory permitting) all results into one .pt file, though as this is a development utility, it makes sense if that's more of an inconvenience for the time being. (reading through the modifications it looks like the scripts now support both, so that's an all around improvement :) )

[markcamera]

At least locally before the new changes are added I don't experience an issue with the sizing.

Adding a print(preds["warp_AB"].shape) into the save_batch block I see by default (16,200,200,2), precise (16,200,200,2), turbo (16,80,80,2) where 16 is the batch I'm running the script at. So is this by chance a regression with the added commits?

[cdcseacave]

yes, for ex at the default setting choice best, the resolution of the output is 160x160

That would be nice, but I use the output in C++ and I did not find a library that can load PyTorch files. Plus I tested the time with PT and NPZ and is identical.

[sentry]

Bug: The coarse_match method calls self.coarse_forward, which is not defined, leading to an AttributeError at runtime.
_{Severity: CRITICAL}

Suggested Fix

Replace the call to the non-existent self.coarse_forward with a call to an existing method that performs the coarse matching logic, such as self.forward.

Prompt for AI Agent

Review the code at the location below. A potential bug has been identified by an AI
agent.
Verify if this is a real issue. If it is, propose a fix; if not, explain why it's not
valid.

Location: src/romav2/romav2.py#L409

Potential issue: The `coarse_match` method on line 409 attempts to call
`self.coarse_forward(img_A_lr, img_B_lr)`. However, the `coarse_forward` method is not
defined in the `RoMaV2` class or its parent `nn.Module`. Any call to `coarse_match` will
immediately raise an `AttributeError`, causing the program to crash. The intended
behavior was likely to call the existing `self.forward` method.

[sentry]

Bug: The script crashes with an AssertionError if run with --batch-size > 1 due to hardcoded assertions, despite documentation encouraging larger batch sizes.
_{Severity: CRITICAL}

Suggested Fix

Either remove or update the assertions on lines 402-405 to correctly handle batch sizes greater than 1, or prevent users from setting --batch-size to a value other than 1 and update the documentation accordingly.

Prompt for AI Agent

Review the code at the location below. A potential bug has been identified by an AI
agent.
Verify if this is a real issue. If it is, propose a fix; if not, explain why it's not
valid.

Location: scripts/batch_match_min.py#L402-L405

Potential issue: The `batch_match_min.py` script allows setting a `--batch-size` greater
than 1 via command-line arguments, and its documentation provides an example with
`--batch-size 4`. However, the script contains assertions that hardcode a requirement
for a batch size of 1. When run with a batch size greater than 1, the script will
concatenate features, creating a batch dimension larger than 1. This leads to a crash
with an `AssertionError` when `preds["warp_AB"].shape[0]` is checked.

[cdcseacave]

@markcamera I've found the reason why we get very small resolution in the new ROMA2, since the refinement step at the end is not run anymore, the resolution returned for the disparity map is 4x downsampled resolution, as that is the resolution of the coarse step, which is to be expected, makes sense now. So I added 2 new options to the script, to either add a fast upscaling and refining step, but it is 4 times slower, or a way to overwrite the initial resolution, so the 4x smaller coarse resolution is still large enough. As to be expected, the best results I got for the refined version, but also the coarse version delivers good enough matches. Pls have a look over the changes, I am sure I must have made some errors in the logic as I am not very familiar with the code.

[markcamera]

It looks like a bit of logic is missing here if you want the bidirectional mode. No problem if you don't need those values (see the if.self.bidirectional ... warp_BA,confidence_BA=_interpolat... block in romav2.py if needed

[cdcseacave]

I do not want the bidirectional mode

[markcamera]

IMO misleading as the override only results in square images, but otherwise makes sense? (there may be additional restrictions with the ViT backbone block size, but that's more of a user value problem than an issue with the code).

[cdcseacave]

not sure what you mean, ROMA2 always estimates square warp images regardless of resolution, so this is not a limitation I am adding, but a ROMA2 limitation

[markcamera]

I pointed out the only two classes of error that I saw. As long as you don't mind square images (with only some resolutions being valid) and you don't actually need the bidirectional mode, this should be all good to go.