harmonization.dictionary

  1import os
  2import warnings
  3
  4import numpy as np
  5import nibabel as nib
  6
  7from itertools import cycle
  8from ast import literal_eval
  9from time import time
 10
 11from scipy.ndimage import zoom
 12
 13from autodmri.blocks import extract_patches
 14from nlsam.angular_tools import angular_neighbors
 15from nlsam.denoiser import greedy_set_finder
 16
 17from harmonization.solver import online_DL, solve_l1
 18from harmonization.recon import depimp_zoom, reconstruct_from_blocks
 19
 20
 21def get_global_D(datasets, outfilename, block_size, block_up, ncores=None, batchsize=32, niter=500,
 22                 use_std=False, positivity=False, fit_intercept=True, center=True, nlambdas=100,
 23                 b0_threshold=20, split_b0s=True, **kwargs):
 24
 25    # get the data shape so we can preallocate some arrays
 26    # we also have to assume all datasets have the same 3D shape obviously
 27    shape = nib.load(datasets[0]['data']).header.get_data_shape()
 28
 29    # block size depends on the upsampling factor, so we pick the biggest one
 30    upsample = np.prod(literal_eval(block_up)) > np.prod(literal_eval(block_size))
 31
 32    if upsample:
 33        current_block_size = literal_eval(block_up)
 34    else:
 35        current_block_size = literal_eval(block_size)
 36
 37    n_atoms = int(np.prod(current_block_size) * 2)
 38    b0_block_size = tuple(current_block_size[:-1]) + ((current_block_size[-1] + 1,))
 39    overlap = b0_block_size
 40    to_denoise = np.empty(shape[:-1] + (current_block_size[-1] + 1,), dtype=np.float32)
 41
 42    train_list = []
 43    variance_large = []
 44
 45    for filename in datasets:
 46
 47        print(f"Now feeding dataset {filename['data']}")
 48
 49        mask = nib.load(filename['mask']).get_fdata(caching='unchanged').astype(bool)
 50        data = nib.load(filename['data']).get_fdata(caching='unchanged', dtype=np.float32)
 51        bvals = np.loadtxt(filename['bval'])
 52        bvecs = np.loadtxt(filename['bvec'])
 53
 54        if np.shape(bvecs)[0] == 3:
 55            bvecs = bvecs.T
 56
 57        if mask.shape != data.shape[:-1]:
 58            error = f'mask shape {mask.shape} does not fit with the data shape {data.shape}'
 59            raise ValueError(error)
 60
 61        data *= mask[..., None]
 62        b0_loc = np.where(bvals <= b0_threshold)[0]
 63        dwis = np.where(bvals > b0_threshold)[0]
 64        num_b0s = len(b0_loc)
 65
 66        # We also convert bvecs associated with b0s to exactly (0,0,0), which
 67        # is not always the case when we hack around with the scanner.
 68        bvecs = np.where(bvals[:, None] <= b0_threshold, 0, bvecs)
 69
 70        # Average all b0s if we don't split them in the training set
 71        if num_b0s > 1 and not split_b0s:
 72            num_b0s = 1
 73            data[..., b0_loc] = np.mean(data[..., b0_loc], axis=-1, keepdims=True)
 74
 75        # Split the b0s in a cyclic fashion along the training data
 76        # If we only had one, cycle just return b0_loc indefinitely,
 77        # else we go through all indexes.
 78        np.random.shuffle(b0_loc)
 79        split_b0s_idx = cycle(b0_loc)
 80        sym_bvecs = np.vstack((bvecs, -bvecs))
 81
 82        neighbors = angular_neighbors(sym_bvecs, current_block_size[-1] - 1) % data.shape[-1]
 83        neighbors = neighbors[:data.shape[-1]]  # everything was doubled for symmetry
 84
 85        full_indexes = [(dwi,) + tuple(neighbors[dwi]) for dwi in range(data.shape[-1]) if dwi in dwis]
 86        indexes = greedy_set_finder(full_indexes)
 87
 88        # If we have more b0s than indexes, then we have to add a few more blocks since
 89        # we won't do a full cycle. If we have more b0s than indexes after that, then it breaks.
 90        if num_b0s > len(indexes):
 91            the_rest = [rest for rest in full_indexes if rest not in indexes]
 92            indexes += the_rest[:(num_b0s - len(indexes))]
 93
 94        if num_b0s > len(indexes):
 95            error = f'Seems like you still have more b0s {num_b0s} than available blocks {len(indexes)}, either average them or deactivate subsampling.'
 96            raise ValueError(error)
 97
 98        # whole global centering
 99        if center:
