Merge pull request #22 from McHaillet/fix_and_test_dtypes_insertion

dtype issue fix for insertion and additional checks

Author: McHaillet

src/torch_image_interpolation/image_interpolation_2d.py

@@ -143,6 +143,9 @@ def insert_into_image_2d(
         raise ValueError('One coordinate pair is required for each value in data.')
     if coordinates_ndim != 2:
         raise ValueError('Coordinates must be 2D with shape (..., 2).')
+    if image.dtype != values.dtype:
+        raise ValueError('Image and values must have the same dtype.')
+
     if weights is None:
         weights = torch.zeros(size=(h, w), dtype=torch.float32, device=image.device)
 
@@ -238,7 +241,7 @@ def _insert_linear_2d(
     # calculate linear interpolation weights for each corner
     y, x = coordinates
     ty, tx = y - y0, x - x0  # fractional position between corners
-    w = torch.empty(size=(b, 2, 2), device=image.device)
+    w = torch.empty(size=(b, 2, 2), device=image.device, dtype=weights.dtype)
     w[:, 0, 0] = (1 - ty) * (1 - tx)   # C00
     w[:, 0, 1] = (1 - ty) * tx         # C01
     w[:, 1, 0] = ty * (1 - tx)         # C10
@@ -254,7 +257,11 @@ def _insert_linear_2d(
     # make sure to do atomic adds
     data = einops.rearrange(data, 'b c -> b c 1 1')
     w = einops.rearrange(w, 'b h w -> b 1 h w')
-    image.index_put_(indices=(idx_c, idx_h, idx_w), values=w * data, accumulate=True)
+    image.index_put_(
+        indices=(idx_c, idx_h, idx_w),
+        values=data * w.to(data.dtype),
+        accumulate=True
+    )
     weights.index_put_(indices=(idx_h, idx_w), values=w, accumulate=True)
 
     return image, weights

src/torch_image_interpolation/image_interpolation_3d.py

@@ -143,6 +143,9 @@ def insert_into_image_3d(
         raise ValueError('One coordinate triplet is required for each value in data.')
     if coordinates_ndim != 3:
         raise ValueError('Coordinates must be 3D with shape (..., 3).')
+    if image.dtype != values.dtype:
+        raise ValueError('Image and values must have the same dtype.')
+
     if weights is None:
         weights = torch.zeros(size=(d, h, w), dtype=torch.float32, device=image.device)
 
@@ -241,7 +244,7 @@ def _insert_linear_3d(
     # calculate trilinear interpolation weights for each corner
     z, y, x = coordinates
     tz, ty, tx = z - z0, y - y0, x - x0  # fractional position between voxel corners
-    w = torch.empty(size=(b, 2, 2, 2), device=image.device)
+    w = torch.empty(size=(b, 2, 2, 2), device=image.device, dtype=weights.dtype)
 
     w[:, 0, 0, 0] = (1 - tz) * (1 - ty) * (1 - tx)  # C000
     w[:, 0, 0, 1] = (1 - tz) * (1 - ty) * tx        # C001
@@ -262,7 +265,11 @@ def _insert_linear_3d(
     # insert weighted data and weight values at each corner
     data = einops.rearrange(data, 'b c -> b c 1 1 1')
     w = einops.rearrange(w, 'b z y x -> b 1 z y x')
-    image.index_put_(indices=(idx_c, idx_z, idx_y, idx_x), values=w * data, accumulate=True)
+    image.index_put_(
+        indices=(idx_c, idx_z, idx_y, idx_x),
+        values=data * w.to(data.dtype),
+        accumulate=True
+    )
     weights.index_put_(indices=(idx_z, idx_y, idx_x), values=w, accumulate=True)
 
     return image, weights

tests/test_image_interpolation_2d.py

@@ -1,6 +1,7 @@
 import einops
 import numpy as np
 import torch
+import pytest
 
 from torch_image_interpolation import sample_image_2d, insert_into_image_2d
 
@@ -153,3 +154,39 @@ def test_insert_into_image_nearest_interp_2d():
     expected = torch.zeros((28, 28)).float()
     expected[11, 14] = 5
     assert torch.allclose(image, expected)
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [torch.float32, torch.float64, torch.complex64, torch.complex128]
+)
+def test_insert_into_image_2d_type_consistency(dtype):
+    image = torch.rand((4, 4), dtype=dtype)
+    coords = torch.tensor(np.random.uniform(low=0, high=3, size=(3, 4, 2)))
+    values = torch.rand(size=(3, 4), dtype=dtype)
+    # cast the dtype to corresponding float for weights
+    weights = torch.zeros_like(image, dtype=torch.float64)
+
+    for mode in ['bilinear', 'nearest']:
+        image, weights = insert_into_image_2d(
+            values,
+            image=image,
+            weights=weights,
+            coordinates=coords,
+            interpolation=mode,
+        )
+        assert image.dtype == dtype
+        assert weights.dtype == torch.float64
+
+
+def test_insert_into_image_3d_type_error():
+    image = torch.rand((4, 4), dtype=torch.complex64)
+    coords = torch.tensor(np.random.uniform(low=0, high=3, size=(3, 4, 2)))
+    values = torch.rand(size=(3, 4), dtype=torch.complex128)
+    # cast the dtype to corresponding float for weights
+    with pytest.raises(ValueError):
+        insert_into_image_2d(
+            values,
+            image=image,
+            coordinates=coords,
+        )

tests/test_image_interpolation_3d.py

@@ -1,6 +1,7 @@
 import numpy as np
 import torch
 import einops
+import pytest
 
 from torch_image_interpolation import sample_image_3d, insert_into_image_3d
 
@@ -154,3 +155,39 @@ def test_insert_multiple_values_into_multichannel_image_2d_nearest():
     # check output shapes
     assert image.shape == (n_channels, 28, 28, 28)
     assert weights.shape == (28, 28, 28)
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [torch.float32, torch.float64, torch.complex64, torch.complex128]
+)
+def test_insert_into_image_2d_type_consistency(dtype):
+    image = torch.rand((4, 4, 4), dtype=dtype)
+    coords = torch.tensor(np.random.uniform(low=0, high=3, size=(3, 4, 5, 3)))
+    values = torch.rand(size=(3, 4, 5), dtype=dtype)
+    # cast the dtype to corresponding float for weights
+    weights = torch.zeros_like(image, dtype=torch.float64)
+
+    for mode in ['bilinear', 'nearest']:
+        image, weights = insert_into_image_3d(
+            values,
+            image=image,
+            weights=weights,
+            coordinates=coords,
+            interpolation=mode,
+        )
+        assert image.dtype == dtype
+        assert weights.dtype == torch.float64
+
+
+def test_insert_into_image_3d_type_error():
+    image = torch.rand((4, 4, 4), dtype=torch.complex64)
+    coords = torch.tensor(np.random.uniform(low=0, high=3, size=(3, 4, 5, 3)))
+    values = torch.rand(size=(3, 4, 5), dtype=torch.complex128)
+    # cast the dtype to corresponding float for weights
+    with pytest.raises(ValueError):
+        insert_into_image_3d(
+            values,
+            image=image,
+            coordinates=coords,
+        )