Lazy transformations of large volumes

[csparker247]

I just came across this project and am very interested in the data augmentations and transforms. However, I tend to work with volumes that don't easily fit in memory (sometimes multiple terabytes in size), so I use h5py or zarr to have a Numpy-like array interface with lazy loading of only the volumetric regions that I need to access. Looking around at the documentation, it seems quite trivial to apply the torchio augmentations and transforms at the subregion level (just load the subregion into an ndarray and go from there), but I was wondering if you've considered support for "lazy transformations" where transforms are applied on-the-fly for only the specific pixel values requested by a slice or getitem operation.

To make this more tangible, it would be great if I could virtually apply an elastic deformation to my entire multi-terabyte volume, but the pixel values would only be interpolated within the bounds of a 3D subregion that I request:

# load my zarr
vol = zarr.open('big-vol.zarr', 'r')

# compose some virtual transforms
tfm = Compose([RandomAffineVirtual(), RandomElasticDeformationVirtual()])

# "apply" the transform to get an interface to a virtually transformed volume
vol = tfm(vol)
vol.shape  # reflects the new dimensions post-transform

# interpolate and return a subvolume
train = vol[200:400, 256:512, 256:512]

# interpolate and return the whole volume
vol = vol[:]

Or maybe this functionality already exists in this project! I would love to know about it if it does. I know that ITK was designed to do this type of thing in the C++ world, but it was not particularly fun to work with, so I'm not sure how much of that paradigm has made it into SimpleITK.

[romainVala]

Hello
That is indeed a great requierement, and there has been several demand on this point, but no implementation I am aware of. it would be a great improvement for torchio, but it is not trivial

A possible workaround would be to apply the transform directly on the patch

import torchio as tio, torch
from tqdm.auto import tqdm, trange

suj  = tio.datasets.Colin27()
sampler = tio.data.GridSampler(suj,92,92-64)
patch_loader = torch.utils.data.DataLoader(sampler)

aggregator = tio.inference.GridAggregator(sampler)
tr = tio.RandomElasticDeformation()

for patches_batch in tqdm(patch_loader, unit='batch'):
    locations = patches_batch[tio.LOCATION]
    #I do not understand why I can not apply directly the transform to patches_batch ... (so I made a new tio.Subjec)
    tmp_suj = tio.Subject(t1=tio.ScalarImage(tensor=patches_batch.t1.data,affine=patches_batch.t1.affine))
    tpm_suj_transformed = tr(tmp_suj)
    
    aggregator.add_batch(tpm_suj_transformed.t1[tio.DATA].unsqueeze(0), locations)

patchwise_output_tensor = aggregator.get_output_tensor()
io = tio.ScalarImage(tensor=patchwise_output_tensor, affine = tpm_suj_transformed.t1.affine)
io.save('/tmp/out.nii')

Looking at the result, (and comparing with the transformed applied directly on the all volume) make it clear that this approach is not equivalent .
The patches are no more coherent (to reconstruct a full volume, even with a patch overlap) but may be it does the job ... since you are only interested in the single patch (not the full transform volume)

My understanding is that it would be quite ok for RandomElastic deformation because this is a 'local" deformation but not for the random affine
imagine you do a translation of 32 of your 64 patches, then you end up with background in a region where you should not ...
(same idea with a rotation)

if you remove the patch overlap it is more clear what the problem is
sampler = tio.data.GridSampler(suj,92)
I put the image, just for the fun

Taking a big enough patch overlap, or adding a crop on the transform patch may solve this "bad" background insertion

Making a transform that would be applied patch by patch and that lead to similar result as applying the transform to the same volume, would be more complex to implement ... not sure if possible ...

[csparker247]

Yes, applying after extracting patches has its problems as the transform is applied within the coordinate frame of the patch rather than the coordinate frame of the originating volume. Affine transform is a perfect and simple example of this: the effect of scale and translation does not depend upon the location of the volume's origin but, as your example shows, the effect of rotation absolutely does.

In principle, this is not a difficult problem. Most transforms (even certain composite transforms) form a bidirectional mapping of coordinates between input and output coordinate frames. When you want to access a value in the output coordinate frame, you apply the inverse transform to the output coordinate to get a coordinate in the original frame, interpolate the value, and return. This is how ITK's ResampleImageFilter works. Implementing it within the scope of an existing project, however... that can get tricky. My guess is that it would require a separate transform and image interface than currently exists.

