from pycuda.autoinit import context
import pycuda.gpuarray as gpuarray
import skcuda.fft as cu_fft
import skcuda
from copy import deepcopy
from ptychoposcorrection import *
from time import time
import Kernels
# Load CUDA kernel modules, params: name, number of pixels, list of values to replace on source code
Kernels.LoadKernel("Scripts/PIEKernels.cu", np.int32(imagApert[0]*imagApert[1]),["#define NumModes X", NumModes])
""" Get CUDA kernel functions """
# Basic ePIE kernels
ReplaceInObject = Kernels.GetFunction("ReplaceInObject")
ApplyDifPad = Kernels.GetFunction("ApplyDifPad")
ExitwaveAndBuffer = Kernels.GetFunction("ExitwaveAndBuffer")
ApertureAbs2 = Kernels.GetFunction("ApertureAbs2")
ObjectAbs2 = Kernels.GetFunction("ObjectAbs2")
CopyFromROI = Kernels.GetFunction("CopyFromROI")
CropObject = Kernels.GetFunction("CropObject")
UpdateProbeAndRspace = Kernels.GetFunction("UpdateProbeAndRspace")
# Only used if bPhaseShift = True
PhaseShiftFunc = Kernels.GetFunction("PhaseShift")
# Multiprobe
ModeSub21 = Kernels.GetFunction("ModeSub21")
ScalarProd = Kernels.GetFunction("ScalarProd")
ModeMultiply = Kernels.GetFunction("ModeMultiply")
iterations = 500
bFORCE_REAL = False # Set to true if the rspace is real.
IM_HALFSIZE = 480 # half the image size to use. Our data is limited to 960x960 (480).
PerfLevel = 3 # The error will be registered every X iterations. Higher number means better performance.
bUsePinholeMask = False # True = Load the pinhole mask. False = Use a msize X msize mask.
bUseCircularMask = True
msize = 98 # size of the mask to be used.
Radius = 49
Run = 0
beta = 0.9
Kernels.LoadKernel("Scripts/HIOKernels.cu",
np.int32(4 * IM_HALFSIZE * IM_HALFSIZE))
ApplyDifPad = Kernels.GetFunction("ApplyDifPad")
HIOStep = Kernels.GetFunction("HIOStep")
Copy = Kernels.GetFunction("Copy")
Error = Kernels.GetFunction("Error1DifCF")
def SaveImage(fname, Image):
np.save(fname + '_hio', Image)
absol = np.absolute(Image).astype(np.float32)
maxv = absol.max()
Normalized = (255 * np.sqrt((1.0 / maxv) * absol)).astype(np.uint8)
cv2.imwrite(fname + '_hio.png', Normalized)
def HIO(HDR_FILE_PATH, SAVE_FILES_PATH, Seed):