-
Notifications
You must be signed in to change notification settings - Fork 7
/
dipy_nlmeans.py
executable file
·92 lines (63 loc) · 2.7 KB
/
dipy_nlmeans.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
#! /usr/bin/env python2.7
# Wrapper for the nlmeans python wrapper function
from __future__ import print_function
import nibabel as nib
import numpy as np
import os
import argparse
from dipy.denoise.non_local_means import non_local_means as nlmeans
from dipy.denoise.noise_estimate import estimate_sigma, piesno
DESCRIPTION = """
Loads an image and launch dipy NLMEANS
"""
def buildArgsParser():
p = argparse.ArgumentParser(description=DESCRIPTION)
p.add_argument('input', action='store', metavar='DIR',
help='path of the image file to denoise')
p.add_argument('-std', action='store', dest='std',
metavar='std', required=False, default=0, type=int,
help='standard deviation of the noise')
p.add_argument('-mask', action='store', dest='mask',
metavar='mask', required=False, default=None, type=str,
help='Mask to accelerate computation.')
p.add_argument('-o', action='store', dest='savename',
metavar='savename', required=False, default=None, type=str,
help='path and prefix for the saved denoised file. The name is always appended with _denoised.nii.gz')
p.add_argument('-avg', action='store', dest='avg',
metavar='averaging', required=False, default=1, type=int,
help='Set to 0 to use Standard weighting. Set to 1 to use Rician weighting')
return p
def main():
parser = buildArgsParser()
args = parser.parse_args()
if args.savename is None:
temp, ext = str.split(os.path.basename(args.input), '.', 1)
filename = os.path.dirname(os.path.realpath(args.input)) + '/' + temp
else:
filename = args.savename
# If the file already exists, check if the user want to overwrite it. If not, simply exit.
filename += '_denoised.nii.gz'
print ("Now denoising", os.path.realpath(args.input))
vol = nib.load(args.input)
affine = vol.get_affine()
data = vol.get_data()
if args.std is 0:
print ("Compute sigma with piesno")
sigma = piesno(data, N=1, return_mask=False)
print ("sigma is: ", sigma)
sigma = np.mean(sigma)
print ("sigma is: ", sigma)
else:
print ("Sigma is: ", args.std)
sigma = args.std
if args.mask is None:
print ("No mask.. you should")
else:
print ("Mask is: ", args.mask)
img_mask = nib.load(args.mask)
mask_t1 = img_mask.get_data()
img_denoised = nlmeans(data, sigma, mask=mask_t1)
nib.save(nib.Nifti1Image(img_denoised, affine), filename)
print ("Denoised file was saved as", filename)
if __name__ == "__main__":
main()