images source (tracel-ai#2558)

* Keep the image source as a field of ImageDatasetItem

* add image source to batch

* Update data.rs

* Fix fmt

---------

Co-authored-by: Guillaume Lagrange <[email protected]>

crates/burn-dataset/src/vision/image_folder.rs

@@ -97,6 +97,9 @@ pub struct ImageDatasetItem {
 
     /// Annotation for the image.
     pub annotation: Annotation,
+
+    /// Original image source.
+    pub image_path: String,
 }
 
 /// Raw annotation types.
@@ -250,6 +253,7 @@ impl Mapper<ImageDatasetItemRaw, ImageDatasetItem> for PathToImageDatasetItem {
         ImageDatasetItem {
             image: img_vec,
             annotation,
+            image_path: item.image_path.display().to_string(),
         }
     }
 }

examples/custom-image-dataset/src/data.rs

@@ -47,6 +47,7 @@ pub struct ClassificationBatcher<B: Backend> {
 pub struct ClassificationBatch<B: Backend> {
     pub images: Tensor<B, 4>,
     pub targets: Tensor<B, 1, Int>,
+    pub images_path: Vec<String>,
 }
 
 impl<B: Backend> ClassificationBatcher<B> {
@@ -83,6 +84,9 @@ impl<B: Backend> Batcher<ImageDatasetItem, ClassificationBatch<B>> for Classific
             })
             .collect();
 
+        // Original sample path
+        let images_path: Vec<String> = items.iter().map(|item| item.image_path.clone()).collect();
+
         let images = items
             .into_iter()
             .map(|item| TensorData::new(image_as_vec_u8(item), Shape::new([32, 32, 3])))
@@ -100,6 +104,10 @@ impl<B: Backend> Batcher<ImageDatasetItem, ClassificationBatch<B>> for Classific
 
         let images = self.normalizer.normalize(images);
 
-        ClassificationBatch { images, targets }
+        ClassificationBatch {
+            images,
+            targets,
+            images_path,
+        }
     }
 }