Merge pull request #606 from MarcCote/enh_conversion_scripts

ENH: adding TRK <=> TCK streamlines conversion scripts

bin/nib-tck2trk

@@ -0,0 +1,18 @@
+#!python
+# emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
+# vi: set ft=python sts=4 ts=4 sw=4 et:
+### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
+#
+#   See COPYING file distributed along with the NiBabel package for the
+#   copyright and license terms.
+#
+### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
+"""
+Convert tractograms (TCK -> TRK).
+"""
+
+from nibabel.cmdline.tck2trk import main
+
+
+if __name__ == '__main__':
+    main()

bin/nib-trk2tck

@@ -0,0 +1,18 @@
+#!python
+# emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
+# vi: set ft=python sts=4 ts=4 sw=4 et:
+### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
+#
+#   See COPYING file distributed along with the NiBabel package for the
+#   copyright and license terms.
+#
+### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
+"""
+Convert tractograms (TRK -> TCK).
+"""
+
+from nibabel.cmdline.trk2tck import main
+
+
+if __name__ == '__main__':
+    main()

nibabel/cmdline/tck2trk.py

@@ -0,0 +1,56 @@
+"""
+Convert tractograms (TCK -> TRK).
+"""
+import os
+import argparse
+
+import nibabel as nib
+
+from nibabel.streamlines import Field
+from nibabel.orientations import aff2axcodes
+
+
+def parse_args():
+    DESCRIPTION = "Convert tractograms (TCK -> TRK)."
+    parser = argparse.ArgumentParser(description=DESCRIPTION)
+    parser.add_argument("anatomy",
+                        help="reference anatomical image (.nii|.nii.gz.")
+    parser.add_argument("tractograms", metavar="tractogram", nargs="+",
+                        help="list of tractograms (.tck).")
+    parser.add_argument("-f", "--force", action="store_true",
+                        help="overwrite existing output files.")
+
+    args = parser.parse_args()
+    return args, parser
+
+
+def main():
+    args, parser = parse_args()
+
+    try:
+        nii = nib.load(args.anatomy)
+    except Exception:
+        parser.error("Expecting anatomical image as first agument.")
+
+    for tractogram in args.tractograms:
+        tractogram_format = nib.streamlines.detect_format(tractogram)
+        if tractogram_format is not nib.streamlines.TckFile:
+            print("Skipping non TCK file: '{}'".format(tractogram))
+            continue
+
+        filename, _ = os.path.splitext(tractogram)
+        output_filename = filename + '.trk'
+        if os.path.isfile(output_filename) and not args.force:
+            msg = "Skipping existing file: '{}'. Use -f to overwrite."
+            print(msg.format(output_filename))
+            continue
+
+        # Build header using infos from the anatomical image.
+        header = {}
+        header[Field.VOXEL_TO_RASMM] = nii.affine.copy()
+        header[Field.VOXEL_SIZES] = nii.header.get_zooms()[:3]
+        header[Field.DIMENSIONS] = nii.shape[:3]
+        header[Field.VOXEL_ORDER] = "".join(aff2axcodes(nii.affine))
+
+        tck = nib.streamlines.load(tractogram)
+        nib.streamlines.save(tck.tractogram, output_filename, header=header)

nibabel/cmdline/trk2tck.py

@@ -0,0 +1,39 @@
+"""
+Convert tractograms (TRK -> TCK).
+"""
+
+import os
+import argparse
+
+import nibabel as nib
+
+
+def parse_args():
+    DESCRIPTION = "Convert tractograms (TRK -> TCK)."
+    parser = argparse.ArgumentParser(description=DESCRIPTION)
+    parser.add_argument("tractograms", metavar="tractogram", nargs="+",
+                        help="list of tractograms (.trk).")
+    parser.add_argument("-f", "--force", action="store_true",
+                        help="overwrite existing output files.")
+
+    args = parser.parse_args()
+    return args, parser
+
+
+def main():
+    args, parser = parse_args()
+    for tractogram in args.tractograms:
+        tractogram_format = nib.streamlines.detect_format(tractogram)
+        if tractogram_format is not nib.streamlines.TrkFile:
+            print("Skipping non TRK file: '{}'".format(tractogram))
+            continue
+
+        filename, _ = os.path.splitext(tractogram)
+        output_filename = filename + '.tck'
+        if os.path.isfile(output_filename) and not args.force:
+            msg = "Skipping existing file: '{}'. Use -f to overwrite."
+            print(msg.format(output_filename))
+            continue
+
+        trk = nib.streamlines.load(tractogram)
+        nib.streamlines.save(trk.tractogram, output_filename)

nibabel/tests/test_scripts.py

@@ -8,21 +8,23 @@
 
 import sys
 import os
+import shutil
 from os.path import (dirname, join as pjoin, abspath, splitext, basename,
                      exists)
 import csv
 from glob import glob
 
 import numpy as np
 
+import nibabel as nib
 from ..tmpdirs import InTemporaryDirectory
 from ..loadsave import load
 from ..orientations import flip_axis, aff2axcodes, inv_ornt_aff
 
 from nose.tools import assert_true, assert_false, assert_equal
 from nose import SkipTest
 
