add test augmentation and remove illumination module

fish-quant · Jan 18, 2022 · 999de09 · 999de09
1 parent d56b134
commit 999de09
Show file tree

Hide file tree

Showing 3 changed files with 84 additions and 103 deletions.
diff --git a/bigfish/stack/augmentation.py b/bigfish/stack/augmentation.py
@@ -30,11 +30,7 @@ def augment_2d(image):
     check_array(image, ndim=[2, 3])
 
     # randomly choose an operator
-    operations = [_identity,
-                  _flip_h, _flip_v,
-                  _transpose, _transpose_inverse,
-                  _rotation_90, _rotation_180, _rotation_270]
-    random_operation = np.random.choice(operations)
+    random_operation = augment_2d_function()
 
     # augment the image
     image_augmented = random_operation(image)

diff --git a/bigfish/stack/illumination.py b/bigfish/stack/illumination.py
diff --git a/bigfish/stack/tests/test_augmentation.py b/bigfish/stack/tests/test_augmentation.py
@@ -140,3 +140,86 @@ def test_augment_2d_dtype(dtype):
                   [0, 0, 0, 0, 0]], dtype=dtype)
     y = stack.augment_2d(x)
     assert y.dtype == dtype
+
+
+def test_augment_2d_function():
+    operations = [_identity,
+                  _flip_h, _flip_v,
+                  _transpose, _transpose_inverse,
+                  _rotation_90, _rotation_180, _rotation_270]
+    bytecodes = []
+    for f in operations:
+        bytecode = f.__code__.co_code
+        bytecodes.append(bytecode)
+    f = stack.augment_2d_function()
+    assert f.__code__.co_code in bytecodes
+    f = stack.augment_2d_function(identity=True)
+    assert f.__code__.co_code == _identity.__code__.co_code
+
+
+def test_augment_8_times():
+    # one channel
+    expected_y_identity = x.copy()
+    expected_y_flip_h = np.array([[0, 0, 0, 0, 0],
+                                  [0, 1, 1, 1, 0],
+                                  [0, 1, 0, 0, 0],
+                                  [0, 1, 0, 0, 0],
+                                  [1, 0, 0, 0, 0]], dtype=np.uint8)
+    expected_y_flip_v = np.array([[0, 0, 0, 0, 1],
+                                  [0, 0, 0, 1, 0],
+                                  [0, 0, 0, 1, 0],
+                                  [0, 1, 1, 1, 0],
+                                  [0, 0, 0, 0, 0]], dtype=np.uint8)
+    expected_y_transpose = np.array([[1, 0, 0, 0, 0],
+                                     [0, 1, 1, 1, 0],
+                                     [0, 0, 0, 1, 0],
+                                     [0, 0, 0, 1, 0],
+                                     [0, 0, 0, 0, 0]], dtype=np.uint8)
+    expected_y_transpose_inverse = np.array([[0, 0, 0, 0, 0],
+                                             [0, 1, 0, 0, 0],
+                                             [0, 1, 0, 0, 0],
+                                             [0, 1, 1, 1, 0],
+                                             [0, 0, 0, 0, 1]], dtype=np.uint8)
+    expected_y_rotation_90 = np.array([[0, 0, 0, 0, 1],
+                                       [0, 1, 1, 1, 0],
+                                       [0, 1, 0, 0, 0],
+                                       [0, 1, 0, 0, 0],
+                                       [0, 0, 0, 0, 0]], dtype=np.uint8)
+    expected_y_rotation_180 = np.array([[0, 0, 0, 0, 0],
+                                        [0, 1, 1, 1, 0],
+                                        [0, 0, 0, 1, 0],
+                                        [0, 0, 0, 1, 0],
+                                        [0, 0, 0, 0, 1]], dtype=np.uint8)
+    expected_y_rotation_270 = np.array([[0, 0, 0, 0, 0],
+                                        [0, 0, 0, 1, 0],
+                                        [0, 0, 0, 1, 0],
+                                        [0, 1, 1, 1, 0],
+                                        [1, 0, 0, 0, 0]], dtype=np.uint8)
+    expected_y = [expected_y_identity,
+                  expected_y_flip_h, expected_y_flip_v,
+                  expected_y_transpose, expected_y_transpose_inverse,
+                  expected_y_rotation_90, expected_y_rotation_180,
+                  expected_y_rotation_270]
+    augmented_arrays = stack.augment_8_times(x)
+    assert isinstance(augmented_arrays, list)
+    assert len(augmented_arrays) == len(expected_y)
+    for a, b in zip(augmented_arrays, expected_y):
+        assert_array_equal(a, b)
+
+    # multichannel
+    xx = x[..., np.newaxis]
+    expected_yy = [y[..., np.newaxis] for y in expected_y]
+    augmented_arrays = stack.augment_8_times(xx)
+    assert isinstance(augmented_arrays, list)
+    assert len(augmented_arrays) == len(expected_yy)
+    for a, b in zip(augmented_arrays, expected_yy):
+        assert_array_equal(a, b)
+
+
+def test_augment_8_times_reversed():
+    y = stack.augment_8_times(x)
+    y_reversed = stack.augment_8_times_reversed(y)
+    assert isinstance(y_reversed, list)
+    assert len(y_reversed) == 8
+    for a in y_reversed:
+        assert_array_equal(a, x.copy())