Commit 7e359259 authored by Andrei-Claudiu Roibu's avatar Andrei-Claudiu Roibu 🖥
Browse files

fixed bug: mask now selecting right channel + being passed to cpu

parent 04272673
......@@ -355,7 +355,7 @@ def _generate_volume_map(file_path, model, device, cuda_available, brain_mask_pa
MNI152_T1_2mm_brain_mask = torch.from_numpy(Image(brain_mask_path).data)
if mean_reduction == True:
mean_mask = torch.from_numpy(Image(mean_mask_path).data)
mean_mask = torch.from_numpy(Image(mean_mask_path).data[:, :, :, 0])
if cuda_available and (type(device) == int):
volume = volume.cuda(device)
......@@ -368,7 +368,7 @@ def _generate_volume_map(file_path, model, device, cuda_available, brain_mask_pa
if mean_reduction==True:
predicted_complete_volume = torch.add(output, mean_mask)
predicted_complete_volume = (complete_output.cpu().numpy()).astype('float32')
predicted_complete_volume = (predicted_complete_volume.cpu().numpy()).astype('float32')
predicted_complete_volume = np.squeeze(predicted_complete_volume)
else:
predicted_complete_volume = None
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment