lamino-filters/tests/tests.py

import os
import shutil
import tempfile

try:
    import unittest2 as unittest
except ImportError:
    import unittest

import numpy as np
from gi.repository import Ufo
from libtiff import TIFF


def data_path(suffix=''):
    return os.path.join(os.path.abspath('./data'), suffix)


class BasicTests(unittest.TestCase):
    def setUp(self):
        config = Ufo.Config(paths=['../build/src', '../build_dir/src'])
        self.pm = Ufo.PluginManager(config=config)
        self.graph = Ufo.TaskGraph()
        self.sched = Ufo.Scheduler(config=config)
        self.tmpdir = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.tmpdir)

    def tmp_path(self, suffix=''):
        return os.path.join(self.tmpdir, suffix)

    def get_task(self, name, **kwargs):
        plugin = self.pm.get_task(name)
        plugin.set_properties(**kwargs)
        return plugin

    def test_ramp(self):
        outname = self.tmp_path('rmp-%05i.tif')
        ramp = self.pm.get_task('lamino-ramp')
        writer = self.pm.get_task('writer')

        ramp.set_properties(width=4096, fwidth=1458, height=2048,
                            theta=1.0921048, tau=0.3)
        writer.set_properties(filename=outname)

        self.graph.connect_nodes(ramp, writer)
        self.sched.run(self.graph)

        files = os.listdir(self.tmpdir)
        self.assertIn('rmp-00000.tif', files)
        self.assertEqual(len(files), 1)

    def test_bug_lamino_reco(self):
        reader = self.get_task('reader', path=data_path('bug/projections/*.tif'))
        writer = self.get_task('null')
        reco = self.get_task('lamino-bp')
        padding = self.get_task('padding-2d')
        fft_input = self.get_task('fft', dimensions=2)
        ramp = self.get_task('lamino-ramp')
        fft_ramp = self.get_task('fft', dimensions=2)
        conv = self.get_task('lamino-conv')
        ifft = self.get_task('ifft', dimensions=2)

        THETA = -1.570796326
        reco.set_properties(theta=THETA, psi=0, angle_step=0.01256637,
		                    vol_sx=832, vol_sy=832, vol_sz=400,
		                    vol_ox=416, vol_oy=416, vol_oz=200,
		                    proj_ox=592, proj_oy=205)

        ramp.set_properties(width=4096, fwidth=1186, height=512,
                            theta=THETA, tau=1)

        self.graph.connect_nodes(reader, padding)
        self.graph.connect_nodes(padding, fft_input)
        self.graph.connect_nodes(ramp, fft_ramp)

        self.graph.connect_nodes_full(fft_ramp, conv, 0)
        self.graph.connect_nodes_full(fft_input, conv, 1)

        self.graph.connect_nodes(conv, ifft)
        self.graph.connect_nodes(ifft, reco)
        self.graph.connect_nodes(reco, writer)
        self.sched.run(self.graph)

    def test_3d_edf(self):
        reader = self.get_task('reader', path=data_path('bug/projections'), count=1)
        writer = self.get_task('edf-3d-writer', filename=self.tmp_path('vol-%05i.edf'))
        reco = self.get_task('lamino-bp')

        vx, vy, vz = 256, 256, 128
        reco.set_properties(vol_sx=vx, vol_sy=vy, vol_sz=vz,
		                    vol_ox=0, vol_oy=0, vol_oz=0,
		                    proj_ox=592, proj_oy=205)

        self.graph.connect_nodes(reader, reco)
        self.graph.connect_nodes(reco, writer)
        self.sched.run(self.graph)

        files = os.listdir(self.tmpdir)
        self.assertIn('vol-00000.edf', files)
        self.assertEqual(len(files), 1)

        filesize = os.path.getsize(self.tmp_path('vol-00000.edf'))
        expected = vx * vy * vz * 4 + 512   # floats + header
        self.assertEqual(filesize, expected)