add make_channels_last

protonn/data/imaging/__init__.py

@@ -0,0 +1 @@
+from . import misc

protonn/data/imaging/misc.py

@@ -5,3 +5,9 @@ def make_channels_first(img):
     axes = img.shape
     img = np.rollaxis(img, axes.index(min(axes)), 0)
     return img
+
+
+def make_channels_last(img):
+    axes = img.shape
+    img = np.rollaxis(img, 0, len(axes))
+    return img

tests/test_image.py

@@ -0,0 +1,23 @@
+"""Tests for imaging"""
+import unittest
+import logging
+import protonn
+import numpy as np
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.DEBUG)
+
+
+class Tests(unittest.TestCase):
+
+    def test_channels(self):
+        logger.info("test image channels")
+        img = np.random.random((80, 100, 3)).astype(np.float32)
+        shape1 = img.shape
+        logger.info("original shape: {}".format(img.shape))
+        img = protonn.data.imaging.misc.make_channels_first(img)
+        logger.info("new shape: {}".format(img.shape))
+        img = protonn.data.imaging.misc.make_channels_last(img)
+        logger.info("resored shape: {}".format(img.shape))
+        shape2 = img.shape
+        assert shape1 == shape2