fixing & testing rgb convertion

Author: jakeret

AUTHORS.rst

@@ -9,7 +9,7 @@ Development Lead
 
 Contributors
 ------------
-- None yet
+* `@tdrobbins <https://github.com/tdrobbins>`_
 
 Citations
 ---------

src/unet/utils.py

@@ -43,28 +43,30 @@ def crop(image, label):
 
 def to_rgb(img: np.array):
     """
-    Converts the given array into a RGB image. If the number of channels is less
-    than 3, the array is tiled such that it has 3 channels. If the number of
-    channels is greater than 3, only the first 3 channels are used. Finally, the
-    values are rescaled to [0,255)
+    Converts the given array into a RGB image and normalizes the values to [0, 1).
+    If the number of channels is less than 3, the array is tiled such that it has 3 channels.
+    If the number of channels is greater than 3, only the first 3 channels are used
 
-    :param img: the array to convert [nx, ny, channels]
+    :param img: the array to convert [bs, nx, ny, channels]
 
-    :returns img: the rgb image [nx, ny, 3]
+    :returns img: the rgb image [bs, nx, ny, 3]
     """
     img = img.astype(np.float32)
     img = np.atleast_3d(img)
 
     channels = img.shape[-1]
-    if channels < 3:
+    if channels == 1:
         img = np.tile(img, 3)
+
+    elif channels == 2:
+        img = np.concatenate((img, img[..., :1]), axis=-1)
+
     elif channels > 3:
-        img = img[:,:,:3]
+        img = img[..., :3]
 
     img[np.isnan(img)] = 0
     img -= np.amin(img)
     if np.amax(img) != 0:
         img /= np.amax(img)
 
-    # img *= 255
     return img

tests/test_utils.py

@@ -0,0 +1,21 @@
+import numpy as np
+import pytest
+
+from unet import utils
+
+
+@pytest.mark.parametrize("channels", [
+    1,2,3,4
+])
+def test_to_rgb(channels):
+    tensor = np.random.normal(size=(5, 12, 12, channels))
+
+    tensor[1, 5, 5, 0] = np.nan
+
+    rgb_img = utils.to_rgb(tensor)
+
+    assert rgb_img.shape[:2] == tensor.shape[:2]
+    assert rgb_img.shape[3] == 3
+
+    assert rgb_img.min() == 0
+    assert rgb_img.max() == 1