By the by, this sort of composable transform framework would immediately handle situations like #1052 where you want to avoid the bottleneck of multiple interpolations.

[romainVala]

By the by, this sort of composable transform framework would immediately handle situations like #1052 where you want to avoid the bottleneck of multiple interpolations.

I do see how it solve it. for me it is just a combine affine+elastic transform to have 1 instead of 2 interpolation. but it only work for a given (whole) image

[romainVala]

Yes I agree, it can be implement properly. It is easy with affine transform, and I would like also to have it for the resampling. The tricky part is how to insert it properly. My guess is that we need to add the transform to the torchio sampler because for a given output location one want a different (driven by the affine) input location ...

But the main problem is that there is no "lazy data reading in torchio : (or I miss something ?...) if you want a patch torchio will read the whole volume.

[csparker247]

But the main problem is that there is no "lazy data reading in torchio

Oh, I wasn't thinking that TorchIO should add the lazy data reading. That's already handled by h5py, zarr, and the like. You would just extend TorchIO's existing functionality for loading from an ndarray or Tensor:

# Load a volume from an HDF
hdf = h5py.File('foo.hdf', mode='r')
vol = hdf['volume']  # h5py.Dataset

# vol behaves like a numpy array
print(vol.shape, vol.dtype)

# create a scalar image
t = torchio.ScalarImage(tensor=vol)

As long as vol in this example never needs to get converted from h5py.Dataset, it will handle the lazy loading every time you do a slice operation.

I do see how it solve it. for me it is just a combine affine+elastic transform to have 1 instead of 2 interpolation. but it only work for a given (whole) image

I've looked at the library a little more, and since TorchIO is already using sitk's resampler, you could use its functionality to define a custom resampling grid [example]. TorchIO is already doing this in some sense, but it's defining a grid which covers the entire transformed image. Instead, you'd define an output grid which only covers the output area of a slice operation.

[fepegar]

@csparker247 does this help (you can use an image as reference for resampling, so maybe we can add a reference to e.g. RandomAffine)?

import torchio as tio

colin = tio.datasets.Colin27(2008)
crop = tio.CropOrPad(100)
resample = tio.Resample(3)
create_reference = tio.Compose([crop, resample])
reference = create_reference(colin.t1)
resample_roi = tio.Resample(reference)
colin_roi = resample_roi(colin)
colin_roi.plot()

[fepegar]

Ah, it seems that you're proposing something like this in #1224 (reply in thread)

[romainVala]

I like the idea, (and I also need it , I start having memory issue with high resolution above 512^3)
So I am motivate to give a try ... (let see next week)
and for sure that would be a great improvement for torchio ...

