Better Integration for Neural Embedding based workflowes (eg: SAM (Segment Anything))

[kunaltyagi]

My actions before raising this issue

• Read/searched the docs
• Searched past issues

On large images, SAM latency per query adds up.

Possible Solution

Is there a way we can store the encoded image embedding in CVAT db, perhaps a separate table and then run the interactive prompt encoder + decoder with onnx.js?

The table would have a primary key on the project/task/job/frame with columns for different model embeddings. We can use a simple binary blob to store the data. This table can be populated by either by task (workflow already exists) or per frame if the db doesn't have an entry.

This would require adding support for arbitrary scripts/webasm binaries in the browser to enable use of onnx.js

Context

Scaling DB is much easier (and cheaper) than scaling GPU HW, specially since the GPU compute is idempotent

Current Implementation and Thoughts

1. Create a table manually in the database: Easy, but official support would be appreciated
2. Create a detector model which gets the request and populates the table, and doesn't detect anything ever
   • Feels like a hack to have 2 models, one for pre-process and one for use
   • Need to access poorly documented event.body payload
3. Create an interactor which runs pytorch for the prompt encoder and decoder as a serverless function
   • Sadly, making this use onnx.js would take too much effort (and know-how of cvat and js)
   • We add a checksum to the db to ensure the image is the same, else generate embeddings again

[bsekachev]

Hi @kunaltyagi

I am working on client-side optimizations for SAM in #6019

[ryanalexmartin]

Excellent idea, @kunaltyagi

@bsekachev A frontend solution is fantastic, and I'm very much looking forward to using your work once it's finished, but I would also love to see a backend implementation. as some users would prefer to use more powerful GPUs available on their cloud instances, and may be connecting from thin clients. Perhaps the Nuclio function itself could store and look-up the embeddings for images in a mounted volume.

I think this per-image workflow would feel good and responsive:

1. User selects the SAM tool, via the magic wand, and clicks "Interact", but is not yet allowed to click on the image. Red cross hair does not yet appear. A request is sent to the Nuclio serverless function. The function checks a table/db/volume/etc to see if embeddings for the image name (and matching checksum) exist.

2. The modal that usually says "Waiting for a response from Segment Anything..." instead says "generating embeddings for image..." or, if the embeddings have been previously generated, are loaded from wherever the Nuclio function stores them. The embedding of image's location is returned to the client.

3. The embeddings for the image are now generated, or loaded from the database. Red crosshair appears, and now we are able to annotate the image. The resource locator of the embeddings are included with these requests, so the model no longer has to generate them every time we click, saving the world tons of money in compute and reducing CO2 emissions.

I'll try my hand at writing this, but if anybody sees any glaring issues with what I described above, I'd appreciate a heads-up.

[kunaltyagi]

In step 2, I'd provide a button to try again (in case the image changed due to some db operation)