added unet example #2786
added unet example #2786
Conversation
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #2786 +/- ##
==========================================
- Coverage 83.64% 81.69% -1.95%
==========================================
Files 826 851 +25
Lines 108746 113982 +5236
==========================================
+ Hits 90956 93122 +2166
- Misses 17790 20860 +3070 ☔ View full report in Codecov by Sentry. |
| pub struct BrainTumorDataset { | ||
| dataset: InMemDataset<BrainTumorItem>, | ||
| } |
There was a problem hiding this comment.
I was sure that you were going to use the segmentation dataset you added in #2426 😅
Any particular reason why you didn't? Looks like it should be compatible 🤔
There was a problem hiding this comment.
There are a few reasons for implementing the dataset in this example, but I plan to test and build more with the ImageFolderDataset in the future.
- It was easier to convert PyTorch code using input and output tensors.
- I'm still learning rust (15 years python; 3 years c++ in undergrad and grad school). My background is in geophysics, not computer science.
- This approach allowed me to dive a little deeper into how datasets are built and how the data loader is working. There are still some traits and advanced rust features to explore and learn here.
There was a problem hiding this comment.
All very good reasons! Was mostly curious since you were the one to add support hehe
We can definitely leave it as is and maybe update it in a follow-up PR if you're up for it.
| @@ -0,0 +1,12 @@ | |||
| [package] | |||
There was a problem hiding this comment.
If you want you can add your username to the authors field since this is basically your contribution 🙂
There was a problem hiding this comment.
Thanks, I will.
Unfortunately, I'm struggling to get some good results. The inferred segmentation image is all zeros. I believe it is due to the class imbalance since there are many more pixels that are not tumors than pixels that are. I'm re-reading the original paper and a few other references right now to refactor the code.
There was a problem hiding this comment.
I'm re-reading the original paper and a few other references right now to refactor the code.
No problem!
In this case, I'll mark the PR as draft for now. But when the example works as expected you can mark it as ready! Feel free to poke me if you have any questions.
| The training data used in the example can be downloaded from: [kaggle - nikhilroxtomar - brain-tumor-segmentation ](https://www.kaggle.com/datasets/nikhilroxtomar/brain-tumor-segmentation/data). | ||
|
|
||
| > [!NOTE] | ||
| > Not all images are the same resolution. Most are 512 x 512 pixels. |
There was a problem hiding this comment.
In the sections below you describe how to resize the images, is this a required step for the data to be used in this example?
Everything else works for 512 x 512 and I didn't download the data yet to see.
If that's the case, I think we should probably provide the script to pre-process the images (instead of having it hidden in the DATA.md file). Even better would be to automatically download the data as we do for other examples but that's not a hard requirement.
There was a problem hiding this comment.
We can remove this too. I had to downsample because my laptop (Apple M2 Pro; 32GB RAM) couldn't really handle the full size images.
Pull Request Template
Checklist
run-checks allscript has been executed.Related Issues/PRs
This is just an example for the community. No issues were created.
Changes
New example added illustrating how to create and train a UNet architecture for image segmentation.
Testing
No changes to the library code itself; no tests added.