# -*- coding: utf-8 -*- from copy import deepcopy import os import psutil from tempfile import TemporaryDirectory import zipfile import pytest import numpy as np from numpy.testing import assert_allclose try: from dipy.io.streamline import save_tractogram, load_tractogram dipy_available = True except ImportError: dipy_available = False import trx.trx_file_memmap as tmm from trx.trx_file_memmap import TrxFile from trx.io import load, save, get_trx_tmp_dir from trx.fetcher import (get_testing_files_dict, fetch_data, get_home) fetch_data(get_testing_files_dict(), keys=['gold_standard.zip']) @pytest.mark.parametrize("path", [("gs.trk"), ("gs.tck"), ("gs.vtk")]) @pytest.mark.skipif(not dipy_available, reason='Dipy is not installed.') def test_seq_ops_sft(path): with TemporaryDirectory() as tmp_dir: gs_dir = os.path.join(get_home(), 'gold_standard') path = os.path.join(tmp_dir, path) obj = load(os.path.join(gs_dir, 'gs.trx'), os.path.join(gs_dir, 'gs.nii')) sft_1 = obj.to_sft() save_tractogram(sft_1, path) obj.close() save_tractogram(sft_1, os.path.join(tmp_dir, 'tmp.trk')) sft_2 = load_tractogram(os.path.join(tmp_dir, 'tmp.trk'), 'same') def test_seq_ops_trx(): with TemporaryDirectory() as tmp_dir: gs_dir = os.path.join(get_home(), 'gold_standard') path = os.path.join(gs_dir, 'gs.trx') trx_1 = tmm.load(path) tmm.save(trx_1, os.path.join(tmp_dir, 'tmp.trx')) trx_1.close() trx_2 = tmm.load(os.path.join(tmp_dir, 'tmp.trx')) trx_2.close() @pytest.mark.parametrize("path", [("gs.trx"), ("gs.trk"), ("gs.tck"), ("gs.vtk")]) @pytest.mark.skipif(not dipy_available, reason='Dipy is not installed.') def test_load_vox(path): gs_dir = os.path.join(get_home(), 'gold_standard') path = os.path.join(gs_dir, path) coord = np.loadtxt(os.path.join(get_home(), 'gold_standard', 'gs_vox_space.txt')) obj = load(path, os.path.join(gs_dir, 'gs.nii')) sft = obj.to_sft() if isinstance(obj, TrxFile) else obj sft.to_vox() assert_allclose(sft.streamlines._data, coord, rtol=1e-04, atol=1e-06) if isinstance(obj, TrxFile): obj.close() @pytest.mark.parametrize("path", [("gs.trx"), ("gs.trk"), ("gs.tck"), ("gs.vtk")]) @pytest.mark.skipif(not dipy_available, reason='Dipy is not installed.') def test_load_voxmm(path): gs_dir = os.path.join(get_home(), 'gold_standard') path = os.path.join(gs_dir, path) coord = np.loadtxt(os.path.join(get_home(), 'gold_standard', 'gs_voxmm_space.txt')) obj = load(path, os.path.join(gs_dir, 'gs.nii')) sft = obj.to_sft() if isinstance(obj, TrxFile) else obj sft.to_voxmm() assert_allclose(sft.streamlines._data, coord, rtol=1e-04, atol=1e-06) if isinstance(obj, TrxFile): obj.close() @pytest.mark.parametrize("path", [("gs.trk"), ("gs.trx"), ("gs_fldr.trx")]) @pytest.mark.skipif(not dipy_available, reason='Dipy is not installed.') def test_multi_load_save_rasmm(path): with TemporaryDirectory() as tmp_gs_dir: gs_dir = os.path.join(get_home(), 'gold_standard') basename, ext = os.path.splitext(path) path = os.path.join(gs_dir, path) coord = np.loadtxt(os.path.join(get_home(), 'gold_standard', 'gs_rasmm_space.txt')) obj = load(path, os.path.join(gs_dir, 'gs.nii')) for i in range(3): out_path = os.path.join( tmp_gs_dir, '{}_tmp{}_{}'.format(basename, i, ext)) save(obj, out_path) if isinstance(obj, TrxFile): obj.close() obj = load(out_path, os.path.join(gs_dir, 'gs.nii')) assert_allclose(obj.streamlines._data, coord, rtol=1e-04, atol=1e-06) if isinstance(obj, TrxFile): obj.close() @pytest.mark.parametrize("path", [("gs.trx"), ("gs_fldr.trx")]) @pytest.mark.skipif(not dipy_available, reason='Dipy is not installed.') def test_delete_tmp_gs_dir(path): gs_dir = os.path.join(get_home(), 'gold_standard') path = os.path.join(gs_dir, path) trx1 = tmm.load(path) if os.path.isfile(path): tmp_gs_dir = deepcopy(trx1._uncompressed_folder_handle.name) assert os.path.isdir(tmp_gs_dir) sft = trx1.to_sft() trx1.close() coord_rasmm = np.loadtxt(os.path.join(get_home(), 'gold_standard', 'gs_rasmm_space.txt')) coord_vox = np.loadtxt(os.path.join(get_home(), 'gold_standard', 'gs_vox_space.txt')) # The folder trx representation does not need tmp files if os.path.isfile(path): assert not os.path.isdir(tmp_gs_dir) assert_allclose(sft.streamlines._data, coord_rasmm, rtol=1e-04, atol=1e-06) # Reloading the TRX and checking its data, then closing trx2 = tmm.load(path) assert_allclose(trx2.streamlines._data, sft.streamlines._data, rtol=1e-04, atol=1e-06) trx2.close() sft.to_vox() assert_allclose(sft.streamlines._data, coord_vox, rtol=1e-04, atol=1e-06) trx3 = tmm.load(path) assert_allclose(trx3.streamlines._data, coord_rasmm, rtol=1e-04, atol=1e-06) trx3.close() @pytest.mark.parametrize("path", [("gs.trx")]) @pytest.mark.skipif(not dipy_available, reason='Dipy is not installed.') def test_close_tmp_files(path): gs_dir = os.path.join(get_home(), 'gold_standard') path = os.path.join(gs_dir, path) trx = tmm.load(path) process = psutil.Process(os.getpid()) open_files = process.open_files() expected_content = ['offsets.uint32', 'positions.3.float32', 'header.json', 'random_coord.3.float32', 'color_y.float32', 'color_x.float32', 'color_z.float32'] count = 0 for open_file in open_files: basename = os.path.basename(open_file.path) if basename in expected_content: count += 1 assert count == 6 trx.close() open_files = process.open_files() count = 0 for open_file in open_files: basename = os.path.basename(open_file.path) if basename in expected_content: count += 1 assert not count @pytest.mark.parametrize("tmp_path", [("~"), ("use_working_dir")]) def test_change_tmp_dir(tmp_path): gs_dir = os.path.join(get_home(), 'gold_standard') path = os.path.join(gs_dir, 'gs.trx') if tmp_path == 'use_working_dir': os.environ['TRX_TMPDIR'] = 'use_working_dir' else: os.environ['TRX_TMPDIR'] = os.path.expanduser(tmp_path) trx = tmm.load(path) tmp_gs_dir = deepcopy(trx._uncompressed_folder_handle.name) if tmp_path == 'use_working_dir': assert os.path.dirname(tmp_gs_dir) == os.getcwd() else: assert os.path.dirname(tmp_gs_dir) == os.path.expanduser(tmp_path) trx.close() assert not os.path.isdir(tmp_gs_dir) @pytest.mark.parametrize("path", [("gs.trx"), ("gs_fldr.trx")]) def test_complete_dir_from_trx(path): gs_dir = os.path.join(get_home(), 'gold_standard') path = os.path.join(gs_dir, path) trx = tmm.load(path) if trx._uncompressed_folder_handle is None: dir_to_check = path else: dir_to_check = trx._uncompressed_folder_handle.name file_paths = [] for dirpath, _, filenames in os.walk(dir_to_check): for filename in filenames: full_path = os.path.join(dirpath, filename) cut_path = full_path.split(dir_to_check)[1][1:].replace('\\', '/') file_paths.append(cut_path) expected_content = ['offsets.uint32', 'positions.3.float32', 'header.json', 'dps/random_coord.3.float32', 'dpv/color_y.float32', 'dpv/color_x.float32', 'dpv/color_z.float32'] assert set(file_paths) == set(expected_content) def test_complete_zip_from_trx(): gs_dir = os.path.join(get_home(), 'gold_standard') path = os.path.join(gs_dir, 'gs.trx') with zipfile.ZipFile(path, mode="r") as zf: zip_file_list = zf.namelist() expected_content = ['offsets.uint32', 'positions.3.float32', 'header.json', 'dps/random_coord.3.float32', 'dpv/color_y.float32', 'dpv/color_x.float32', 'dpv/color_z.float32'] assert set(zip_file_list) == set(expected_content)