Correct handling of slices that extend beyond cloudvolume boundaries

- Only grabs areas that exist
- Regions beyond the bounds of the volume are set to 0.

Author: den-sq

python/ouroboros/helpers/volume_cache.py

@@ -1,8 +1,9 @@
+from dataclasses import astuple
 import os
 import sys
 import traceback
 
-from cloudvolume import CloudVolume, VolumeCutout
+from cloudvolume import CloudVolume, VolumeCutout, Bbox
 import numpy as np
 
 from .bounding_boxes import BoundingBox, boxes_dim_range
@@ -158,16 +159,26 @@ def remove_volume(self, volume_index: int, destroy_shared: bool = False):
     def download_volume(
         self, volume_index: int, bounding_box: BoundingBox, parallel=False
     ) -> VolumeCutout:
-        bbox = bounding_box.to_cloudvolume_bbox()
+        bbox = bounding_box.to_cloudvolume_bbox().astype(int)
+        vol_shape = NGOrder(*bbox.size3(), self.cv.cv.num_channels)
+
+        # Limit size of area we are grabbing, in case we go out of bounds.
+        dl_box = Bbox.intersection(self.cv.cv.bounds, bbox)
+        local_min = [int(start) for start in np.subtract(dl_box.minpt, bbox.minpt)]
+
+        local_bounds = np.s_[*[slice(start, stop) for start, stop in
+                               zip(local_min, np.sum([local_min, dl_box.size3()], axis=0))],
+                             :]
 
         # Download the bounding box volume
         if self.use_shared:
-            vol_shape = NGOrder(*bbox.astype(int).size3(), self.cv.cv.num_channels)
             volume = self.shm_host.SharedNPArray(vol_shape, np.float32)
             with volume as volume_data:
-                volume_data[:] = self.cv.cv.download(bbox, mip=self.mip, parallel=parallel)
+                volume_data[:] = 0     # Prob not most efficient but makes math much easier
+                volume_data[local_bounds] = self.cv.cv.download(dl_box, mip=self.mip, parallel=parallel)
         else:
-            volume = self.cv.cv.download(bbox, mip=self.mip, parallel=parallel)
+            volume = np.zeros(astuple(vol_shape))
+            volume[local_bounds] = self.cv.cv.download(dl_box, mip=self.mip, parallel=parallel)
 
         # Store the volume in the cache
         self.volumes[volume_index] = volume

python/test/helpers/test_volume_cache.py

@@ -1,5 +1,9 @@
 import pytest
 from unittest.mock import MagicMock, patch
+
+from cloudvolume import Bbox
+import numpy as np
+
 from ouroboros.helpers.volume_cache import (
     VolumeCache,
     CloudVolumeInterface,
@@ -19,6 +23,8 @@ def mock_cloud_volume():
         mock_cv.shape = (100, 100, 100, 3)
         mock_cv.cache.flush = MagicMock()
         mock_cv.mip_volume_size = lambda mip: (100, 100, 100)
+        mock_cv.bounds = Bbox((0, 0, 0), (100, 100, 100))
+        mock_cv.download.return_value = np.zeros((0, 10, 0, 1))
         yield mock_cv
 
 
@@ -173,8 +179,8 @@ def test_request_volume_for_slice(volume_cache):
         volume_data, bounding_box = volume_cache.request_volume_for_slice(slice_index)
 
         mock_volume_index.assert_called_once_with(slice_index)
-        assert bounding_box == volume_cache.bounding_boxes[1]
-        assert volume_data == volume_cache.volumes[1]
+        assert np.all(bounding_box == volume_cache.bounding_boxes[1])
+        assert np.all(volume_data == volume_cache.volumes[1])
 
 
 def test_create_processing_data(volume_cache):
@@ -234,7 +240,7 @@ def test_volume_cache_remove_volume(volume_cache):
         volume_data, bounding_box = volume_cache.request_volume_for_slice(slice_index)
 
         mock_volume_index.assert_called_once_with(slice_index)
-        assert volume_data == volume_cache.volumes[1]
+        assert np.all(volume_data == volume_cache.volumes[1])
         volume_cache.remove_volume(1)
         assert volume_cache.volumes[1] is None