Skip to content

Commit 2a10bf7

Browse files
MarcCoteMarcCote
committed
ENH: adding TRK <=> TCK streamlines conversion scripts
1 parent e48b746 commit 2a10bf7

File tree

4 files changed

+205
-1
lines changed

4 files changed

+205
-1
lines changed

‎bin/nib-tck2trk

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
#!python
2+
# emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
3+
# vi: set ft=python sts=4 ts=4 sw=4 et:
4+
5+
"""
6+
Convert tractograms (TCK -> TRK).
7+
"""
8+
import os
9+
import argparse
10+
11+
import nibabel as nib
12+
13+
from nibabel.streamlines import Field
14+
from nibabel.orientations import aff2axcodes
15+
16+
17+
def parse_args():
18+
DESCRIPTION = "Convert tractograms (TCK -> TRK)."
19+
parser = argparse.ArgumentParser(description=DESCRIPTION)
20+
parser.add_argument("anatomy",
21+
help="reference anatomical image (.nii|.nii.gz.")
22+
parser.add_argument("tractograms", metavar="tractogram", nargs="+",
23+
help="list of tractograms (.tck).")
24+
parser.add_argument("-f", "--force", action="store_true",
25+
help="overwrite existing output files.")
26+
27+
args = parser.parse_args()
28+
return args, parser
29+
30+
31+
def main():
32+
args, parser = parse_args()
33+
34+
try:
35+
nii = nib.load(args.anatomy)
36+
except Exception:
37+
parser.error("Expecting anatomical image as first agument.")
38+
39+
for tractogram in args.tractograms:
40+
tractogram_format = nib.streamlines.detect_format(tractogram)
41+
if tractogram_format is not nib.streamlines.TckFile:
42+
print("Skipping non TCK file: '{}'".format(tractogram))
43+
continue
44+
45+
filename, _ = os.path.splitext(tractogram)
46+
output_filename = filename + '.trk'
47+
if os.path.isfile(output_filename) and not args.force:
48+
msg = "Skipping existing file: '{}'. Use -f to overwrite."
49+
print(msg.format(output_filename))
50+
continue
51+
52+
# Build header using infos from the anatomical image.
53+
header = {}
54+
header[Field.VOXEL_TO_RASMM] = nii.affine.copy()
55+
header[Field.VOXEL_SIZES] = nii.header.get_zooms()[:3]
56+
header[Field.DIMENSIONS] = nii.shape[:3]
57+
header[Field.VOXEL_ORDER] = "".join(aff2axcodes(nii.affine))
58+
59+
tck = nib.streamlines.load(tractogram)
60+
nib.streamlines.save(tck.tractogram, output_filename, header=header)
61+
62+
63+
if __name__ == '__main__':
64+
main()

‎bin/nib-trk2tck

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
#!python
2+
# emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
3+
# vi: set ft=python sts=4 ts=4 sw=4 et:
4+
5+
"""
6+
Convert tractograms (TRK -> TCK).
7+
"""
8+
9+
import os
10+
import argparse
11+
12+
import nibabel as nib
13+
14+
15+
def parse_args():
16+
DESCRIPTION = "Convert tractograms (TRK -> TCK)."
17+
parser = argparse.ArgumentParser(description=DESCRIPTION)
18+
parser.add_argument("tractograms", metavar="tractogram", nargs="+",
19+
help="list of tractograms (.trk).")
20+
parser.add_argument("-f", "--force", action="store_true",
21+
help="overwrite existing output files.")
22+
23+
args = parser.parse_args()
24+
return args, parser
25+
26+
27+
def main():
28+
args, parser = parse_args()
29+
for tractogram in args.tractograms:
30+
tractogram_format = nib.streamlines.detect_format(tractogram)
31+
if tractogram_format is not nib.streamlines.TrkFile:
32+
print("Skipping non TRK file: '{}'".format(tractogram))
33+
continue
34+
35+
filename, _ = os.path.splitext(tractogram)
36+
output_filename = filename + '.tck'
37+
if os.path.isfile(output_filename) and not args.force:
38+
msg = "Skipping existing file: '{}'. Use -f to overwrite."
39+
print(msg.format(output_filename))
40+
continue
41+
42+
trk = nib.streamlines.load(tractogram)
43+
nib.streamlines.save(trk.tractogram, output_filename)
44+
45+
46+
if __name__ == '__main__':
47+
main()

‎nibabel/tests/test_scripts.py

Lines changed: 92 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,21 +8,23 @@
88

99
import sys
1010
import os
11+
import shutil
1112
from os.path import (dirname, join as pjoin, abspath, splitext, basename,
1213
exists)
1314
import csv
1415
from glob import glob
1516

1617
import numpy as np
1718

19+
import nibabel as nib
1820
from ..tmpdirs import InTemporaryDirectory
1921
from ..loadsave import load
2022
from ..orientations import flip_axis, aff2axcodes, inv_ornt_aff
2123

2224
from nose.tools import assert_true, assert_false, assert_equal
2325
from nose import SkipTest
2426

25-
from numpy.testing import assert_almost_equal
27+
from numpy.testing import assert_almost_equal, assert_array_equal
2628

