start_epoch = 1
mode = opt.MODEL.MODE
session = opt.MODEL.SESSION
result_dir = os.path.join(opt.TRAINING.SAVE_DIR, mode, 'results', session)
model_dir = os.path.join(opt.TRAINING.SAVE_DIR, mode, 'models', session)
utils.mkdir(result_dir)
utils.mkdir(model_dir)
train_dir = opt.TRAINING.TRAIN_DIR
val_dir = opt.TRAINING.VAL_DIR
######### Model ###########
model_restoration = MPRNet()
model_restoration.cuda()
device_ids = [i for i in range(torch.cuda.device_count())]
if torch.cuda.device_count() > 1:
print("\n\nLet's use", torch.cuda.device_count(), "GPUs!\n\n")
new_lr = opt.OPTIM.LR_INITIAL
optimizer = optim.Adam(model_restoration.parameters(),
lr=new_lr,
betas=(0.9, 0.999),
eps=1e-8)
######### Scheduler ###########
warmup_epochs = 3
type=str,
help='Path to weights')
parser.add_argument(
'--dataset', default='GoPro', type=str,
help='Test Dataset') # ['GoPro', 'HIDE', 'RealBlur_J', 'RealBlur_R']
parser.add_argument('--gpus',
default='0',
type=str,
help='CUDA_VISIBLE_DEVICES')
args = parser.parse_args()
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
model_restoration = MPRNet()
utils.load_checkpoint(model_restoration, args.weights)
print("===>Testing using weights: ", args.weights)
model_restoration.cuda()
model_restoration = nn.DataParallel(model_restoration)
model_restoration.eval()
dataset = args.dataset
rgb_dir_test = os.path.join(args.input_dir, dataset, 'test', 'input')
test_dataset = get_test_data(rgb_dir_test, img_options={})
test_loader = DataLoader(dataset=test_dataset,
batch_size=1,
shuffle=False,
num_workers=4,
drop_last=False,
parser = argparse.ArgumentParser(description='parser for video prediction')
parser.add_argument('--dataset', type=str, required=True, choices=['gopro', 'reds'], help='datasets')
parser.add_argument('--data_root', type=str, default='dataset', help='input image directory')
parser.add_argument('--deblur', type=str, default='models/center_v3.pth', help='deblurring directory')
parser.add_argument('--output_path', type=str, required=True, help='directory to save outputs')
parser.add_argument('--seq_len', type=int, required=True, help='Length of sequence')
parser.add_argument('--cuda', action='store_true', help='use cuda')
parser.add_argument('--save_img', action='store_true', help='save images')
args = parser.parse_args()
device = 'cuda' if args.cuda else 'cpu'
# load model
if 'mprnet' in args.deblur:
model1 = MPRNet()
else:
model1 = centerEsti()
model2 = F35_N8()
model3 = F26_N9()
model4 = F17_N9()
flowComp = superslomo.UNet(6, 4)
flowComp.to(device)
ArbTimeFlowIntrp = superslomo.UNet(20, 5)
ArbTimeFlowIntrp.to(device)
validationFlowBackWarp = superslomo.backWarp(1280, 704, device)
validationFlowBackWarp = validationFlowBackWarp.to(device)
if 'center_v3' in args.deblur:
checkpoint = torch.load('models/center_v3.pth')
[[-98 << 16, 6640., 1 / 771], [68 << 19, 1152., 0],
[53 << 18, 1088., 0]]])
mode_switch = {
'dehaze': ('./model/dehaze/AOD_net_epoch_relu_10.pth', AODnet, ramCoef[0],
1, 8, normalize),
'sun':
('./model/demoire/sun_epoch_200.pth', SUNNet, ramCoef[1], 9, 32, identity),
'moire_obj':
('./model/demoire/moire_obj.pth', ObjNet, ram2[0], 9, 128, identity),
'moire_screen_gan': ('./model/demoire/moire_screen_gan.pth', GANNet,
ram2[1], 17, 512, identity),
'MPRNet_deblurring': ('./model/MPRNet/model_deblurring.pth', MPRNet,
ramCoef[2], 7, 8, identity),
'MPRNet_deraining':
('./model/MPRNet/model_deraining.pth',
(lambda: MPRNet(n_feat=40, scale_unetfeats=20, scale_orsnetfeats=16)),
ramCoef[3], 7, 8, identity),
}
def getOpt(optDe):
model = optDe.get('model', 'dehaze')
opt = Option()
modelPath, opt.modelDef, ramCoef, opt.padding, opt.align, opt.prepare = mode_switch[
model]
opt.ramCoef = ramCoef[config.getRunType()]
opt.model = modelPath
opt.modelCached = initModel(opt, modelPath, model)
return opt
# -*- coding:utf-8 -*-
import numpy as np
from imageProcess import initModel, Option
from models import NetDN, SEDN
from MPRNet import MPRNet
from config import config
ramCoef = .9 / np.array([[2700., 2400., 1253.4], [4106.9, 7405., 4304.2], [60493., 8400., 4288.]])
mode_switch = {
'15': ('./model/l15/model_new.pth', SEDN, ramCoef[1], 1),
'25': ('./model/l25/model_new.pth', SEDN, ramCoef[1], 1),
'50': ('./model/l50/model_new.pth', SEDN, ramCoef[1], 1),
'lite5': ('./model/dn_lite5/model_new.pth', NetDN, ramCoef[0], 1),
'lite10': ('./model/dn_lite10/model_new.pth', NetDN, ramCoef[0], 1),
'lite15': ('./model/dn_lite15/model_new.pth', NetDN, ramCoef[0], 1),
'MPRNet_denoising': ('./model/MPRNet/model_denoising.pth', (lambda: MPRNet(n_feat=80, scale_unetfeats=48, scale_orsnetfeats=32)), ramCoef[2], 0),
}
##################################
def getOpt(optDN):
model = optDN['model']
if not model in mode_switch:
return
opt = Option(mode_switch[model][0])
opt.modelDef = mode_switch[model][1]
opt.ramCoef = mode_switch[model][2][config.getRunType()]
opt.cropsize = config.getConfig()[1 if model[:4] == 'lite' else 2]
opt.modelCached = initModel(opt, opt.model, 'DN' + model)
sd = mode_switch[model][3]