def load_burst_n2n_model(cfg):
# -- init --
init_lr = cfg.init_lr
# -- load kpn model --
kpn,_ = load_model_kpn_1f(cfg)
# -- load unet info --
unet_info = edict()
unet_info.model = UNet_n2n( 1,3,3)
cfg.init_lr = 1e-4
unet_info.optim = load_optimizer(cfg,unet_info.model)
unet_info.S = None
# -- load noise critic info --
disc_model = DCGAN_D(64, -1,3,64,1,0)
disc_model = disc_model.cuda(cfg.gpuid)
cfg.init_lr = 1*1e-4
sim_params = edict({'mean':0,'std':25./255,'noise_type':'gaussian'})
disc_optim = load_optimizer_gan(cfg,disc_model)
noise_critic = NoiseCriticModel(disc_model,disc_optim,sim_params,cfg.device)
# -- create burstaligned model --
model = BurstAlignN2N(kpn,unet_info,noise_critic)
criterion = BurstRecLoss(noise_critic,alpha=1.0)
# -- model models to cuda --
model = model.cuda(cfg.gpuid)
model.unet_info.model = model.unet_info.model.cuda(cfg.gpuid)
# -- finally --
cfg.init_lr = init_lr
return model,criterion
def load_burst_kpn_model(cfg):
# -- init --
init_lr = cfg.init_lr
# -- load alignment model --
align_info = edict()
align_info.model, _ = load_model_kpn_1f_cls_cascade(cfg)
align_info.optim = load_optimizer(cfg, align_info.model)
align_info.S = None
# -- load denoising model --
denoiser_info = edict()
denoiser_info.model, _ = load_model_kpn_cascade(cfg)
# denoiser_info.model,_ = load_model_kpn(cfg,cfg.N)
denoiser_info.optim = load_optimizer(cfg, denoiser_info.model)
denoiser_info.S = None
# -- load unet model --
unet_info = edict()
unet_info.model = load_model_unet(cfg).to(cfg.device)
# unet_info.model = load_model_unet(cfg).to(cfg.device)
# unet_info.model = init_net(unet_info.model)
# unet_info.model,_ = load_model_kpn(cfg,cfg.N)
unet_info.optim = load_optimizer(cfg, unet_info.model)
unet_info.S = None
# -- create burstaligned model --
use_align = cfg.burst_use_alignment
use_unet = cfg.burst_use_unet
use_unet_only = cfg.burst_use_unet_only
criterion = BurstRecLoss(alpha=1.0)
model = BurstAlignSG(align_info,
denoiser_info,
unet_info,
use_alignment=use_align,
use_unet=use_unet,
use_unet_only=use_unet_only,
kpn_num_frames=cfg.kpn_num_frames)
# -- load noise critic info --
disc_model = DCGAN_D(64, -1, 3, 64, 1, 0)
disc_model = disc_model.cuda(cfg.gpuid)
cfg.init_lr = 1 * 1e-4
sim_params = edict({'mean': 0, 'std': 25. / 255, 'noise_type': 'gaussian'})
disc_optim = load_optimizer_gan(cfg, disc_model)
p_lambda = 10
noise_critic = NoiseCriticModel(disc_model, disc_optim, sim_params,
cfg.device, p_lambda)
# -- model models to cuda --
model = model.cuda(cfg.gpuid)
model.denoiser_info.model = model.denoiser_info.model.cuda(cfg.gpuid)
model.align_info.model = model.align_info.model.cuda(cfg.gpuid)
# -- finally --
cfg.init_lr = init_lr
return model, noise_critic, criterion
def load_model_disc(cfg, use_simclr=False):
if use_simclr:
model = DiscriminatorSimCLR(32, 2, 64, 1, 0)
else:
model = DCGAN_D(32, 100, 3, 64, 1, 0)
# model = Discriminator(32,3)
model = model.to(cfg.device)
return model
def load_burst_stn_model(cfg):
# -- init --
init_lr = cfg.init_lr
# -- load alignment model --
align_info = edict()
align_info.model, _ = load_model_stn(cfg)
align_info.optim = load_optimizer(cfg, align_info.model)
align_info.S = None
# -- load denoising model --
denoiser_info = edict()
denoiser_info.model, _ = load_model_kpn(cfg)
denoiser_info.optim = load_optimizer(cfg, denoiser_info.model)
denoiser_info.S = None
# -- create burstaligned model --
model = BurstAlignSTN(align_info, denoiser_info)
criterion = BurstRecLoss(alpha=1.0)
# -- load noise critic info --
disc_model = DCGAN_D(64, -1, 3, 64, 1, 0)
disc_model = disc_model.cuda(cfg.gpuid)
cfg.init_lr = 1 * 1e-4
sim_params = edict({'mean': 0, 'std': 25. / 255, 'noise_type': 'gaussian'})
disc_optim = load_optimizer_gan(cfg, disc_model)
p_lambda = 10
noise_critic = NoiseCriticModel(disc_model, disc_optim, sim_params,
cfg.device, p_lambda)
# -- model models to cuda --
model = model.cuda(cfg.gpuid)
model.denoiser_info.model = model.denoiser_info.model.cuda(cfg.gpuid)
model.align_info.model = model.align_info.model.cuda(cfg.gpuid)
# -- finally --
cfg.init_lr = init_lr
return model, noise_critic, criterion