[Feat] Add new burn-vision crate with one initial op

[wingertge]

Pull Request Template

Checklist

• Confirmed that run-checks all script has been executed.
• Made sure the book is up to date with changes in this PR.

Changes

Adds a new burn-vision crate and an initial implementation of connected_components and connected_components_with_stats (as in opencv).

The algorithms are WIP and rely on tracel-ai/cubecl#446, but I'd like some early feedback on the structure of the new crate and the backend implementations.

Testing

Adds new connected components tests.

[laggui]

Since this is a backend extension, what do you think about keeping the implementation within the vision crate instead?

Similar to how you've added a default cpu_impl, the vision crate would look something like this:

burn-vison/src/
└── ops/
    ├── cpu_impl/
    ├── jit/
    ├── mod.rs
    └── base.rs

It would only move your current additions in burn-jit/src/kernel/vision to burn-vision/src/ops/jit (and change the imports to use burn_jit). The jit module can be under a feature flag.

That way, the backend implementations are kept to the required stuff only and the vision trait implementations are isolated within the crate.

[wingertge]

Since this is a backend extension, what do you think about keeping the implementation within the vision crate instead?

Similar to how you've added a default cpu_impl, the vision crate would look something like this:

burn-vison/src/
└── ops/
    ├── cpu_impl/
    ├── jit/
    ├── mod.rs
    └── base.rs

It would only move your current additions in burn-jit/src/kernel/vision to burn-vision/src/ops/jit (and change the imports to use burn_jit). The jit module can be under a feature flag.

That way, the backend implementations are kept to the required stuff only and the vision trait implementations are isolated within the crate.

That's how I originally had it, I wasn't sure which way was correct. I guess that's one vote for keeping it in the crate.
I think if backends had more support for these ops it would make sense to implement it in each backend, but there's really no vision op support for any of the other backends. They all just rely on opencv. So I guess as long as we only have two implementations it makes more sense to keep it in burn-vision.

Anyways I'm still working on CPU support, all the fast algorithms use decision forests and they're full of goto and #define. I'm trying to modify the GRAPHGEN library to generate the popular algorithms' decision trees in a Rust-compatible way, but it's rough 😅
Linking to opencv would be another option but has its own problems (C++ dependency, converting between burn tensors and opencv Mat, etc). Not sure what the correct solution there is.

[laggui]

That's how I originally had it, I wasn't sure which way was correct. I guess that's one vote for keeping it in the crate. I think if backends had more support for these ops it would make sense to implement it in each backend, but there's really no vision op support for any of the other backends. They all just rely on opencv. So I guess as long as we only have two implementations it makes more sense to keep it in burn-vision.

Yeah I was debating that as well when going over the code. But I think this is too specific to live in the backend implementations like you said, it will probably stay restricted to our own backends and not third-party backends. So it makes more sense in this state to have it as an extension in the vision crate.

Anyways I'm still working on CPU support, all the fast algorithms use decision forests and they're full of goto and #define. I'm trying to modify the GRAPHGEN library to generate the popular algorithms' decision trees in a Rust-compatible way, but it's rough 😅 Linking to opencv would be another option but has its own problems (C++ dependency, converting between burn tensors and opencv Mat, etc). Not sure what the correct solution there is.

Yeah I'm sure opencv would be a pain to use here 😬 but at the same time it offers so much for image processing lol. Must be a pain to link though.

A default CPU implementation is not a hard requirement either for the vision crate (that is, unless you also need it for your own application lol). This is an extension after all..

[wingertge]

A default CPU implementation is not a hard requirement either for the vision crate (that is, unless you also need it for your own application lol). This is an extension after all..

My vision (hah) for this crate is to provide easy access to image pre and postprocessing for burn, since the current solutions are very awkward (imageproc is pretty incomplete and doesn't support GPU acceleration, opencv is a pain to build and transferring data across FFI is hard). Part of that is always having a fallback implementation available, so you know it always at least works, and uses acceleration where possible.

Ideally, if someone is missing an operation, we could relatively easily port the default opencv implementation from C++ to Rust and set it as the default, then if someone wants to add GPU support it can be done later. Once I get GRAPHGEN up and running it should be fairly easy to do even with decision forest algorithms. That's why I'm attempting to modify the codegen instead of manually translating the generated code.

[laggui]

Not gonna lie that would be pretty awesome 👀

[codecov]

Codecov Report

Attention: Patch coverage is 27.40113% with 2313 lines in your changes missing coverage. Please review.

Project coverage is 81.89%. Comparing base (6755f14) to head (d4fa1c0).
Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
...nts/spaghetti/Spaghetti_center_line_forest_code.rs	12.13%	869 Missing ⚠️
...s/jit/connected_components/hardware_accelerated.rs	3.12%	527 Missing ⚠️
...nents/spaghetti/Spaghetti_last_line_forest_code.rs	6.43%	407 Missing ⚠️
...rc/backends/jit/connected_components/prefix_sum.rs	0.00%	232 Missing ⚠️
...ents/spaghetti/Spaghetti_first_line_forest_code.rs	27.41%	90 Missing ⚠️
...nts/spaghetti/Spaghetti_single_line_forest_code.rs	0.00%	54 Missing ⚠️
...backends/cpu/connected_components/spaghetti/mod.rs	73.51%	49 Missing ⚠️
...ates/burn-vision/src/tests/connected_components.rs	76.25%	38 Missing ⚠️
...ision/src/backends/jit/connected_components/mod.rs	0.00%	35 Missing ⚠️
...rn-vision/src/backends/cpu/connected_components.rs	96.81%	5 Missing ⚠️
... and 2 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2753      +/-   ##
==========================================
- Coverage   83.43%   81.89%   -1.55%     
==========================================
  Files         833      849      +16     
  Lines      110097   113252    +3155     
==========================================
+ Hits        91860    92745     +885     
- Misses      18237    20507    +2270     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[laggui]

I'm not familiar with these CV algorithms so I can't really comment on that, but otherwise the vision crate as an extension looks good!

The codegen is not pretty but glad to see that you managed to get it working 😅

Some small stuff below.

[laggui]

@nathanielsimard new burn-vision crate! Do we want to add it to the published crates already?

[nathanielsimard]

Another point to keep all the kernels in burn-vision. Compiling kernels is quite expensive, and if they are all within the same crate, they can't be parallelized by Cargo. I plan to move the fusion portion of burn-jit into its own crate for that reason as well, so less stuff in burn-jit is better.

[nathanielsimard]

I don't think we need to export the tests, since they run in this crate

[wingertge]

Integration tests count as running in a separate crate, and we need to use integration tests because of annoying rust limitations around macro_export (which is generated by burn-tensor-testgen).

[nathanielsimard]

I don't think we're using burn-tensor-testgen in burn-core, and yeah we're probably not running the tests as integration tests for that reason.

[wingertge]

We might want to parameterize some of the tests at some point, so I think using testgen right away is better. I just didn't add a parameterized test yet because connected components are almost never used with anything other than u32. I believe opencv also supports u16 but I don't think it's work adding parameterized tests until there's more ops. It's good to have the infrastructure for it regardless.