def RotateImage(img, deltaPhi):
mt = img.memType
img.MoveToGPU()
img.ReIm2AmPh()
piBy4 = np.pi / 4
deltaPhiRad = deltaPhi * np.pi / 180.0
rMax = np.sqrt((img.width / 2.0) ** 2 + (img.height / 2.0) ** 2)
if (deltaPhiRad // piBy4 + 1) % 2:
angle = piBy4 - (deltaPhiRad % piBy4)
else:
angle = deltaPhiRad % piBy4
rotWidth = int(np.ceil(2 * rMax * np.cos(angle)))
rotHeight = rotWidth
imgRotated = Image(rotHeight, rotWidth, img.cmpRepr, img.memType)
filled = cuda.to_device(np.zeros(imgRotated.amPh.am.shape, dtype=np.int32))
# rotation
blockDim, gridDim = ccfg.DetermineCudaConfigNew(img.amPh.am.shape)
RotateImage_dev[gridDim, blockDim](img.amPh.am, imgRotated.amPh.am, filled, deltaPhiRad)
imgRotated.MoveToGPU()
blockDim, gridDim = ccfg.DetermineCudaConfigNew(imgRotated.amPh.am.shape)
InterpolateMissingPixels_dev[gridDim, blockDim](imgRotated.amPh.am, filled)
imgRotated = CreateImageExpFromImage(imgRotated)
img.ChangeMemoryType(mt)
return imgRotated
def MultAmPhMatrices(ap1, ap2):
blockDim, gridDim = ccfg.DetermineCudaConfig(ap1.am.shape[0])
apRes = ComplexAmPhMatrix(ap1.am.shape[0], ap1.am.shape[1],
ComplexAmPhMatrix.mem['GPU'])
MultAmPhMatrices_dev[gridDim, blockDim](ap1.am, ap1.ph, ap2.am, ap2.ph,
apRes.am, apRes.ph)
return apRes
def ConjugateAmPhMatrix(ap):
blockDim, gridDim = ccfg.DetermineCudaConfig(ap.am.shape[0])
apConj = ComplexAmPhMatrix(ap.am.shape[0], ap.am.shape[1],
ComplexAmPhMatrix.mem['GPU'])
ConjugateAmPhMatrix_dev[gridDim, blockDim](ap.am, ap.ph, apConj.am,
apConj.ph)
return apConj
def InsertAperture(img, ap_diameter):
img_dim = img.amPh.am.shape[0]
blockDim, gridDim = ccfg.DetermineCudaConfig(img_dim)
img.MoveToGPU()
img_with_aperature = imsup.CopyImage(img)
ap_radius = ap_diameter // 2
InsertAperture_dev[gridDim, blockDim](img_with_aperature.amPh.am, img_with_aperature.amPh.ph, img_dim, ap_radius)
return img_with_aperature
def MagnifyImage(img, factor):
img.MoveToGPU()
img.ReIm2AmPh()
magHeight = int(factor * img.height)
magWidth = int(factor * img.width)
imgScaled = Image(magHeight, magWidth, img.cmpRepr, img.memType)
filled = cuda.to_device(np.zeros(imgScaled.amPh.am.shape, dtype=np.int32))
# magnification
blockDim, gridDim = ccfg.DetermineCudaConfigNew(img.amPh.am.shape)
MagnifyImage_dev[gridDim, blockDim](img.amPh.am, imgScaled.amPh.am, filled, factor)
imgScaled.MoveToGPU()
blockDim, gridDim = ccfg.DetermineCudaConfigNew(imgScaled.amPh.am.shape)
InterpolateMissingPixels_dev[gridDim, blockDim](imgScaled.amPh.am, filled)
return imgScaled
def CalcTransferFunction(imgDim, pxDim, defocusChange):
blockDim, gridDim = ccfg.DetermineCudaConfig(imgDim)
ctf = imsup.Image(imgDim, imgDim, imsup.Image.cmp['CAP'],
imsup.Image.mem['GPU'])
ctfCoeff = np.pi * const.ewfLambda * defocusChange
CalcTransferFunction_dev[gridDim, blockDim](ctf.amPh.am, ctf.amPh.ph,
imgDim, pxDim, ctfCoeff)
ctf.defocus = defocusChange
return ctf
def PasteROIToImage(img, roi, roiOrig):
dt = img.cmpRepr
img.AmPh2ReIm()
imgNew = Image(img.height, img.width, Image.cmp['CRI'], Image.mem['GPU'])
imgNew.reIm = img.reIm
roiOrig_d = cuda.to_device(np.array(roiOrig))
blockDim, gridDim = ccfg.DetermineCudaConfigNew(roi.reIm.shape)
PasteROIToImage_dev[gridDim, blockDim](imgNew.reIm, roi.reIm, roiOrig_d)
img.ChangeComplexRepr(dt)
return imgNew
def AmPh2ReIm(self):
if self.cmpRepr == self.cmp['CRI']:
return
mt = self.memType
self.MoveToGPU()
blockDim, gridDim = ccfg.DetermineCudaConfigNew((self.height, self.width))
AmPh2ReIm_dev[gridDim, blockDim](self.amPh.am, self.amPh.ph, self.reIm)
self.cmpRepr = self.cmp['CRI']
if mt == self.mem['CPU']:
self.MoveToCPU()
def ShiftImageAmpBuffer(img, shift):
img.shift = [x + dx for x, dx in zip(img.shift, shift)]
fillValue = np.max(img.amPh.am)
img.MoveToGPU()
imgShifted = imsup.Image(img.height, img.width, imsup.Image.cmp['CAP'],
imsup.Image.mem['GPU'])
shift_d = cuda.to_device(np.array(img.shift))
blockDim, gridDim = ccfg.DetermineCudaConfigNew(img.buffer.shape)
ShiftImage_dev[gridDim, blockDim](img.amPh.am, imgShifted.amPh.am, shift_d,
fillValue)
img.buffer.copy_to_device(imgShifted.amPh.am)
img.MoveToCPU()
def FFT2Diff(fft):
mt = fft.memType
dt = fft.cmpRepr
fft.MoveToGPU()
fft.AmPh2ReIm()
diff = imsup.Image(fft.height, fft.width, imsup.Image.cmp['CRI'],
imsup.Image.mem['GPU'])
blockDim, gridDim = ccfg.DetermineCudaConfig(fft.width)
FFT2Diff_dev[gridDim, blockDim](fft.reIm, diff.reIm, fft.width)
diff.defocus = fft.defocus
fft.ChangeComplexRepr(dt)
fft.ChangeMemoryType(mt)
return diff
def calc_dists_from_ellipse(data, ellipse):
dists_d = cuda.to_device(np.zeros(data.shape, dtype=np.float32))
# zero_fix można uwzględnić w F1_d, F2_d zamiast w kernelu
F1_d = cuda.to_device(np.array(ellipse.F1).astype(np.float32))
F2_d = cuda.to_device(np.array(ellipse.F2).astype(np.float32))
blockDim, gridDim = ccfg.DetermineCudaConfigNew(data.shape)
calc_dists_from_ellipse_dev[gridDim, blockDim](dists_d, F1_d, F2_d)
dists = dists_d.copy_to_host()
dists = np.abs(dists - ellipse.s)
return dists
def CropImageROICoords(img, coords):
roiHeight = coords[3] - coords[1]
roiWidth = coords[2] - coords[0]
dt = img.cmpRepr
img.AmPh2ReIm()
roi = Image(roiHeight, roiWidth, img.cmpRepr, Image.mem['GPU'])
topLeft_d = cuda.to_device(np.array(coords[:2], dtype=np.int32))
blockDim, gridDim = ccfg.DetermineCudaConfigNew((roiHeight, roiWidth))
CropImageROICoords_dev[gridDim, blockDim](img.reIm, roi.reIm, topLeft_d)
img.ChangeComplexRepr(dt)
roi.ChangeComplexRepr(dt)
roi.defocus = img.defocus # !!!
roi.numInSeries = img.numInSeries # !!!
return roi
def CropImageROI(img, roiOrig, roiDims, isOrigTopLeft):
dt = img.cmpRepr
img.AmPh2ReIm()
roi = Image(roiDims[0], roiDims[1], img.cmpRepr, Image.mem['GPU'])
roiOrig_d = cuda.to_device(np.array(roiOrig))
blockDim, gridDim = ccfg.DetermineCudaConfigNew(roiDims)
if isOrigTopLeft:
CropImageROITopLeft_dev[gridDim, blockDim](img.reIm, roi.reIm, roiOrig_d)
else:
CropImageROIMid_dev[gridDim, blockDim](img.reIm, roi.reIm, roiOrig_d)
img.ChangeComplexRepr(dt)
roi.ChangeComplexRepr(dt)
roi.defocus = img.defocus # !!!
return roi
def crop_am_ph_roi(img, coords):
mt = img.memType
img.MoveToGPU()
roi_h = coords[3] - coords[1]
roi_w = coords[2] - coords[0]
roi = ImageExp(roi_h, roi_w, img.cmpRepr, img.memType)
top_left_d = cuda.to_device(np.array(coords[:2], dtype=np.int32))
block_dim, grid_dim = ccfg.DetermineCudaConfigNew((roi_h, roi_w))
CropImageROICoords_dev[grid_dim, block_dim](img.amPh.am, roi.amPh.am, top_left_d)
CropImageROICoords_dev[grid_dim, block_dim](img.amPh.ph, roi.amPh.ph, top_left_d)
img.ChangeMemoryType(mt)
roi.ChangeMemoryType(mt)
return roi
def ShiftImage(img, shift):
mt = img.memType
dt = img.cmpRepr
img.MoveToGPU()
img.AmPh2ReIm()
imgShifted = imsup.Image(img.height, img.width, img.cmpRepr, img.memType)
shift_d = cuda.to_device(np.array(shift))
blockDim, gridDim = ccfg.DetermineCudaConfigNew(img.reIm.shape)
ShiftImage_dev[gridDim, blockDim](img.reIm, imgShifted.reIm, shift_d, 0.0)
img.ChangeComplexRepr(dt)
img.ChangeMemoryType(mt)
imgShifted.ChangeComplexRepr(dt)
imgShifted.ChangeMemoryType(mt)
return imgShifted
def PasteROIToImageTopLeft(img, roi, roiTopLeft, spot=False):
img.MoveToGPU()
roi.MoveToGPU()
dt = img.cmpRepr
imgNew = CopyImage(img)
imgNew.AmPh2ReIm()
roi.AmPh2ReIm()
roiOrig = np.array(roiTopLeft) + np.array(roi.reIm.shape) // 2 # !!!
roiOrig_d = cuda.to_device(roiOrig)
blockDim, gridDim = ccfg.DetermineCudaConfigNew(roi.reIm.shape)
if not spot:
PasteROIToImage_dev[gridDim, blockDim](imgNew.reIm, roi.reIm, roiOrig_d)
else:
PasteSpotToImage_dev[gridDim, blockDim](imgNew.reIm, roi.reIm, roiOrig_d)
imgNew.ChangeComplexRepr(dt)
return imgNew
def calc_ctf(img_dim, px_dim, defocus, Cs=const.Cs, A1=ab.PolarComplex(const.A1_amp, const.A1_phs),
df_spread=const.df_spread, conv_angle=const.conv_angle, A1_dir=1):
df_coeff = np.pi * const.ewfLambda * defocus
Cs_coeff = 0.5 * np.pi * (const.ewfLambda ** 3) * Cs
Cs_spat_coeff = Cs * const.ewfLambda ** 2
conv_ang_coeff = (np.pi * conv_angle) ** 2
df_spread_coeff = 0.5 * np.pi * const.ewfLambda * df_spread
A1_re_coeff = A1_dir * np.pi * const.ewfLambda * A1.real()
A1_im_coeff = A1_dir * 2 * np.pi * const.ewfLambda * A1.imag()
block_dim, grid_dim = ccfg.DetermineCudaConfig(img_dim)
ctf = imsup.ImageExp(img_dim, img_dim, imsup.Image.cmp['CAP'], imsup.Image.mem['GPU'])
calc_ctf_dev[grid_dim, block_dim](ctf.amPh.am, ctf.amPh.ph, img_dim, px_dim, defocus, df_coeff, Cs_coeff, A1_re_coeff,
A1_im_coeff, Cs_spat_coeff, conv_ang_coeff, df_spread_coeff)
ctf.defocus = defocus
ctf.MoveToCPU()
return ctf
def FindMinInImage(img):
dimSize = img.width
arrReduced = img.amPh.am
blockDim, gridDim = ccfg.DetermineCudaConfig(dimSize // 2)
while dimSize > 2:
dimSize //= 2
# arrReducedNew = cuda.device_array((dimSize, dimSize), dtype=np.float32)
arrReducedNew = cuda.to_device(
np.zeros((dimSize, dimSize), dtype=np.float32))
ReduceArrayToFindMin_dev[gridDim, blockDim](arrReduced, arrReducedNew)
arrReduced = arrReducedNew
if gridDim[0] > 1:
gridDim = [gridDim[0] // 2] * 2
else:
blockDim = [blockDim[0] // 2] * 2
imgMin = np.min(arrReduced.copy_to_host())
# imgMin = arrReduced.copy_to_host()[0, 0]
return imgMin
def shift_am_ph_image(img, shift):
mt = img.memType
img.MoveToGPU()
img_shifted = imsup.ImageWithBuffer(img.height, img.width, img.cmpRepr,
img.memType)
shift_d = cuda.to_device(np.array(shift))
blockDim, gridDim = ccfg.DetermineCudaConfigNew(img.amPh.am.shape)
ShiftImage_dev[gridDim, blockDim](img.amPh.am, img_shifted.amPh.am,
shift_d, 0.0)
ShiftImage_dev[gridDim, blockDim](img.amPh.ph, img_shifted.amPh.ph,
shift_d, 0.0)
img.ChangeMemoryType(mt)
img_shifted.ChangeMemoryType(mt)
if img.cos_phase is not None:
img_shifted.update_cos_phase()
return img_shifted
def ShiftImage(img, shift):
dt = img.cmpRepr
img.AmPh2ReIm()
imgShifted = imsup.Image(img.height, img.width, img.cmpRepr,
imsup.Image.mem['GPU'])
shift_d = cuda.to_device(np.array(shift))
blockDim, gridDim = ccfg.DetermineCudaConfigNew(img.reIm.shape)
ShiftImage_dev[gridDim, blockDim](img.reIm, imgShifted.reIm, shift_d, 0.0)
img.ChangeComplexRepr(dt)
imgShifted.ChangeComplexRepr(dt)
# imgShifted.prev = img.prev
# imgShifted.next = img.next
# if imgShifted.prev is not None:
# imgShifted.prev.next = imgShifted
# if imgShifted.next is not None:
# imgShifted.next.prev = imgShifted
# print('Image was shifted by ({0}, {1}) px'.format(shift[1], shift[0]))
return imgShifted
def run_iwfr(imgs_to_iwfr, n_iters):
ewf_dir = 'results/ewf/'
amp_name = 'amp'
phs_name = 'phs'
amp_path_base = '{0}{1}_00.png'.format(ewf_dir, amp_name)
phs_path_base = '{0}{1}_00.png'.format(ewf_dir, phs_name)
print('Starting IWFR...')
exit_wave = imsup.copy_am_ph_image(imgs_to_iwfr[0])
for i in range(0, n_iters):
print('Iteration no {0}...'.format(i+1))
imgs_to_iwfr, exit_wave = run_iteration_of_iwfr(imgs_to_iwfr)
ewf_amp_path = amp_path_base.replace('00', '0{0}'.format(i+1) if i < 10 else '{0}'.format(i+1))
ewf_phs_path = phs_path_base.replace('00', '0{0}'.format(i+1) if i < 10 else '{0}'.format(i+1))
imsup.SaveAmpImage(exit_wave, ewf_amp_path)
imsup.SavePhaseImage(exit_wave, ewf_phs_path)
ccfg.GetGPUMemoryUsed()
print('All done')
return exit_wave
from numba import cuda
import Constants as const
import CudaConfig as ccfg
import GUI as gui
cuda.select_device(0)
dev = cuda.get_current_device()
print('CUDA device in use: ' + dev.name.decode())
ccfg.GetGPUFreeMemory()
ccfg.GetGPUMemoryUsed()
gui.RunEwrWindow(const.gridDim)
def AddTwoArrays(arr1, arr2):
blockDim, gridDim = ccfg.DetermineCudaConfig(arr1.shape[0])
# arrSum = cuda.device_array(arr1.shape, dtype=arr1.dtype)
arrSum = cuda.to_device(np.zeros(arr1.shape, dtype=arr1.dtype))
AddTwoArrays_dev[gridDim, blockDim](arr1, arr2, arrSum)
return arrSum
def CalcSqrtOfArray(arr):
blockDim, gridDim = ccfg.DetermineCudaConfig(arr.shape[0])
arrSqrt = cuda.device_array(arr.shape, dtype=arr.dtype)
CalcSqrtOfArray_dev[gridDim, blockDim](arr, arrSqrt)
return arrSqrt
def MultArrayByScalar(arr, scalar):
blockDim, gridDim = ccfg.DetermineCudaConfig(arr.shape[0])
# arrRes = cuda.device_array(arr.shape, dtype=arr.dtype)
MultArrayByScalar_dev[gridDim, blockDim](arr, scalar)
# arr = cuda.to_device(arrRes)
return arr
def AddArrayToArray(arr1, arr2):
blockDim, gridDim = ccfg.DetermineCudaConfig(arr1.shape[0])
AddArrayToArray_dev[gridDim, blockDim](arr1, arr2)
return arr1
def cos_dev_wrapper(arr):
blockDim, gridDim = ccfg.DetermineCudaConfig(arr.shape[0])
cos_arr = cuda.to_device(np.zeros(arr.shape, dtype=np.float32))
cos_dev[gridDim, blockDim](arr, cos_arr)
return cos_arr
def PerformIWFR(images, N):
imgHeight, imgWidth = images[0].height, images[0].width
ewfResultsDir = 'results/ewf/'
ewfAmName = 'am'
ewfPhName = 'ph'
# backCTFunctions = []
# forwCTFunctions = []
print('Starting IWFR...')
# storing contrast transfer functions
# for img in images:
# # print(imgWidth, img.px_dim, -img.defocus)
# ctf = CalcTransferFunction(imgWidth, img.px_dim, -img.defocus)
# ctf.AmPh2ReIm()
# # ctf.reIm = ccc.Diff2FFT(ctf.reIm)
# backCTFunctions.append(ctf)
#
# ctf = CalcTransferFunction(imgWidth, img.px_dim, img.defocus)
# ctf.AmPh2ReIm()
# # ctf.reIm = ccc.Diff2FFT(ctf.reIm)
# forwCTFunctions.append(ctf)
# start IWFR procedure
exitWave = imsup.Image(imgHeight, imgWidth, imsup.Image.cmp['CRI'],
imsup.Image.mem['GPU'])
for i in range(0, N):
print('Iteration no {0}...'.format(i + 1))
imsup.ClearImageData(exitWave)
# backpropagation (to in-focus plane)
ccfg.GetGPUMemoryUsed() # !!!
for img, idx in zip(images, range(0, len(images))):
# img = PropagateWave(img, backCTFunctions[idx]) # faster, but uses more memory
img = PropagateToFocus(img) # slower, but uses less memory
img.AmPh2ReIm()
exitWave.reIm = arrsup.AddTwoArrays(exitWave.reIm, img.reIm)
exitWave.reIm = arrsup.MultArrayByScalar(exitWave.reIm,
1 / len(images))
ewfAmPath = ewfResultsDir + ewfAmName + str(i + 1) + '.png'
ewfPhPath = ewfAmPath.replace(ewfAmName, ewfPhName)
ewfAmplPhPath = ewfPhPath.replace(ewfPhName, ewfPhName + 'Ampl')
imsup.SaveAmpImage(exitWave, ewfAmPath)
imsup.SavePhaseImage(exitWave, ewfPhPath)
amplifExitPhase = FilterPhase(exitWave)
imsup.SavePhaseImage(amplifExitPhase, ewfAmplPhPath)
# forward propagation (to the original focus plane)
ccfg.GetGPUMemoryUsed() # !!!
for img, idx in zip(images, range(0, len(images))):
# images[idx] = PropagateWave(exitWave, forwCTFunctions[idx])
images[idx] = PropagateBackToDefocus(exitWave, img.defocus)
images[idx].amPh.am = img.amPh.am # restore original amplitude
print('All done')
return exitWave