wip

src/pancreas_ai/dataset/ds_augmentation/rotate.py

@@ -1,30 +1,37 @@
 import numpy as np
 import cv2
+import scipy.ndimage as ndimage
 
-def __rotate_CV(image, angel , interpolation):
+def __rotate_CV(image, angle: float):
     #in OpenCV we need to form the tranformation matrix and apply affine calculations
     #interpolation cv2.INTER_CUBIC (slow) & cv2.INTER_LINEAR
+    image = image[..., np.newaxis].astype(np.uint8)
     h,w = image.shape[:2]
     cX,cY = (w//2,h//2)
-    M = cv2.getRotationMatrix2D((cX,cY),angel,1)
-    rotated = cv2.warpAffine(image,M , (w,h),flags=interpolation)
+    M = cv2.getRotationMatrix2D((cX,cY),angle,1)
+
+    rotated = cv2.warpAffine(image, M, (w,h), flags=cv2.INTER_LINEAR)
     return rotated
 
-def __rotate_data(angle: float, axis: int, data: np.ndarray) -> np.ndarray:
+def __rotate_scipy(data: np.ndarray, angle: float) -> np.ndarray:
+    return ndimage.rotate(data, angle, reshape=False, order=1)
+
+def rotate_data(angle: float, axis: int, data: np.ndarray) -> np.ndarray:
+    assert data.ndim == 3
+
     unstacked = np.split(data, data.shape[axis], axis)
     unstacked = [np.squeeze(x) for x in unstacked]
-    rotated = [__rotate_CV(x, angle, cv2.INTER_LINEAR) for x in unstacked]
+    rotated = [__rotate_scipy(x, angle) for x in unstacked]
     return np.stack(rotated, axis)
 
-def random_rotate_data_and_label(data: np.ndarray, label: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
-    for i in range(len(data.shape)):
-        if np.random.rand() > 0.5:
-            data = np.flip(data, i)
-            label = np.flip(label, i)
-    return data, label
+def random_rotate_data_and_label(angle1: float, angle2: float, data: np.ndarray, label: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    angle = np.random.uniform(angle1, angle2)
+    axis = np.random.randint(0, data.ndim)
+
+    return rotate_data(angle, axis, data), rotate_data(angle, axis, label)
 
 def random_rotate_data(angle1: float, angle2: float, data: np.ndarray) -> np.ndarray:
     angle = np.random.uniform(angle1, angle2)
     axis = np.random.randint(0, data.ndim)
 
-    return __rotate_data(angle, axis, data)
+    return rotate_data(angle, axis, data)

tests/test_3d_rotate.py

@@ -1,55 +1,110 @@
-import tensorflow as tf
+import numpy as np
 import unittest
+from PIL import Image
 
 import context
 context.setup()
-from pancreas_ai.tools import resize_3d
+from pancreas_ai.dataset.ds_augmentation import rotate
 
-class Resize_3d(unittest.TestCase):
+class Rotate_3d(unittest.TestCase):
     '''
     Unit test to resize_3d function
     '''
-    def test_resize_to_same_size(self):
-        arr1 = tf.range(0, 27, dtype = tf.int32)
-        arr2 = tf.reshape(arr1, [3, 3, 3])
+    def test_rotate_90_deg(self):
+        arr1 = np.arange(0, 3 * 3 * 2, dtype = np.int32)
+        arr2 = np.reshape(arr1, [3, 3, -1])
 
-        arr3 = resize_3d.resize_3d_image(arr2, tf.constant([3, 3, 3]))
-        result = tf.logical_not(tf.cast(tf.reshape(arr2-arr3, [-1]), dtype=tf.bool)).numpy().all()
+        min = np.min(arr2)
+        max = np.max(arr2)
+        arr2 = (arr2 - min) / (max - min) * 2 - 1
 
-        self.assertEqual(result, True, msg = "resize to the same size failed")
+        arr3 = rotate.rotate_data(90, 2, arr2)
+        arr4 = rotate.rotate_data(-270, 2, arr2)
 
-    def test_resize_twice_and_back(self):
-        arr1 = tf.range(0, 27, dtype = tf.int32)
-        arr2 = tf.reshape(arr1, [3, 3, 3])
+        self.assertEqual(np.all(arr3 == arr4), True, msg = "rotation 90 and -270 failed")
 
-        arr3 = resize_3d.resize_3d_image(arr2, tf.constant([6, 6, 6]))
-        arr4 = resize_3d.resize_3d_image(arr3, tf.constant([3, 3, 3]))
+    def test_rotate_45_deg_images_axis_9(self):
+        img1 = np.array(Image.open("tests/data/images/one.jpg"))
+        img2 = np.array(Image.open("tests/data/images/two.jpg"))
+        img3 = np.array(Image.open("tests/data/images/three.jpg"))
 
-        result = tf.logical_not(tf.cast(tf.reshape(arr2-arr4, [-1]), dtype=tf.bool)).numpy().all()
+        # make it grayscale
+        img1 = np.average(img1, axis = 2).astype(np.uint8)
+        img2 = np.average(img2, axis = 2).astype(np.uint8)
+        img3 = np.average(img3, axis = 2).astype(np.uint8)
 
-        self.assertEqual(result, True, msg = "resize twice and back failed")
+        img4 = np.stack([img1, img2, img3], axis = 2)
 
-    def test_resize_triple_and_back(self):
-        arr1 = tf.range(0, 27, dtype = tf.int32)
-        arr2 = tf.reshape(arr1, [3, 3, 3])
+        # normalize to [-1, 1]
+        min = np.min(img4)
+        max = np.max(img4)
+        img4 = (img4 - min) / (max - min) * 2 - 1
 
-        arr3 = resize_3d.resize_3d_image(arr2, tf.constant([9, 9, 9]))
-        arr4 = resize_3d.resize_3d_image(arr3, tf.constant([3, 3, 3]))
+        img5 = rotate.rotate_data(45, 2, img4)
+        img6 = rotate.rotate_data(-(360-45), 2, img4)
 
-        result = tf.logical_not(tf.cast(tf.reshape(arr2-arr4, [-1]), dtype=tf.bool)).numpy().all()
+        for i in range(3):
+            # rescale back to [0, 255]
+            img_to_save1 = img5[:, :, i] * 127.5 + 127.5
+            img_to_save2 = img6[:, :, i] * 127.5 + 127.5
 
-        self.assertEqual(result, True, msg = "resize triple and back failed")
+            rotated_img1 = Image.fromarray(img_to_save1.astype(np.uint8))
+            rotated_img2 = Image.fromarray(img_to_save2.astype(np.uint8))
 
-    def test_resize_triple_and_back_xl(self):
-        cube_size = 64
-        arr2 = tf.random.uniform([cube_size, cube_size, cube_size], dtype = tf.float32)
+            rotated_img1.save(f"tests/data/images/rotated_{i}_45.jpg")
+            rotated_img2.save(f"tests/data/images/rotated_{i}_{360-45}.jpg")
 
-        arr3 = resize_3d.resize_3d_image(arr2, tf.constant([cube_size*3, cube_size*3, cube_size*3]))
-        arr4 = resize_3d.resize_3d_image(arr3, tf.constant([cube_size, cube_size, cube_size]))
+        self.assertAlmostEqual(np.sum(img5 - img6), 0, places=2, msg = "rotation 45 and -45 failed")
 
-        result = tf.logical_not(tf.cast(tf.reshape(arr2-arr4, [-1]), dtype=tf.bool)).numpy().all()
 
-        self.assertEqual(result, True, msg = "resize triple and back failed")
+    def test_rotate_45_deg_images_axis_1(self):
+        img1 = np.array(Image.open("tests/data/images/one.jpg"))
+        img2 = np.array(Image.open("tests/data/images/two.jpg"))
+        img3 = np.array(Image.open("tests/data/images/three.jpg"))
+
+        # make it grayscale
+        img1 = np.average(img1, axis = 2).astype(np.uint8)
+        img2 = np.average(img2, axis = 2).astype(np.uint8)
+        img3 = np.average(img3, axis = 2).astype(np.uint8)
+
+        img4 = np.stack([img1, img2, img3], axis = 2)
+
+        # normalize to [-1, 1]
+        min = np.min(img4)
+        max = np.max(img4)
+        img4 = (img4 - min) / (max - min) * 2 - 1
+
+        img5 = rotate.rotate_data(45, 0, img4)
+        img6 = rotate.rotate_data(-(360-45), 0, img4)
+
+        # saving images does not make sense here, due to rotation around axis 1
+
+        self.assertAlmostEqual(np.sum(img5 - img6), 0, places=2, msg = "rotation 45 and -45 failed")
+
+
+
+
+    # def test_resize_triple_and_back(self):
+    #     arr1 = tf.range(0, 27, dtype = tf.int32)
+    #     arr2 = tf.reshape(arr1, [3, 3, 3])
+
+    #     arr3 = resize_3d.resize_3d_image(arr2, tf.constant([9, 9, 9]))
+    #     arr4 = resize_3d.resize_3d_image(arr3, tf.constant([3, 3, 3]))
+
+    #     result = tf.logical_not(tf.cast(tf.reshape(arr2-arr4, [-1]), dtype=tf.bool)).numpy().all()
+
+    #     self.assertEqual(result, True, msg = "resize triple and back failed")
+
+    # def test_resize_triple_and_back_xl(self):
+    #     cube_size = 64
+    #     arr2 = tf.random.uniform([cube_size, cube_size, cube_size], dtype = tf.float32)
+
+    #     arr3 = resize_3d.resize_3d_image(arr2, tf.constant([cube_size*3, cube_size*3, cube_size*3]))
+    #     arr4 = resize_3d.resize_3d_image(arr3, tf.constant([cube_size, cube_size, cube_size]))
+
+    #     result = tf.logical_not(tf.cast(tf.reshape(arr2-arr4, [-1]), dtype=tf.bool)).numpy().all()
+
+    #     self.assertEqual(result, True, msg = "resize triple and back failed")
 
 
 if __name__ == "__main__":