def define_G(opt): opt_net = opt['network_G'] which_model = opt_net['which_model_G'] # image restoration if which_model == 'MSRResNet': netG = SRResNet_arch.MSRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'], upscale=opt_net['scale']) elif which_model == 'RRDBNet': netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb']) # 16x superresolution elif which_model == 'RRDBNet_16x': netG = RRDBNet_arch.RRDBNet_16x(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb']) # 16x superresolution with transposed conv elif which_model == 'RRDBNetTRConv_16x': netG = RRDBNet_arch.RRDBNetTRConv_16x(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb']) # video restoration elif which_model == 'EDVR': netG = EDVR_arch.EDVR(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) else: raise NotImplementedError('Generator model [{:s}] not recognized'.format(which_model)) return netG
def define_G(opt): opt_net = opt['network_G'] which_model = opt_net['which_model_G'] # image restoration if which_model == 'MSRResNet': netG = SRResNet_arch.MSRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'], upscale=opt_net['scale']) elif which_model == 'RRDBNet': netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb']) elif which_model == 'ORDSRNet': netG = ORDSRModel(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], N=opt_net['base_size'], S=opt_net['stride'], upscale=opt_net['scale']) else: raise NotImplementedError( 'Generator model [{:s}] not recognized'.format(which_model)) return netG
def define_G(opt): opt_net = opt['network_G'] which_model = opt_net['which_model_G'] # image restoration if which_model == 'MSRResNet': netG = SRResNet_arch.MSRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'], upscale=opt_net['scale']) elif which_model == 'RRDBNet': netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb']) # video restoration elif which_model == 'EDVR': netG = EDVR_arch.EDVR(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA'], scale=opt['scale']) elif which_model == 'DUF': if opt_net['layers'] == 16: netG = DUF_arch.DUF_16L(scale=opt['scale'], adapt_official=True) elif opt_net['layers'] == 28: netG = DUF_arch.DUF_28L(scale=opt['scale'], adapt_official=True) else: netG = DUF_arch.DUF_52L(scale=opt['scale'], adapt_official=True) elif which_model == 'TOF': netG = TOF_arch.TOFlow(adapt_official=True) else: raise NotImplementedError('Generator model [{:s}] not recognized'.format(which_model)) return netG
def define_G(opt): opt_net = opt['network_G'] which_model = opt_net['which_model_G'] # image restoration if which_model == 'MSRResNet': netG = SRResNet_arch.MSRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'], upscale=opt_net['scale']) elif which_model == 'RRDBNet': netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb']) elif which_model == 'AdaFMNet': netG = AdaFMNet_arch.AdaFMNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'], adafm_ksize=opt_net['adafm_ksize']) elif which_model == 'CResMDNet': netG = CResMDNet_arch.CResMDNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'], cond_dim=opt_net['cond_dim']) elif which_model == 'BaseNet': netG = CResMDNet_arch.BaseNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb']) elif which_model == 'CondNet': netG = CResMDNet_arch.CondNet(in_nc=opt_net['in_nc'], nf=opt_net['nf']) # video restoration elif which_model == 'EDVR': netG = EDVR_arch.EDVR(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) else: raise NotImplementedError( 'Generator model [{:s}] not recognized'.format(which_model)) return netG
def define_G(opt): opt_net = opt['network_G'] which_model = opt_net['which_model_G'] # image restoration if which_model == 'MSRResNet': netG = SRResNet_arch.MSRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'], upscale=opt_net['scale']) elif which_model == 'RRDBNet': netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb']) # video restoration elif which_model == 'EDVR': import models.archs.EDVR_arch as EDVR_arch netG = EDVR_arch.EDVR(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA'], w_GCB=opt_net['w_GCB']) # elif which_model == 'EDVR_woDCN': # import models.archs.EDVR_woDCN_arch as EDVR_arch # netG = EDVR_arch.EDVR(nf=opt_net['nf'], nframes=opt_net['nframes'], # groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], # back_RBs=opt_net['back_RBs'], center=opt_net['center'], # predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], # w_TSA=opt_net['w_TSA'], w_GCB=opt_net['w_GCB']) elif which_model == 'MGANet': netG = Gen_Guided_UNet(input_size=opt_net['input_size']) elif which_model == 'Unet': import repo.CycleGAN.networks as unet_networks netG = unet_networks.define_G(2 * 3, 1, opt_net['nf'], opt_net['G_type'], opt_net['norm'], opt_net['dropout'], opt_net['init_type'], opt_net['init_gain']) else: raise NotImplementedError( 'Generator model [{:s}] not recognized'.format(which_model)) return netG
def _torch_infer(self, img): if torch.cuda.is_available(): device = torch.device('cuda') else: device = torch.device('cpu') torch_input = torch.from_numpy(img).to(device) model = arch.RRDBNet(in_nc=3, out_nc=3, nf=64, nb=23, upscale=2) model_bytes = torch.load(self._model_path) model.load_state_dict(model_bytes, strict=False) model.eval() model = model.to(device) with torch.no_grad(): pred = model(torch_input) return pred
def define_G(opt): opt_net = opt['network_G'] which_model = opt_net['which_model_G'] # image restoration if which_model == 'CNLRN': netG = CNLRN_arch.CNLRN(n_colors=opt_net['n_colors'], n_deblur_blocks=opt_net['n_deblur_blocks'], n_nlrgs_body=opt_net['n_nlrgs_body'], n_nlrgs_up1=opt_net['n_nlrgs_up1'], n_nlrgs_up2=opt_net['n_nlrgs_up2'], n_subgroups=opt_net['n_subgroups'], n_rcabs=opt_net['n_rcabs'], n_feats=opt_net['n_feats'], nonlocal_psize=opt_net['nonlocal_psize'], scale=opt_net['scale']) elif which_model == 'PreDeblur': netG = PreDeblur_arch.PreDeblur( n_colors=opt_net['n_colors'], n_deblur_blocks=opt_net['n_deblur_blocks'], n_feats=opt_net['n_feats']) elif which_model == 'RRDBNet': netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb']) # video restoration elif which_model == 'EDVR': netG = EDVR_arch.EDVR(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) else: raise NotImplementedError( 'Generator model [{:s}] not recognized'.format(which_model)) return netG
def define_G(opt): opt_net = opt['network_G'] which_model = opt_net['which_model_G'] # image restoration if which_model == 'MSRResNet': netG = SRResNet_arch.MSRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'], upscale=opt_net['scale']) elif which_model == 'RRDBNet': netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb']) # video restoration elif which_model == 'EDVR': netG = EDVR_arch.EDVR(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) elif which_model == 'JASRNet': netG = JASRNet_arch.JASR(n_Parts=opt_net['n_Parts'], n_resblocks=opt_net['n_resblocks'], n_feats=opt_net['n_feats'], scale=opt_net['scale'], rgb_range=opt_net['rgb_range'], n_colors=opt_net['n_colors']) else: raise NotImplementedError( 'Generator model [{:s}] not recognized'.format(which_model)) return netG
def define_G(opt): opt_net = opt['network_G'] which_model = opt_net['which_model_G'] # image restoration if which_model == 'MSRResNet': netG = SRResNet_arch.MSRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'], upscale=opt_net['scale']) elif which_model == 'RRDBNet': netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb']) elif which_model == 'RCAN': netG = RCAN_arch.RCAN(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], n_features=opt_net['nf'], n_resgroups=opt_net['ng'], n_resblocks=opt_net['nb'], reduction=opt_net['reduction'], scale=opt_net['scale'], res_scale=opt_net['res_scale']) # video restoration elif which_model == 'EDVR': netG = EDVR_arch.EDVR(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) elif which_model == 'EDVR_DN': netG = EDVR_arch.EDVR_DN(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) elif which_model == 'EDVR_pyramid': netG = EDVR_arch.EDVR_pyramid(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) elif which_model == 'PFNL': netG = PFNL_arch.PFNL(nf=opt_net['nf'], nc=opt_net['nc'], nt=opt_net['nt'], r=opt_net['r'], scale=opt_net['scale']) else: raise NotImplementedError( 'Generator model [{:s}] not recognized'.format(which_model)) return netG
def define_G(opt): opt_net = opt['network_G'] which_model = opt_net['which_model_G'] # image restoration if which_model == 'MSRResNet': netG = SRResNet_arch.MSRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'], upscale=opt_net['scale']) elif which_model == 'RRDBNet': netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb']) # video restoration elif which_model == 'EDVR': netG = EDVR_arch.EDVR(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) elif which_model == 'EDVR2X': netG = EDVR_arch.EDVR2X(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) elif which_model == 'EDVRImg': netG = EDVR_arch.EDVRImage(nf=opt_net['nf'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], down_scale=opt_net['down_scale']) elif which_model == 'EDVR3D': netG = EDVR_arch.EDVR3D(nf=opt_net['nf'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], down_scale=opt_net['down_scale']) elif which_model == 'UPEDVR': netG = EDVR_arch.UPEDVR(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], w_TSA=opt_net['w_TSA'], down_scale=opt_net['down_scale'], align_target=opt_net['align_target'], ret_valid=opt_net['ret_valid']) elif which_model == 'UPContEDVR': netG = EDVR_arch.UPControlEDVR( nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], w_TSA=opt_net['w_TSA'], down_scale=opt_net['down_scale'], align_target=opt_net['align_target'], ret_valid=opt_net['ret_valid'], multi_scale_cont=opt_net['multi_scale_cont']) elif which_model == 'FlowUPContEDVR': netG = EDVR_arch.FlowUPControlEDVR( nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], w_TSA=opt_net['w_TSA'], down_scale=opt_net['down_scale'], align_target=opt_net['align_target'], ret_valid=opt_net['ret_valid'], multi_scale_cont=opt_net['multi_scale_cont']) # video SR for multiple target frames elif which_model == 'MultiEDVR': netG = EDVR_arch.MultiEDVR(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) # arbitrary magnification video super-resolution elif which_model == 'MetaEDVR': netG = EDVR_arch.MetaEDVR(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA'], fix_edvr=opt_net['fix_edvr']) else: raise NotImplementedError( 'Generator model [{:s}] not recognized'.format(which_model)) return netG
def define_G(opt): opt_net = opt['network_G'] which_model = opt_net['which_model_G'] # image restoration if which_model == 'MSRResNet': netG = SRResNet_arch.MSRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'], upscale=opt_net['scale']) elif which_model == 'RRDBNet': netG = RRDBNet_arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb']) # video restoration elif which_model == 'EDVR': netG = EDVR_arch.EDVR(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) elif which_model == 'MY_EDVR_FusionDenoise': netG = my_EDVR_arch.MYEDVR_FusionDenoise( nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) elif which_model == 'MY_EDVR_RES': netG = my_EDVR_arch.MYEDVR_RES(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) elif which_model == 'MY_EDVR_PreEnhance': netG = my_EDVR_arch.MYEDVR_PreEnhance(nf=opt_net['nf'], nframes=opt_net['nframes'], groups=opt_net['groups'], front_RBs=opt_net['front_RBs'], back_RBs=opt_net['back_RBs'], center=opt_net['center'], predeblur=opt_net['predeblur'], HR_in=opt_net['HR_in'], w_TSA=opt_net['w_TSA']) elif which_model == 'Recurr_ResBlocks': netG = Recurr_arch.Recurr_ResBlocks( nf=opt_net['nf'], N_RBs=opt_net['N_RBs'], N_flow_lv=opt_net['N_flow_lv'], pretrain_flow=opt_net['pretrain_flow']) else: raise NotImplementedError( 'Generator model [{:s}] not recognized'.format(which_model)) return netG