def __call__(self, target_ksp, target_im, attrs, fname, slice):
kspace_np = target_ksp
target_im = transforms.to_tensor(target_im)
target_ksp = transforms.to_tensor(target_ksp)
if self.args.coil_compress_coils:
target_ksp = transforms.coil_compress(target_ksp, self.args.coil_compress_coils)
if self.args.calculate_offsets_directly:
krow = kspace_np.sum(axis=(0,1)) # flatten to a single row
width = len(krow)
offset = (krow != 0).argmax()
acq_start = offset
acq_end = width - (krow[::-1] != 0).argmax() #exclusive
else:
offset = None # Mask will pick randomly
if self.partition == 'val' and 'mask_offset' in attrs:
offset = attrs['mask_offset']
acq_start = attrs['padding_left']
acq_end = attrs['padding_right']
#pdb.set_trace()
seed = None if not self.use_seed else tuple(map(ord, fname))
input_ksp, mask, num_lf = transforms.apply_mask(
target_ksp, self.mask_func,
seed, offset,
(acq_start, acq_end))
#pdb.set_trace()
sens_map = torch.Tensor(0)
if self.args.compute_sensitivities:
start_of_center_mask = (kspace_np.shape[-1] - num_lf + 1) // 2
end_of_center_mask = start_of_center_mask + num_lf
sens_map = est_sens_maps(kspace_np, start_of_center_mask, end_of_center_mask)
sens_map = transforms.to_tensor(sens_map)
if self.args.grappa_input:
with h5py.File(self.args.grappa_input_path / self.partition / fname, 'r') as hf:
kernel = transforms.to_tensor(hf['kernel'][slice])
input_ksp = transforms.apply_grappa(input_ksp, kernel, target_ksp, mask)
grappa_kernel = torch.Tensor(0)
if self.args.grappa_path is not None:
with h5py.File(self.args.grappa_path / self.partition / fname, 'r') as hf:
grappa_kernel = transforms.to_tensor(hf['kernel'][slice])
if self.args.grappa_target:
with h5py.File(self.args.grappa_target_path / self.partition / fname, 'r') as hf:
kernel = transforms.to_tensor(hf['kernel'][slice])
target_ksp = transforms.apply_grappa(target_ksp.clone(), kernel, target_ksp, mask, sample_accel=2)
target_im = transforms.root_sum_of_squares(transforms.complex_abs(transforms.ifft2(target_ksp)))
input_im = transforms.ifft2(input_ksp)
if not self.args.scale_inputs:
scale = torch.Tensor([1.])
else:
abs_input = transforms.complex_abs(input_im)
if self.args.scale_type == 'max':
scale = torch.max(abs_input)
else:
scale = torch.mean(abs_input)
input_ksp /= scale
target_ksp /= scale
target_im /= scale
scale = scale.view([1, 1, 1])
attrs_dict = dict(**attrs)
return OrderedDict(
input = input_ksp,
target = target_ksp,
target_im = target_im,
mask = mask,
grappa_kernel = grappa_kernel,
scale = scale,
attrs_dict = attrs_dict,
fname = fname,
slice = slice,
num_lf = num_lf,
sens_map = sens_map,
)
def forward(self, x, input):
return T.apply_grappa(x, input[f'grappa_{self.acceleration}'], input['kspace'], input['mask'].float(), sample_accel=self.acceleration)