From eac60d760bc4d5e15531ccd775d0065a24763f2a Mon Sep 17 00:00:00 2001
From: Anthony Torlucci <anthonytorlucci@gmail.com>
Date: Mon, 18 Nov 2024 07:15:15 -0600
Subject: [PATCH] Add segmentation mask to burn book (#2495)

---
 burn-book/src/building-blocks/dataset.md | 31 +++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/burn-book/src/building-blocks/dataset.md b/burn-book/src/building-blocks/dataset.md
index 8fb842e60..00a82a46f 100644
--- a/burn-book/src/building-blocks/dataset.md
+++ b/burn-book/src/building-blocks/dataset.md
@@ -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`: