[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.18779% with 2338 lines in your changes missing coverage. Please review.

Project coverage is 82.10%. Comparing base (6015823) to head (5a2931f).

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 ⚠️
crates/burn-vision/src/backends/cpu/ops.rs	0.00%	25 Missing ⚠️
... and 3 more

❌ Your patch check has failed because the patch coverage (27.18%) is below the target coverage (80.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2753      +/-   ##
==========================================
- Coverage   83.67%   82.10%   -1.58%     
==========================================
  Files         832      849      +17     
  Lines      109524   112704    +3180     
==========================================
+ Hits        91643    92532     +889     
- Misses      17881    20172    +2291     

☔ 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]

Is there no partial derivative for those ops? In this case, should we actually implement the autodiff backend?

[wingertge]

The reason I did that is because Tensor is generic over the backend, and it would be much more ergonomic to be able to use autodiff tensors for pre/post-processing. Not sure how big the actual use case is, but you might need some vision ops for at least evaluation so having it implemented makes things a lot nicer.

[nathanielsimard]

Can it be unchecked? I saw an early return previously. Same question for all launch kernels.

[wingertge]

Probably, I'm just a bit paranoid after that big OOB write issue we had a while ago. As long as we keep it in mind in case any future issues pop up, it should be fine to run unchecked.

[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).