batch_size, width, height, depth = 1, 128, 128, 32
num_channels = 32
# Input data
input_data = torch.rand(batch_size, num_channels, depth, height, width).float()
sampling_grid = (torch.rand(batch_size, 256, 256, 3) - 0.5)*2.0

# create interpolation layer
trilinear_interpolation = TrilinearIntepolation()
# apply_interpolation
interpolated_data = trilinear_interpolation(input_data, sampling_grid)
print(interpolated_data.shape)

# before
sampled_in_2d = F.grid_sample(input=input_feats.view(batch_size, num_chans * num_d, height, width),
                                      grid=xy_grid, mode='nearest').view(batch_size, num_chans, num_d, grid_height,
                                                                         grid_width)

# after
sampled_in_2d = F.grid_sample(input=input_feats.view(batch_size, num_chans * num_d, height, width),
                                      grid=xy_grid, mode='nearest', align_corners=True).view(batch_size, num_chans, num_d, grid_height,
                                                                         grid_width)

from interpolation import TrilinearIntepolation
import torch

if __name__ == "__main__":
    batch = 1
    feature_channel_num = 1
    data_points = torch.tensor([1., 2., 3., 4., 5., 6., 7., 8.]).reshape(batch, feature_channel_num, 2, 2, 2)
    sample_point_pos = torch.tensor([
        [0., 0., 0.],
        [1., 0., 0.],
        [0., 1., 0.],
        [1., 1., 0.],
        [0., 0., 1.],
        [1., 0., 1.],
        [0., 1., 1.],
        [1., 1., 1.],
    ]).reshape(batch, 4, 2, 3) * 2.0 - 1.0
    trilinear_interpolation = TrilinearIntepolation()
    output = trilinear_interpolation(data_points, sample_point_pos)
    print(output)