dnrops
/
burn
mirror of https://github.com/tracel-ai/burn.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add segmentation mask to burn book (#2495)

This commit is contained in:
Anthony Torlucci 2024-11-18 07:15:15 -06:00 committed by GitHub
parent 56dc4c0252
commit eac60d760b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 30 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
burn-book/src/building-blocks/dataset.md
View File

@ -137,7 +137,7 @@ those are the only requirements.
### Images
`ImageFolderDataset` is a generic vision dataset used to load images from disk. It is currently
available for multi-class and multi-label classification tasks.
available for multi-class and multi-label classification tasks as well as semantic segmentation tasks.
```rust, ignore
// Create an image classification dataset from the root folder,
@ -168,6 +168,35 @@ let dataset = ImageFolderDataset::new_multilabel_classification_with_items(
.unwrap();
```
```rust, ignore
// Create a segmentation mask dataset from a list of items, where each
// item is a tuple `(image path, mask path)` and a list of classes
// corresponding to the integer values in the mask.
let items = vec![
(
"path/to/images/image0.png",
"path/to/annotations/mask0.png",
),
(
"path/to/images/image1.png",
"path/to/annotations/mask1.png",
),
(
"path/to/images/image2.png",
"path/to/annotations/mask2.png",
),
];
let dataset = ImageFolderDataset::new_segmentation_with_items(
items,
&[
"cat", // 0
"dog", // 1
"background", // 2
],
)
.unwrap();
```
### Comma-Separated Values (CSV)
Loading records from a simple CSV file in-memory is simple with the `InMemDataset`: