probmap.py

import sys
import numpy as np
import nibabel as nib

import tractconverter as tc

from dipy.tracking.streamlinespeed import length, set_number_of_points

from dipy.tracking.vox2track import track_counts


# nifti, nifti, tck/trk, txt, INT
def main(roi_anat, outfile, tractfile, recofile, mode):

	tracts_format = tc.detect_format(tractfile)
	hdr = tc.formats.header.get_header_from_anat(roi_anat)
	tracts_file = tracts_format(tractfile, anatFile=roi_anat)
	streamlines = np.array([s for s in tracts_file])

	map_img = nib.load(roi_anat)
	voxel_dim = map_img.get_header()['pixdim'][1:4]
	anat_dim = map_img.get_header().get_data_shape()

	# desired_stepsize = 0.5
	# streamlines_resamp = [set_number_of_points(t, int(np.ceil(length(t)/desired_stepsize))) for t in streamlines]
	# streamlines = streamlines_resamp

	recoscore = np.genfromtxt(recofile)

	tcs, tes = track_counts(streamlines, anat_dim, voxel_dim, True)

	newmap = np.zeros(anat_dim)
	mode = int(mode)
	if mode == 0:
		# tractcount
		print('mode == 0')
		newmap = tcs
	elif mode == 1:
		# reco score sum
		print('mode == 1')
		for vox in tes.keys():
			fiberlist = tes[vox]
			newmap[vox] = np.sum(recoscore[fiberlist])
	elif mode == 2:
		# normalized reco score sum
		print('mode == 2')
		for vox in tes.keys():
			fiberlist = tes[vox]
			newmap[vox] = np.mean(recoscore[fiberlist])
	elif mode == 3:
		# max reco score
		print('mode == 3')
		for vox in tes.keys():
			fiberlist = tes[vox]
			newmap[vox] = np.max(recoscore[fiberlist])


	aff = map_img.get_affine()
	img = nib.nifti1.Nifti1Image(newmap.astype(np.float32), aff)
	nib.save(img, outfile)


if __name__ == "__main__":
	main(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4], sys.argv[5])