-from numpy.testing import assert_almost_equal
+from numpy.testing import assert_almost_equal, assert_array_equal
 
 from .scriptrunner import ScriptRunner
 from .nibabel_data import needs_nibabel_data
@@ -357,3 +359,92 @@ def test_parrec2nii_with_data():
         assert_equal(sorted(csv_keys), ['diffusion b value number',
                                         'gradient orientation number'])
         assert_equal(nlines, 8)  # 8 volumes present in DTI.PAR
+
+
+@script_test
+def test_nib_trk2tck():
+    simple_trk = pjoin(DATA_PATH, "simple.trk")
+    standard_trk = pjoin(DATA_PATH, "standard.trk")
+
+    with InTemporaryDirectory() as tmpdir:
+        # Copy input files to convert.
+        shutil.copy(simple_trk, tmpdir)
+        shutil.copy(standard_trk, tmpdir)
+        simple_trk = pjoin(tmpdir, "simple.trk")
+        standard_trk = pjoin(tmpdir, "standard.trk")
+        simple_tck = pjoin(tmpdir, "simple.tck")
+        standard_tck = pjoin(tmpdir, "standard.tck")
+
+        # Convert one file.
+        cmd = ["nib-trk2tck", simple_trk]
+        code, stdout, stderr = run_command(cmd)
+        assert_equal(len(stdout), 0)
+        assert_true(os.path.isfile(simple_tck))
+        trk = nib.streamlines.load(simple_trk)
+        tck = nib.streamlines.load(simple_tck)
+        assert_array_equal(tck.streamlines.data, trk.streamlines.data)
+        assert_true(isinstance(tck, nib.streamlines.TckFile))
+
+        # Skip non TRK files.
+        cmd = ["nib-trk2tck", simple_tck]
+        code, stdout, stderr = run_command(cmd)
+        assert_true("Skipping non TRK file" in stdout)
+
+        # By default, refuse to overwrite existing output files.
+        cmd = ["nib-trk2tck", simple_trk]
+        code, stdout, stderr = run_command(cmd)
+        assert_true("Skipping existing file" in stdout)
+
+        # Convert multiple files and with --force.
+        cmd = ["nib-trk2tck", "--force", simple_trk, standard_trk]
+        code, stdout, stderr = run_command(cmd)
+        assert_equal(len(stdout), 0)
+        trk = nib.streamlines.load(standard_trk)
+        tck = nib.streamlines.load(standard_tck)
+        assert_array_equal(tck.streamlines.data, trk.streamlines.data)
+
+
+@script_test
+def test_nib_tck2trk():
+    anat = pjoin(DATA_PATH, "standard.nii.gz")
+    standard_tck = pjoin(DATA_PATH, "standard.tck")
+
+    with InTemporaryDirectory() as tmpdir:
+        # Copy input file to convert.
+        shutil.copy(standard_tck, tmpdir)
+        standard_trk = pjoin(tmpdir, "standard.trk")
+        standard_tck = pjoin(tmpdir, "standard.tck")
+
+        # Anatomical image not found as first argument.
+        cmd = ["nib-tck2trk", standard_tck, anat]
+        code, stdout, stderr = run_command(cmd, check_code=False)
+        assert_equal(code, 2)  # Parser error.
+        assert_true("Expecting anatomical image as first agument" in stderr)
+
+        # Convert one file.
+        cmd = ["nib-tck2trk", anat, standard_tck]
+        code, stdout, stderr = run_command(cmd)
+        assert_equal(len(stdout), 0)
+        assert_true(os.path.isfile(standard_trk))
+        tck = nib.streamlines.load(standard_tck)
+        trk = nib.streamlines.load(standard_trk)
+        assert_array_equal(trk.streamlines.data, tck.streamlines.data)
+        assert_true(isinstance(trk, nib.streamlines.TrkFile))
+
+        # Skip non TCK files.
+        cmd = ["nib-tck2trk", anat, standard_trk]
+        code, stdout, stderr = run_command(cmd)
+        assert_true("Skipping non TCK file" in stdout)
+
+        # By default, refuse to overwrite existing output files.
+        cmd = ["nib-tck2trk", anat, standard_tck]
+        code, stdout, stderr = run_command(cmd)
+        assert_true("Skipping existing file" in stdout)
+
+        # Convert multiple files and with --force.
+        cmd = ["nib-tck2trk", "--force", anat, standard_tck, standard_tck]
+        code, stdout, stderr = run_command(cmd)
+        assert_equal(len(stdout), 0)
+        tck = nib.streamlines.load(standard_tck)
+        trk = nib.streamlines.load(standard_trk)
+        assert_array_equal(tck.streamlines.data, trk.streamlines.data)

setup.py

@@ -117,6 +117,8 @@ def main(**extra_args):
                           pjoin('bin', 'nib-ls'),
                           pjoin('bin', 'nib-dicomfs'),
                           pjoin('bin', 'nib-nifti-dx'),
+                          pjoin('bin', 'nib-tck2trk'),
+                          pjoin('bin', 'nib-trk2tck'),
                           ],
           cmdclass = cmdclass,
           **extra_args