test_flip_image.py 1.32 KB
Newer Older
BO ZHANG's avatar
BO ZHANG committed
1
2
3
4
5
import unittest

import numpy as np

from csst_proto.top_level_interface import flip_image, read_test_image
BO ZHANG's avatar
BO ZHANG committed
6
from csst_proto.flip_image import flip_multiple_images_jl, flip_multiple_images_mp
BO ZHANG's avatar
BO ZHANG committed
7
8
9
10


class FlipImageTestCase(unittest.TestCase):
    def test_flip_image(self):
BO ZHANG's avatar
BO ZHANG committed
11
        """ test flip image """
BO ZHANG's avatar
BO ZHANG committed
12
13
14
15
16
17
18
        self.assertTrue(
            np.all(flip_image(read_test_image()) == np.array([[4, 3], [2, 1]]))
        )

        # the code fails for 1D array
        with self.assertRaises(AssertionError):
            flip_image(np.array([1, 2, 3, 4]))
BO ZHANG's avatar
BO ZHANG committed
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

    def test_flip_multiple_images_mp(self):
        """ test flip multiple images with multiprocessing """
        n_jobs = 10
        imgs = [read_test_image() for i_job in range(n_jobs)]
        flipped_imgs = flip_multiple_images_mp(imgs, n_jobs)
        for i_job in range(n_jobs):
            self.assertTrue(
                np.all(flipped_imgs[i_job] == np.array([[4, 3], [2, 1]]))
            )

    def test_flip_multiple_images_jl(self):
        """ test flip multiple images with joblib """
        n_jobs = 10
        imgs = [read_test_image() for i_job in range(n_jobs)]
        flipped_imgs = flip_multiple_images_jl(imgs, n_jobs)
        for i_job in range(n_jobs):
            self.assertTrue(
                np.all(flipped_imgs[i_job] == np.array([[4, 3], [2, 1]]))
            )