@fepegar what about adding h5py support ?
or the npz format from nnunet, (I can't find the thread where I read about the advantage of lazy reading for their internal format.

@csparker247 : does the HDF also handle the affine with the same convention as the nifti ?... (because you forget the affine in your last code line ex)

[csparker247]

Afaik, HDF doesn't support any sort of internal transformation of the coordinate system. You should think of it (and zarr) as an array storage format. Any sort of metadata about the array (transform to some world coordinate frame, pixel spacing, etc) would need to be provided manually in the same way it would if you passed a raw tensor/ndarray.

Anyway, I didn't show affine because the exact transform doesn't really matter. Affine, bspline, composite of multiple transforms - these all compose as transforms before you do anything to the data in the array.

[romainVala]

Anyway, I didn't show affine because the exact transform doesn't really matter. Affine, bspline, composite of multiple transforms - these all compose as transforms before you do anything to the data in the array.

Yes affine and elastic are relative to the input volume, so we do not care on the exat volume position (and so not need of the affine). But for other transform, like resampling, one need this information.

[csparker247]

Yes affine and elastic are relative to the input volume, so we do not care on the exat volume position (and so not need of the affine). But for other transform, like resampling, one need this information.

Both HDF5 and Zarr support metadata, but to my knowledge there is no standardized metadata for volumetric data in either of these formats as they are generic array storage formats. This is in contrast to NifTI which was explicitly designed for FMRI, and thus has appropriate spatial information.

I haven't used NifTI much, but it sounds like you're talking about a transform to map from pixel coordinates into a real world coordinate frame. If that's important and required, then users will need to provide that to TorchIO manually as I expect they would with a raw tensor. In the absence of such information, the best you can do is assume an identity transform and know that it's not in reference to any real world coordinate frame. Or maybe I'm missing something here?

[romainVala]

Sorry for the confusion affine transform and affine of a volume (or the torchio.ScalarImage) are two different thing. The later beeing indeed the affine mapping from pixel to real world.
For 3D MRI volume this information is very important for coregistration and resampling. Even though we do not really need them for some transform, it is important to keep track of it. This is why it is a mandatory input to construct a torchio image

but you are right, it will be just fine to do

# create a scalar image
t = torchio.ScalarImage(tensor=vol, affine = np.eyes(4) )

I understand this is not include in HDF5, but there is a need to keep this affine information. I'll do it separatly to start with.

[csparker247]

I should mention that HDF5 and zarr both support storing multiple arrays in a single file, so you could store that information alongside the volume quite trivially:

import h5py

# load HDF file handle
hdf = h5py.File('Patient001.hdf', mode='r')
# get vol
vol = hdf['scan01/volume/data']
# slice op [:] loads this entirely into memory because it's only 4x4
pix2world = hdf['scan01/volume/transform'][:]

t = torchio.ScalarImage(tensor=vol, affine=pix2world)

But like I say, there's no standard paths or keys or anything for this unless you impose a standard when you first create the file.

[csparker247]

One tricky implementation detail here is not so much about the resampling as it is the stochastic nature of the transforms in this project. If we compose multiple randomized transforms, when should the parameters for those transforms be updated? It seems like one might want the option of doing it at different times. As an example:

# build transforms
tfm = VirtualCompose(list_of_transforms)

# virtual dataset with virtual transform
tvol = tfm(vol)
for i in range(0, 100, 10):
    # same transform applied to each chunk
    s = tvol[i:i+10, i:i+10, i:i+10]

# option to update params on slice op
tvol = tfm(vol, update_on_slice=True)
for i in range(0, 100, 10):
    # different transform applied to each chunks
    s = tvol[i:i+10, i:i+10, i:i+10]

It seems like the former would probably be more useful in a DataLoader context, but I don't know.

[romainVala]

Yes for the random transform point of view, one do not really care to have a coherent transform among patches.

but there are so use case were we do. I would like to have a unique (random or not) transform and apply it (patch wise) to the whole volume.

Instead of implementing, as you suggest, a Virtual transform, I was thinking to implement a patch sampler, that can also apply transform to the patches.

[csparker247]

As you like, though I think the virtual transform/virtual volume class as I've described it (without the update_on_slice argument) would meet your goal using the existing patch sampler.

[romainVala]

I am open to suggestion

The difficulty I see with your Virtual Transform is how to implement the virtual loading and applying on the fly the adapted transform.
in short how do you implement the fact that
s = tvol[i:i+10, i:i+10, i:i+10]
also apply the pre-registered transfo (tvol = tfm(vol))

On the other hand, going with torchio patche-based pipelines, I can see the steps
-1 Add a transform to the patch_sampler
-2 before accessing the input box (index box of the selected patch),
2.1 apply the inverse transform the a mask representing the input box,
2.2 take the maximum box (add a few slices for interpolation) and compute the "wanted" input box in voxel
2.3 access the data and use out of the box lazy loading
2.4 apply the transform on the patch and return the result (with good crop ! to get the wanted dimension back)

of course it is easier for Affine, since the transform is the same (and only parametrised with 4*4 matrix) on all patch.

where as for Elastic you need to define it on the global volume (num_control_points ). this mean you also compute the bspline on the whole image (their may be memory issue here) and then just slide it (corresponding to the selected patch). For this transform the option update_on_slice=True will be easier to implement. ... (parameter are define at the patch level only ...)

Since the logic is different depending on the transform, so it goes in your direction, to implement a VirtualTransform ...

[csparker247]

The difficulty I see with your Virtual Transform is how to implement the virtual loading and applying on the fly the adapted transform. in short how do you implement the fact that s = tvol[i:i+10, i:i+10, i:i+10] also apply the pre-registered transfo (tvol = tfm(vol))

By using the built-in functionality of sitk.ResampleImageFilter.

I believe that every current transform class in this project uses sitk.ResampleImageFilter to perform the actual resampling. Certainly affine and elastic do. The basic idea in both is:

1. Construct an sitk.Transform of any type
2. Construct the coordinates of the output region, currently by using the same region as the input volume
3. Set up ResamplelmageFilter with the transform and the output bounds
4. Compute

When you set the output region to be the same as the input region, you get something like the following:

But if you inserted yourself into the process and changed the output region to only the region specified by a slice operation (see here), ImageResampleFilter would only resample the pixels that are contained within the output region:

In terms of how this works when you combine affine and elastic transforms, well, that's built-in to SimpleITK and you'd use sitk.CompositeTransform to combine the multiple transforms into a single transform that's passed to the image resampler (see here). Yes, all transforms would be defined on the coordinate frame of the global volume (including bspline), but the transform is only every evaluated locally.

So after all of that, tou need a new class for a VirtualVolume that's something like:

class VirtualVolume:
  _src: ArrayLike
  _tfm: sitk.Transform
  
  # returned when you do tfm(vol)
  def __init__(src, tfm):
    _src = src
    _tfm = tfm
  
  # called on slice op
  def __getitem__(*args, **kwargs):
    z_range, y_range, x_range = *args # simple assumption of 3 slices
    out_origin =[z_range.start, y_range.start, x_range.start]
    out_size = [
        z_range.end - z_range.start,
        y_range.end - y_range.start,
        x_range.end - x_range.start,
    ] 
    
   resampler = sitk.ResampleImageFilter()
   resampler.SetOutputOrigin(out_origin)
   resampler.SetOutputSize(out_size)
   resampler.SetTransform(self._tfm)
   resampled = resampler.Execute(self._src)
   
   np_array = sitk.GetArrayFromImage(resampled)
   np_array = np_array.transpose()  # ITK to NumPy
   tensor = torch.as_tensor(np_array)
   return tensor

Lots of edge cases need to be handled there, but that's more or less what I was thinking.

[csparker247]

Note: This also assumes that sitk doesn't need to convert the source array or that it too just wraps the source array and never requests the whole thing.

If it doesn't, then things get trickier, and you would have to preload some of the original volume first, figuring out the input subregion which corresponds to the output region. Sitk has methods for this too.

[romainVala]

I am not sure sitk will accept a "lazy" array
resampled = resampler.Execute(self._src)
did you test it ?

the way it is implemented in torchio is through the nib_to_sitk
and this will access the data (and properly the array ordering as expected by sitk)

So I think one need to construct the sitk volume, with only the input subregion needed. (invert transform of output region and takes the smallest box that contain it (since it can be rotated or deformed ...)

[csparker247]

sigh I just tried in both SimpleITK and ITK Python, and both attempt to load the full data into memory when you go to make an Image. SimpleITK is pretty clear about always copying the data, but I'm not sure that's the same with ITK. But I also can't find a way in Python to construct an "unmanaged" array. I know it's a wrapped C++ library, so I shouldn't expect too much, but quite annoying!

So I think one need to construct the sitk volume, with only the input subregion needed. (invert transform of output region and takes the smallest box that contain it (since it can be rotated or deformed ...)

Yes to all of this, apparently.

[romainVala]

good motivation then for torchio to offer this functionality ....

[romainVala]

The problem I can not solve yet is how to build the affine that we want to apply to the input patch, so that the output result would be the same as the affine apply to the whole volume.
rotations are closely related to the image reference frame. And the difficulty is that when considering a small patch, the reference frame is different than the one of the all image. So how to properly adjust the rotation ... ?

Let's consider to start with, a simple rotation 20 ° around z axis at (whole) image center.
how to define the rotation for a given patch

I first though that one should take the same rotation 20° around z axis and only change the rotation center so that it correspond (in the patch referential) to the rotation center define on the whole image

Since it does not work, I look at rotation applyed on the whole image, and just changing the center, with sitk SetRotation method of Euler3DTransform but it does not do what I expect ...

[csparker247]

After you create an sitk.Image from the patch, are you setting its origin to its original location in the full volume's coordinate frame?

You shouldn't need to adjust the transform at all.

[romainVala]

Yes I try to set the input (and output) patch affine correctly so that the get the same spatial position as the full volume.
and I change torchio Affine so that I can set a different output Reference for the resampler

Yes it would be much better not to change the transform (so that it can be any transform)

but when trying with rotation, I get confuse ...

I was concerned about this line

torchio/src/torchio/transforms/augmentation/spatial/random_affine.py

Line 312 in 8fc9173

transform.SetCenter(center_lps)

So I just try without the default value for center (so that sitk center is not change)
tr = tio.Affine(scales=1,degrees=(0,0,20), translation=0, center='origin')

but I still get black parts in my recomposed image, I need to triple check ...

[csparker247]

When I said "the transform shouldn't have to change", what I meant was that a sitk.Transform defined for the full volume shouldn't need to be changed when applied to a properly constructed patch. Maybe this is what you're trying to do already, but quickly glancing at the tio class, it seems like there are lots of places where the logic of the existing tio transform classss would need to be rethought.

For example, what do you pass to apply_transforn? The patch? Then it looks like all of the existing logic will use the patch's coordinate frame to create the transform, when it should be using the original volume's.

I suppose making the rotation relative to the origin should work, but is tio using the world origin, the patch's world space origin (the location of the patch's 0,0,0 pixel, but in the parent's world space coordinates) or its pixel own pixel origin (0,0,0)? Only the first is the correct origin, but I suspect it's one of the latter.

[romainVala]

Thanks for your comments @csparker247 , it helps me to focus and find out my mistakes.
After many try and test, I finally succeed to apply a given affine patch by patch and obtain the same as performing the affine once on the whole volume !
I need now more time to rearrange my code to propose a PR.

[romainVala]

Ok, now about lazy loading : It is almost ok if you use torchio patch sampler.

In order to make it work I had to made those changes
data/image.py add at the begining
from h5py import Dataset
and near line 506

        elif isinstance(tensor, Dataset):
            #no check no casting to avoid reading all the data
            return tensor

This change allow to create a torchio image, but hdf5 numpy like array

Torchio was implemented so that all data are always store as torch tensor. (which we want to avoid here in order not to reall the all file)
Patch sampler rely on crop transform (which is all you need for a lazy reading) so I had to also modify line 51 in transforms.preprocessing.spatial.crop.py
change
image.set_data(image.data[:, i0:i1, j0:j1, k0:k1].clone()) by

            if isinstance(image.data,torch.Tensor):
                image.set_data(image.data[:, i0:i1, j0:j1, k0:k1].clone())
            else:
                image.set_data(image.data[:, i0:i1, j0:j1, k0:k1])

(and add import torch at the begining of the file)

With the preceding changes in torchio you can now do lazy loading with patch sampler (@fepegar would a PR with those changes be likely accepted ?)

# reading hdf5 
f = h5py.File("Colin_0125mm.hdf5", "r")
all_data = f['data'] # but do not read all !
affine = f['affine'][:]
suj_hdf5 = tio.Subject({'t1': tio.ScalarImage(tensor = all_data, affine = affine)})
                        
sampler = tio.data.GridSampler(suj_hdf5,256,0)

for nump, one_patch in enumerate(sampler): #no DataLoader so that we keep Subject (no batch dimentsion)
    locations = one_patch[tio.LOCATION]
    loc = (locations * scale_factor).to(int)
    print(f'P {nump} loc {locations}')

the for loop run without loading the whole data from disk
and within the for loop one_patch is a torchio subject containing only a patch of the original volume.
so at this point you can apply any torchio transform to your patch. (one_patch)

if you use random transform, then you get a different random transform for each patch, but for data augmentation during training, this may be good enough

Next step is to have "coherent" spatial transform so that applying it at the patch level would lead to the same result as for the whole volume. This is what we discuss previously. I make it work for the affine, but it should work for any sitk transform too.

[romainVala]

Not sure if usefull, but I add here a more detail example with :
resampling colin subject's T1 to 0.125 mm and saving the result to an hdf5 file.
Then using this hdf5 file, and patch sampler, I perform a resampling back to 1 mm (patch by patch on the high res hdf5 files with lazzy loading

test_patch.py.txt