Skip to content
test_flip_image.py 1.96 KiB
Newer Older
import os
BO ZHANG's avatar
BO ZHANG committed
import unittest

import numpy as np

BO ZHANG's avatar
BO ZHANG committed
from csst_proto import flip_image, read_default_image
from csst_proto import flip_multiple_images_jl, flip_multiple_images_mp
BO ZHANG's avatar
BO ZHANG committed


class FlipImageTestCase(unittest.TestCase):
    def test_flip_image(self):
BO ZHANG's avatar
BO ZHANG committed
        """test flip image"""
BO ZHANG's avatar
BO ZHANG committed
        self.assertTrue(
BO ZHANG's avatar
BO ZHANG committed
            np.all(flip_image(read_default_image()) == np.array([[4, 3], [2, 1]]))
BO ZHANG's avatar
BO ZHANG committed
        )

        # 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

    def test_flip_image_on_server(self):
        """test flip image on server"""
BO ZHANG's avatar
BO ZHANG committed
        image_input = np.loadtxt(
            os.path.join(
                os.environ["UNIT_TEST_DATA_ROOT"],
                "csst_proto/test_flip/input/image.txt",
            ),
            delimiter=",",
            dtype=int,
BO ZHANG's avatar
BO ZHANG committed
        image_answer = np.loadtxt(
            os.path.join(
                os.environ["UNIT_TEST_DATA_ROOT"],
                "csst_proto/test_flip/answer/flipped_image.txt",
            ),
            delimiter=",",
            dtype=int,
BO ZHANG's avatar
BO ZHANG committed
        self.assertLess(
            np.linalg.norm(flip_image(image_input) - image_answer),
            1e-6,
            "Test flip image on server failed",
        )

BO ZHANG's avatar
BO ZHANG committed
    def test_flip_multiple_images_mp(self):
BO ZHANG's avatar
BO ZHANG committed
        """test flip multiple images with multiprocessing"""
BO ZHANG's avatar
BO ZHANG committed
        n_jobs = 10
        imgs = [read_default_image() for _ in range(n_jobs)]
BO ZHANG's avatar
BO ZHANG committed
        flipped_imgs = flip_multiple_images_mp(imgs, n_jobs)
        for i_job in range(n_jobs):
BO ZHANG's avatar
BO ZHANG committed
            self.assertTrue(np.all(flipped_imgs[i_job] == np.array([[4, 3], [2, 1]])))
BO ZHANG's avatar
BO ZHANG committed

    def test_flip_multiple_images_jl(self):
BO ZHANG's avatar
BO ZHANG committed
        """test flip multiple images with joblib"""
BO ZHANG's avatar
BO ZHANG committed
        n_jobs = 10
        imgs = [read_default_image() for _ in range(n_jobs)]
BO ZHANG's avatar
BO ZHANG committed
        flipped_imgs = flip_multiple_images_jl(imgs, n_jobs)
        for i_job in range(n_jobs):
BO ZHANG's avatar
BO ZHANG committed
            self.assertTrue(np.all(flipped_imgs[i_job] == np.array([[4, 3], [2, 1]])))