result_path = f'./models'
model_path = opt.model_path # path for new model checkpoints
utils.cond_mkdir(model_path)
phase_path = opt.phase_path # path of precomputed phase pool
data_path = f'./data/train1080' # path of targets
# Hardware setup
camera_prop = PhysicalProp(channel, laser_arduino=True, roi_res=(roi_res[1], roi_res[0]), slm_settle_time=0.15,
range_row=(220, 1000), range_col=(300, 1630),
patterns_path=opt.calibration_path, # path of 21 x 12 calibration patterns, see Supplement.
show_preview=True)
# Model instance to train
# Check propagation_model.py for the default parameter settings!
blur = utils.make_kernel_gaussian(0.85, 3) # Optional, just be consistent with inference.
model = ModelPropagate(distance=prop_dist,
feature_size=feature_size,
wavelength=wavelength,
blur=blur).to(device)
if opt.pretrained_path != '':
print(f' - Start from pre-trained model: {opt.pretrained_model_path}')
checkpoint = torch.load(opt.pretrained_model_path)
model.load_state_dict(checkpoint)
model = model.train()
# Augmented image loader (If you want to shuffle, augment dataset, put options accordingly)
image_loader = ImageLoader(data_path,
channel=channel,
batch_size=opt.batch_size,
# set model instance as naive ASM
model_prop = ModelPropagate(distance=prop_dist, feature_size=feature_size, wavelength=wavelength,
target_field=False, num_gaussians=0, num_coeffs_fourier=0,
use_conv1d_mlp=False, num_latent_codes=[0],
norm=None, blur=None, content_field=False, proptype=opt.proptype).to(device)
zernike_coeffs = None
source_amplitude = None
latent_codes = None
u_t = None
else:
if opt.manual_aberr_corr:
final_phase_num_in = 2 + opt.num_latent_codes
else:
final_phase_num_in = 4
blur = utils.make_kernel_gaussian(0.849, 3)
# load camera model and set it into eval mode
model_prop = ModelPropagate(distance=prop_dist,
feature_size=feature_size,
wavelength=wavelength,
blur=blur).to(device)
model_prop.load_state_dict(torch.load(opt.model_path, map_location=device))
# Here, we crop model parameters to match the Holonet resolution, which is slightly different from 1080p.
# parameters for CITL model
zernike_coeffs = model_prop.coeffs_fourier
source_amplitude = model_prop.source_amp
latent_codes = model_prop.latent_code
latent_codes = utils.crop_image(latent_codes, target_shape=image_res, pytorch=True, stacked_complex=False)