120 lines
4.8 KiB
Python
120 lines
4.8 KiB
Python
"""
|
|
Copyright (C) 2018-2020 Intel Corporation
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
#! [mri_demo:demo]
|
|
import numpy as np
|
|
import cv2 as cv
|
|
import argparse
|
|
import time
|
|
from openvino.inference_engine import IECore
|
|
|
|
|
|
def kspace_to_image(kspace):
|
|
assert(len(kspace.shape) == 3 and kspace.shape[-1] == 2)
|
|
fft = cv.idft(kspace, flags=cv.DFT_SCALE)
|
|
img = cv.magnitude(fft[:,:,0], fft[:,:,1])
|
|
return cv.normalize(img, dst=None, alpha=255, beta=0, norm_type=cv.NORM_MINMAX, dtype=cv.CV_8U)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description='MRI reconstrution demo for network from https://github.com/rmsouza01/Hybrid-CS-Model-MRI (https://arxiv.org/abs/1810.12473)')
|
|
parser.add_argument('-i', '--input', dest='input', help='Path to input .npy file with MRI scan data.')
|
|
parser.add_argument('-p', '--pattern', dest='pattern', help='Path to sampling mask in .npy format.')
|
|
parser.add_argument('-m', '--model', dest='model', help='Path to .xml file of OpenVINO IR.')
|
|
parser.add_argument('-l', '--cpu_extension', dest='cpu_extension', help='Path to extensions library with FFT implementation.')
|
|
parser.add_argument('-d', '--device', dest='device', default='CPU',
|
|
help='Optional. Specify the target device to infer on; CPU, '
|
|
'GPU, HDDL or MYRIAD is acceptable. For non-CPU targets, '
|
|
'HETERO plugin is used with CPU fallbacks to FFT implementation. '
|
|
'Default value is CPU')
|
|
args = parser.parse_args()
|
|
|
|
xml_path = args.model
|
|
assert(xml_path.endswith('.xml'))
|
|
bin_path = xml_path[:xml_path.rfind('.xml')] + '.bin'
|
|
|
|
ie = IECore()
|
|
ie.add_extension(args.cpu_extension, "CPU")
|
|
|
|
net = ie.read_network(xml_path, bin_path)
|
|
|
|
device = 'CPU' if args.device == 'CPU' else ('HETERO:' + args.device + ',CPU')
|
|
exec_net = ie.load_network(net, device)
|
|
|
|
# Hybrid-CS-Model-MRI/Data/stats_fs_unet_norm_20.npy
|
|
stats = np.array([2.20295299e-01, 1.11048916e+03, 4.16997984e+00, 4.71741395e+00], dtype=np.float32)
|
|
# Hybrid-CS-Model-MRI/Data/sampling_mask_20perc.npy
|
|
var_sampling_mask = np.load(args.pattern) # TODO: can we generate it in runtime?
|
|
print('Sampling ratio:', 1.0 - var_sampling_mask.sum() / var_sampling_mask.size)
|
|
|
|
data = np.load(args.input)
|
|
num_slices, height, width = data.shape[0], data.shape[1], data.shape[2]
|
|
pred = np.zeros((num_slices, height, width), dtype=np.uint8)
|
|
data /= np.sqrt(height * width)
|
|
|
|
print('Compute...')
|
|
start = time.time()
|
|
for slice_id, kspace in enumerate(data):
|
|
kspace = kspace.copy()
|
|
|
|
# Apply sampling
|
|
kspace[var_sampling_mask] = 0
|
|
kspace = (kspace - stats[0]) / stats[1]
|
|
|
|
# Forward through network
|
|
input = np.expand_dims(kspace.transpose(2, 0, 1), axis=0)
|
|
outputs = exec_net.infer(inputs={'input_1': input})
|
|
output = next(iter(outputs.values()))
|
|
output = output.reshape(height, width)
|
|
|
|
# Save predictions
|
|
pred[slice_id] = cv.normalize(output, dst=None, alpha=255, beta=0, norm_type=cv.NORM_MINMAX, dtype=cv.CV_8U)
|
|
|
|
print('Elapsed time: %.1f seconds' % (time.time() - start))
|
|
|
|
WIN_NAME = 'MRI reconstruction with OpenVINO'
|
|
|
|
slice_id = 0
|
|
def callback(pos):
|
|
global slice_id
|
|
slice_id = pos
|
|
|
|
kspace = data[slice_id]
|
|
img = kspace_to_image(kspace)
|
|
|
|
kspace[var_sampling_mask] = 0
|
|
masked = kspace_to_image(kspace)
|
|
|
|
rec = pred[slice_id]
|
|
|
|
# Add a header
|
|
border_size = 20
|
|
render = cv.hconcat((img, masked, rec))
|
|
render = cv.copyMakeBorder(render, border_size, 0, 0, 0, cv.BORDER_CONSTANT, value=255)
|
|
cv.putText(render, 'Original', (0, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, color=0)
|
|
cv.putText(render, 'Sampled (PSNR %.1f)' % cv.PSNR(img, masked), (width, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, color=0)
|
|
cv.putText(render, 'Reconstructed (PSNR %.1f)' % cv.PSNR(img, rec), (width*2, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, color=0)
|
|
|
|
cv.imshow(WIN_NAME, render)
|
|
cv.waitKey(1)
|
|
|
|
cv.namedWindow(WIN_NAME, cv.WINDOW_NORMAL)
|
|
print(num_slices)
|
|
cv.createTrackbar('Slice', WIN_NAME, num_slices // 2, num_slices - 1, callback)
|
|
callback(num_slices // 2) # Trigger initial visualization
|
|
cv.waitKey()
|
|
#! [mri_demo:demo]
|