2729
from .scriptrunner import ScriptRunner
2830
from .nibabel_data import needs_nibabel_data
@@ -357,3 +359,92 @@ def test_parrec2nii_with_data():
357359
assert_equal(sorted(csv_keys), ['diffusion b value number',
358360
'gradient orientation number'])
359361
assert_equal(nlines, 8) # 8 volumes present in DTI.PAR
362+
363+
364+
@script_test
365+
def test_nib_trk2tck():
366+
simple_trk = pjoin(DATA_PATH, "simple.trk")
367+
standard_trk = pjoin(DATA_PATH, "standard.trk")
368+
369+
with InTemporaryDirectory() as tmpdir:
370+
# Copy input files to convert.
371+
shutil.copy(simple_trk, tmpdir)
372+
shutil.copy(standard_trk, tmpdir)
373+
simple_trk = pjoin(tmpdir, "simple.trk")
374+
standard_trk = pjoin(tmpdir, "standard.trk")
375+
simple_tck = pjoin(tmpdir, "simple.tck")
376+
standard_tck = pjoin(tmpdir, "standard.tck")
377+
378+
# Convert one file.
379+
cmd = ["nib-trk2tck", simple_trk]
380+
code, stdout, stderr = run_command(cmd)
381+
assert_equal(len(stdout), 0)
382+
assert_true(os.path.isfile(simple_tck))
383+
trk = nib.streamlines.load(simple_trk)
384+
tck = nib.streamlines.load(simple_tck)
385+
assert_array_equal(tck.streamlines.data, trk.streamlines.data)
386+
assert_true(isinstance(tck, nib.streamlines.TckFile))
387+
388+
# Skip non TRK files.
389+
cmd = ["nib-trk2tck", simple_tck]
390+
code, stdout, stderr = run_command(cmd)
391+
assert_true("Skipping non TRK file" in stdout)
392+
393+
# By default, refuse to overwrite existing output files.
394+
cmd = ["nib-trk2tck", simple_trk]
395+
code, stdout, stderr = run_command(cmd)
396+
assert_true("Skipping existing file" in stdout)
397+
398+
# Convert multiple files and with --force.
399+
cmd = ["nib-trk2tck", "--force", simple_trk, standard_trk]
400+
code, stdout, stderr = run_command(cmd)
401+
assert_equal(len(stdout), 0)
402+
trk = nib.streamlines.load(standard_trk)
403+
tck = nib.streamlines.load(standard_tck)
404+
assert_array_equal(tck.streamlines.data, trk.streamlines.data)
405+
406+
407+
@script_test
408+
def test_nib_tck2trk():
409+
anat = pjoin(DATA_PATH, "standard.nii.gz")
410+
standard_tck = pjoin(DATA_PATH, "standard.tck")
411+
412+
with InTemporaryDirectory() as tmpdir:
413+
# Copy input file to convert.
414+
shutil.copy(standard_tck, tmpdir)
415+
standard_trk = pjoin(tmpdir, "standard.trk")
416+
standard_tck = pjoin(tmpdir, "standard.tck")
417+
418+
# Anatomical image not found as first argument.
419+
cmd = ["nib-tck2trk", standard_tck, anat]
420+
code, stdout, stderr = run_command(cmd, check_code=False)
421+
assert_equal(code, 2) # Parser error.
422+
assert_true("Expecting anatomical image as first agument" in stderr)
423+
424+
# Convert one file.
425+
cmd = ["nib-tck2trk", anat, standard_tck]
426+
code, stdout, stderr = run_command(cmd)
427+
assert_equal(len(stdout), 0)
428+
assert_true(os.path.isfile(standard_trk))
429+
tck = nib.streamlines.load(standard_tck)
430+
trk = nib.streamlines.load(standard_trk)
431+
assert_array_equal(trk.streamlines.data, tck.streamlines.data)
432+
assert_true(isinstance(trk, nib.streamlines.TrkFile))
433+
434+
# Skip non TCK files.
435+
cmd = ["nib-tck2trk", anat, standard_trk]
436+
code, stdout, stderr = run_command(cmd)
437+
assert_true("Skipping non TCK file" in stdout)
438+
439+
# By default, refuse to overwrite existing output files.
440+
cmd = ["nib-tck2trk", anat, standard_tck]
441+
code, stdout, stderr = run_command(cmd)
442+
assert_true("Skipping existing file" in stdout)
443+
444+
# Convert multiple files and with --force.
445+
cmd = ["nib-tck2trk", "--force", anat, standard_tck, standard_tck]
446+
code, stdout, stderr = run_command(cmd)
447+
assert_equal(len(stdout), 0)
448+
tck = nib.streamlines.load(standard_tck)
449+
trk = nib.streamlines.load(standard_trk)
450+
assert_array_equal(tck.streamlines.data, trk.streamlines.data)

‎setup.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,8 @@ def main(**extra_args):
115115
pjoin('bin', 'nib-ls'),
116116
pjoin('bin', 'nib-dicomfs'),
117117
pjoin('bin', 'nib-nifti-dx'),
118+
pjoin('bin', 'nib-tck2trk'),
119+
pjoin('bin', 'nib-trk2tck'),
118120
],
119121
cmdclass = cmdclass,
120122
**extra_args

0 commit comments

Comments
 (0)