✨ Implement len function for XbatcherSlicerIterDataPipe

docs/chipping.md

@@ -126,8 +126,7 @@ This should give us about 12 chips in total, 6 from each of the 2 Sentinel-1
 images that were passed in.
 
 ```{code-cell}
-chips = [chip for chip in dp_xbatcher]
-print(f"Number of chips: {len(chips)}")
+print(f"Number of chips: {len(dp_xbatcher)}")
 ```
 
 Now, if you want to customize the sliding window (e.g. do overlapping strides),
@@ -145,14 +144,14 @@ Great, and this overlapping stride method should give us more 512x512 chips 🧮
 than before.
 
 ```{code-cell}
-chips = [chip for chip in dp_xbatcher]
-print(f"Number of chips: {len(chips)}")
+print(f"Number of chips: {len(dp_xbatcher)}")
 ```
 
 Double-check that single chips are of the correct dimensions
 (band: 1, y: 512, x: 512).
 
 ```{code-cell}
+chips = list(dp_xbatcher)
 sample = chips[0]
 sample
 ```
@@ -232,7 +231,7 @@ Then, pass this collate function to
 
 ```{code-cell}
 dp_collate = dp_batch.collate(collate_fn=xr_collate_fn)
-print(f"Number of mini-batches: {len(list(dp_collate))}")
+print(f"Number of mini-batches: {len(dp_collate)}")
 print(f"Mini-batch tensor shape: {list(dp_collate)[0].shape}")
 ```
 

docs/object-detection-boxes.md

@@ -447,7 +447,7 @@ def boximg_collate_fn(samples) -> (list[torch.Tensor], torch.Tensor, list[dict])
 
 ```{code-cell}
 dp_collate = dp_batch.collate(collate_fn=boximg_collate_fn)
-print(f"Number of mini-batches: {len(list(dp_collate))}")
+print(f"Number of mini-batches: {len(dp_collate)}")
 mini_batch_box, mini_batch_img, mini_batch_metadata = list(dp_collate)[1]
 print(f"Mini-batch image tensor shape: {mini_batch_img.shape}")
 print(f"Mini-batch box tensors: {mini_batch_box}")

zen3geo/datapipes/xbatcher.py

@@ -33,7 +33,7 @@ class XbatcherSlicerIterDataPipe(IterDataPipe[Union[xr.DataArray, xr.Dataset]]):
         stacked into one dimension called ``batch``.
 
     kwargs : Optional
-        Extra keyword arguments to pass to :py:func:`xbatcher.BatchGenerator`.
+        Extra keyword arguments to pass to :py:class:`xbatcher.BatchGenerator`.
 
     Yields
     ------
@@ -88,7 +88,7 @@ def __init__(
         self,
         source_datapipe: IterDataPipe[Union[xr.DataArray, xr.Dataset]],
         input_dims: Dict[Hashable, int],
-        **kwargs: Optional[Dict[str, Any]]
+        **kwargs: Optional[Dict[str, Any]],
     ) -> None:
         if xbatcher is None:
             raise ModuleNotFoundError(
@@ -109,5 +109,8 @@ def __iter__(self) -> Iterator[Union[xr.DataArray, xr.Dataset]]:
             ):
                 yield chip
 
-    # def __len__(self) -> int:
-    #     return len(self.source_datapipe)
+    def __len__(self) -> int:
+        return sum(
+            len(dataarray.batch.generator(input_dims=self.input_dims, **self.kwargs))
+            for dataarray in self.source_datapipe
+        )

zen3geo/tests/test_datapipes_xbatcher.py

@@ -28,6 +28,7 @@ def test_xbatcher_slicer_dataarray():
     # Using functional form (recommended)
     dp_xbatcher = dp.slice_with_xbatcher(input_dims={"y": 64, "x": 64})
 
+    assert len(dp_xbatcher) == 4
     it = iter(dp_xbatcher)
     dataarray_chip = next(it)
 
@@ -55,6 +56,7 @@ def test_xbatcher_slicer_dataset():
     # Using functional form (recommended)
     dp_xbatcher = dp.slice_with_xbatcher(input_dims={"y": 16, "x": 16})
 
+    assert len(dp_xbatcher) == 4
     it = iter(dp_xbatcher)
     dataset_chip = next(it)