Skip to content

Commit 9511327

Browse files
committed
RF: make default trk structured array a scalar
When making the default trk header structured array, make a numpy array scalar instead of a 1D array. Adjust the code accordingly. Take opportunity to rename header to st_arr to make clear it is not the 'header' of `self.header` (which is a dict).
1 parent 3dc3d71 commit 9511327

File tree

1 file changed

+18
-16
lines changed

1 file changed

+18
-16
lines changed

nibabel/streamlines/trk.py

Lines changed: 18 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -72,8 +72,8 @@ def get_affine_trackvis_to_rasmm(header):
7272
7373
Parameters
7474
----------
75-
header : dict
76-
Dict containing trackvis header.
75+
header : dict or ndarray
76+
Dict or numpy structured array containing trackvis header.
7777
7878
Returns
7979
-------
@@ -101,9 +101,12 @@ def get_affine_trackvis_to_rasmm(header):
101101
# If the voxel order implied by the affine does not match the voxel
102102
# order in the TRK header, change the orientation.
103103
# voxel (header) -> voxel (affine)
104-
header_ornt = asstr(header[Field.VOXEL_ORDER])
104+
vox_order = header[Field.VOXEL_ORDER]
105+
# Input header can be dict or structured array
106+
if hasattr(vox_order, 'item'): # structured array
107+
vox_order = header[Field.VOXEL_ORDER].item()
105108
affine_ornt = "".join(aff2axcodes(header[Field.VOXEL_TO_RASMM]))
106-
header_ornt = axcodes2ornt(header_ornt)
109+
header_ornt = axcodes2ornt(vox_order.decode('latin1'))
107110
affine_ornt = axcodes2ornt(affine_ornt)
108111
ornt = nib.orientations.ornt_transform(header_ornt, affine_ornt)
109112
M = nib.orientations.inv_ornt_aff(ornt, header[Field.DIMENSIONS])
@@ -265,25 +268,25 @@ def is_correct_format(cls, fileobj):
265268
def _default_structarr(cls):
266269
""" Return an empty compliant TRK header as numpy structured array
267270
"""
268-
header = np.zeros(1, dtype=header_2_dtype)
271+
st_arr = np.zeros((), dtype=header_2_dtype)
269272

270273
# Default values
271-
header[Field.MAGIC_NUMBER] = cls.MAGIC_NUMBER
272-
header[Field.VOXEL_SIZES] = np.array((1, 1, 1), dtype="f4")
273-
header[Field.DIMENSIONS] = np.array((1, 1, 1), dtype="h")
274-
header[Field.VOXEL_TO_RASMM] = np.eye(4, dtype="f4")
275-
header[Field.VOXEL_ORDER] = b"RAS"
276-
header['version'] = 2
277-
header['hdr_size'] = cls.HEADER_SIZE
274+
st_arr[Field.MAGIC_NUMBER] = cls.MAGIC_NUMBER
275+
st_arr[Field.VOXEL_SIZES] = np.array((1, 1, 1), dtype="f4")
276+
st_arr[Field.DIMENSIONS] = np.array((1, 1, 1), dtype="h")
277+
st_arr[Field.VOXEL_TO_RASMM] = np.eye(4, dtype="f4")
278+
st_arr[Field.VOXEL_ORDER] = b"RAS"
279+
st_arr['version'] = 2
280+
st_arr['hdr_size'] = cls.HEADER_SIZE
278281

279-
return header
282+
return st_arr
280283

281284
@classmethod
282285
def create_empty_header(cls):
283286
""" Return an empty compliant TRK header as dict
284287
"""
285-
header_rec = cls._default_structarr()
286-
return dict(zip(header_rec.dtype.names, header_rec))
288+
st_arr = cls._default_structarr()
289+
return dict(zip(st_arr.dtype.names, st_arr.tolist()))
287290

288291
@classmethod
289292
def load(cls, fileobj, lazy_load=False):
@@ -410,7 +413,6 @@ def save(self, fileobj):
410413
nb_scalars = 0
411414
nb_properties = 0
412415

413-
header = header[0]
414416
with Opener(fileobj, mode="wb") as f:
415417
# Keep track of the beginning of the header.
416418
beginning = f.tell()

0 commit comments

Comments
 (0)