100            data -= data.mean(axis=-1, keepdims=True)
101
102        for i, idx in enumerate(indexes):
103            b0_loc = tuple((next(split_b0s_idx),))
104
105            # if we mix datasets, then we may need to change the storage array size
106            if to_denoise.shape[:-1] != data.shape[:-1]:
107                del to_denoise
108                to_denoise = np.empty(data.shape[:-1] + (current_block_size[-1] + 1,), dtype=np.float32)
109
110            to_denoise[..., 0] = data[..., b0_loc].squeeze()
111            to_denoise[..., 1:] = data[..., idx]
112
113            patches = extract_patches(to_denoise, b0_block_size, overlap, flatten=False)
114            axis = tuple(range(patches.ndim//2, patches.ndim))
115            mask_patch = np.count_nonzero(patches, axis=axis) > np.prod(b0_block_size) // 2
116            patches = patches[mask_patch].reshape(-1, np.prod(b0_block_size))
117
118            if use_std:
119                try:
120                    variance = nib.load(filename['std']).get_fdata(caching='unchanged', dtype=np.float32)**2 * mask
121                    variance = np.broadcast_to(variance[..., None], data.shape)
122                    variance = extract_patches(variance, b0_block_size, overlap, flatten=False)
123                    axis = tuple(range(variance.ndim//2, variance.ndim))
124                    variance = np.median(variance, axis=axis)[mask_patch].ravel()
125                except IOError:
126                    warnings.warn(f"Volume {filename['std']} not found! Skipping using variance.")
127                    variance = [None]
128            else:
129                variance = [None]
130
131            # check to build with np.r_ the whole list from stringnames instead
132            train_list += [patches]
133            variance_large += list(variance)
134
135        del data, mask, patches, variance
136
137    print('Fed everything in')
138
139    lengths = [l.shape[0] for l in train_list]
140    train_data = np.empty((np.sum(lengths), np.prod(b0_block_size)))
141
142    step = 0
143    for i in range(len(train_list)):
144        length = lengths[i]
145        idx = slice(step, step + length)
146        train_data[idx] = train_list[i].reshape(-1, np.prod(b0_block_size))
147        step += length
148
149    del train_list
150
151    # we have variance as a N elements list - so check one element to see if it's an array
152    if variance_large[0] is not None:
153        variance_large = np.asarray(variance_large).ravel()
154    else:
155        variance_large = None
156
157    D = online_DL(train_data, ncores=ncores, positivity=positivity, fit_intercept=fit_intercept, standardize=True,
158                  nlambdas=nlambdas, niter=niter, batchsize=batchsize, n_atoms=n_atoms, variance=variance_large,
159                  progressbar=True, use_joblib=True)
160
161    return D
162
163
164def rebuild(data, mask, D, block_size, block_up, ncores=-1, nlambdas=100,
165            positivity=False, variance=None, fix_mean=True, fit_intercept=False, use_crossval=False):
166
167    data = data * mask[..., None]
168
169    if variance is not None:
170        variance *= mask
171
172    if len(block_size) == len(block_up):
173        last = block_up[-1]
174    else:
175        last = block_size[-1]
176
177    factor = np.divide(block_up, block_size)
178    new_shape = (int(data.shape[0] * factor[0]),
179                 int(data.shape[1] * factor[1]),
180                 int(data.shape[2] * factor[2]),
181                 last)
182    overlap = (1, 1, 1, last)
183    new_overlap = overlap
184
185    if block_size == block_up:
186        D_depimpe = np.copy(D)
187    else:
188        D_depimpe = depimp_zoom(D, block_size, block_up)
189
190    blocks = extract_patches(data, block_size, overlap, flatten=False).reshape(-1, np.prod(block_size))
191    del data
192
193    # skip empty rows from training since they are probably masked background
194    mask = np.count_nonzero(blocks, axis=1) > np.prod(block_size) // 2
195
196    # get the variance as blocks
197    if variance is not None:
198        variance = extract_patches(variance, block_size[:-1], overlap[:-1], flatten=True)
199        axis = tuple(range(1, len(block_size) + 1))
200        variance = np.median(variance, axis=axis)
201
202        # if we are on an edge, variance can be 0, so truncate those cases as well
203        np.logical_and(mask, variance > 0, out=mask)
204        variance = variance[mask]
205
206    blocks = blocks[mask]
207
208    tt = time()
209    X_small_denoised, alpha, intercept, _ = solve_l1(blocks, D_depimpe, variance=variance, return_all=True, nlambdas=nlambdas, use_joblib=True,
210                                                     positivity=positivity, fit_intercept=True, standardize=True, progressbar=True, verbose=5,
211                                                     ncores=ncores, use_crossval=use_crossval)
212
213    print(f'total time was {round((time() - tt) / 60, 2)} mins')
214
215    # we actually only want alpha and intercept in this step and throw out X if we do upsampling normally
216    # for ease of use reason, we now use directly X for the reconstruction, but we should
217    # 1. multiply back X
218    # 2a. feed X in an empty array with a mask
219    #
220    #   OR
221    #
222    # 2b. reconstruct the array and average things internally according to a mask
223
224    if block_size != block_up:
225        X_small_denoised = np.dot(D, alpha).T + intercept
226
227    X_final = np.zeros(((mask.shape[0],) + block_up), dtype=np.float32)
228    X_final[mask] = X_small_denoised.reshape(-1, *block_up)
229    recon = reconstruct_from_blocks(X_final, new_shape, block_size, block_up[:-1], new_overlap, weights=None)
230
231    del X_final, mask
232    return recon
233
234
235def harmonize_my_data(dataset, kwargs):
236    outpath = kwargs['outpath']
237    path_D = kwargs['outfilename']
238    center = kwargs['center']
239    block_up = literal_eval(kwargs['block_up'])
240    block_size = literal_eval(kwargs['block_size'])
241    positivity = kwargs['positivity']
242    fit_intercept = kwargs['fit_intercept']
243    fix_mean = kwargs['fix_mean']
244    use_crossval = kwargs['use_crossval']
245    split_b0s = kwargs['split_b0s']
246    b0_threshold = kwargs['b0_threshold']
247    ncores = kwargs['ncores']
248    nlambdas = kwargs['nlambdas']
249    ext = kwargs['ext']
250    overwrite = kwargs['overwrite']
251
252    upsample = np.prod(block_up) > np.prod(block_size)
253
254    print(f"Now rebuilding {dataset['data']}")
255    D = np.load(path_D)
256
257    if dataset['folder'] is None:
258        relfolder = ''
259    else:
260        relfolder = dataset['folder']
261
262    new_suffix = '_' + os.path.splitext(os.path.basename(path_D))[0] + ext
263    new_filename = os.path.basename(dataset['data']).replace(ext, new_suffix)
264    output_filename = os.path.join(outpath, relfolder, new_filename)
265
266    # Create subfolders tree if it does not exist
267    if not os.path.exists(os.path.dirname(output_filename)):
268        os.makedirs(os.path.dirname(output_filename), exist_ok=True)
269
270    if os.path.isfile(output_filename):
271        if overwrite:
272            warnings.warn(f'File already exists! Overwriting {output_filename}')
273            pass
274        else:
275            warnings.warn(f'File already exists! Skipping {output_filename}')
276            return
277
278    vol = nib.load(dataset['data'])
279    data = vol.get_fdata(caching='unchanged', dtype=np.float32)
280    affine = vol.affine
281    header = vol.header
282
283    mask = nib.load(dataset['mask']).get_fdata(caching='unchanged').astype(bool)
284    bvals = np.loadtxt(dataset['bval'])
285    bvecs = np.loadtxt(dataset['bvec'])
286
287    if np.shape(bvecs)[0] == 3:
288        bvecs = bvecs.T
289
290    if mask.shape != data.shape[:-1]:
291            error = f'mask shape {mask.shape} does not fit with the data shape {data.shape}'
292            raise ValueError(error)
293
294    if center:
295        # we need to pull down the 3D volume mean for the upsampling part to make sense though
296        # We use nanmean because we implicitly exclude voxels with a value of 0 from the mean this way
297        data_mean = np.nanmean(np.where(data != 0, data, np.nan), axis=(0, 1, 2), keepdims=True)
298        data -= data_mean
299
300    variance = None
301
302    # number of blocks is without b0, so we +1 the last dimension
303    if len(block_up) < data.ndim:
304        current_block_size = block_size[:-1] + (block_size[-1] + 1,)
305        current_block_up = block_up + (block_size[-1] + 1,)
306    else:
307        current_block_size = block_size[:-1] + (block_size[-1] + 1,)
308        current_block_up = block_up[:-1] + (block_up[-1] + 1,)
309
310    print(f'Output filename is {output_filename}')
311
312    factor = np.divide(current_block_up, current_block_size)
313
314    b0_loc = np.where(bvals <= b0_threshold)[0]
315    dwis = np.where(bvals > b0_threshold)[0]
316    num_b0s = len(b0_loc)
317
318    # We also convert bvecs associated with b0s to exactly (0,0,0), which
319    # is not always the case when we hack around with the scanner.
320    bvecs = np.where(bvals[:, None] <= b0_threshold, 0, bvecs)
321
322    # Average all b0s if we don't split them in the training set
323    if num_b0s > 1 and not split_b0s:
324        num_b0s = 1
325        data[..., b0_loc] = np.mean(data[..., b0_loc], axis=-1, keepdims=True)
326
327    # Split the b0s in a cyclic fashion along the training data
328    # If we only had one, cycle just return b0_loc indefinitely,
329    # else we go through all indexes.
330    np.random.shuffle(b0_loc)
331    split_b0s_idx = cycle(b0_loc)
332    sym_bvecs = np.vstack((bvecs, -bvecs))
333
334    neighbors = angular_neighbors(sym_bvecs, current_block_size[-1] - 2) % data.shape[-1]
335    neighbors = neighbors[:data.shape[-1]]  # everything was doubled for symmetry
336
337    full_indexes = [(dwi,) + tuple(neighbors[dwi]) for dwi in range(data.shape[-1]) if dwi in dwis]
338    indexes = greedy_set_finder(full_indexes)
339
340    # If we have more b0s than indexes, then we have to add a few more blocks since
341    # we won't do a full cycle. If we have more b0s than indexes after that, then it breaks.
342    if num_b0s > len(indexes):
343        the_rest = [rest for rest in full_indexes if rest not in indexes]
344        indexes += the_rest[:(num_b0s - len(indexes))]
345
346    if num_b0s > len(indexes):
347        error = f'Seems like you still have more b0s {num_b0s} than available blocks {len(indexes)}, either average them or deactivate subsampling.'
348        raise ValueError(error)
349
350    # Stuff happens here / we only pimp up dwis, not b0s
351    # actually we can't really pimp 4D stuff, SH are there for that
352    predicted_size = (int(data.shape[0] * factor[0]),
353                      int(data.shape[1] * factor[1]),
354                      int(data.shape[2] * factor[2]),
355                      data.shape[-1])
356
357    predicted = np.zeros(predicted_size, dtype=np.float32)
358    divider = np.zeros(predicted.shape[-1])
359
360    # Put all idx + b0 in this array in each iteration
361    to_denoise = np.empty(data.shape[:-1] + (current_block_size[-1],), dtype=np.float64)
362
363    for i, idx in enumerate(indexes, start=1):
364        b0_loc = tuple((next(split_b0s_idx),))
365        to_denoise[..., 0] = data[..., b0_loc].squeeze()
366        to_denoise[..., 1:] = data[..., idx]
367        divider[list(b0_loc + idx)] += 1
368
369        print(f'Now rebuilding volumes {np.hstack((b0_loc, idx))} / block {i} out of {len(indexes)}.')
370
371        predicted[..., b0_loc + idx] += rebuild(to_denoise,
372                                                mask,
373                                                D,
374                                                block_size=current_block_size,
375                                                block_up=current_block_up,
376                                                ncores=ncores,
377                                                positivity=positivity,
378                                                fix_mean=fix_mean,
379                                                nlambdas=nlambdas,
380                                                # center=center,
381                                                fit_intercept=fit_intercept,
382                                                use_crossval=use_crossval,
383                                                variance=variance)
384    predicted /= divider
385
386    # If we upsample, the mask does not match anymore, so we also resample it
387    if upsample:
388        factors = 1 / np.divide(mask.shape, predicted.shape[:-1])
389        mask = zoom(mask, factors, order=0, mode='grid-constant', grid_mode=True)
390        mask_output_filename = output_filename.replace(ext, '_mask' + ext)
391        nib.Nifti1Image(mask.astype(np.int8), affine, header).to_filename(mask_output_filename)
392
393    if center:
394        np.add(predicted, data_mean, where=mask[..., None], out=predicted)
395
396    # clip negatives, which could happens at the borders
397    predicted[predicted < 0] = 0
398
399    nib.Nifti1Image(predicted, affine, header).to_filename(output_filename)