| import os |
| os.environ['CUDA_VISIBLE_DEVICES'] = '0' |
| import argparse |
| from tqdm import tqdm |
| from data.data import * |
| from torchvision import transforms |
| from torch.utils.data import DataLoader |
| from loss.losses import * |
| from net.IADNet import IADNet |
|
|
| def unwrap_model(net): |
| return net.module if isinstance(net, torch.nn.DataParallel) else net |
|
|
|
|
| def _normalize_state_dict_keys(state_dict, target_has_module_prefix): |
| has_module_prefix = all(k.startswith("module.") for k in state_dict.keys()) |
| if has_module_prefix == target_has_module_prefix: |
| return state_dict |
|
|
| if target_has_module_prefix: |
| return {f"module.{k}": v for k, v in state_dict.items()} |
| return {k[len("module."):]: v if k.startswith("module.") else v for k, v in state_dict.items()} |
|
|
|
|
| def load_checkpoint_flexible(model, model_path): |
| checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) |
| state_dict = checkpoint["state_dict"] if isinstance(checkpoint, dict) and "state_dict" in checkpoint else checkpoint |
| model_state_keys = model.state_dict().keys() |
| target_has_module_prefix = all(k.startswith("module.") for k in model_state_keys) |
| state_dict = _normalize_state_dict_keys(state_dict, target_has_module_prefix) |
| model.load_state_dict(state_dict, strict=True) |
|
|
| def eval(model, testing_data_loader, model_path, output_folder,norm_size=True,LOL=False,v2=False,unpaired=False,alpha=1.0,gamma=1.0): |
| torch.set_grad_enabled(False) |
| model_core = unwrap_model(model) |
| load_checkpoint_flexible(model, model_path) |
| print('Pre-trained model is loaded.') |
| model.eval() |
| print('Evaluation:') |
| if LOL: |
| model_core.trans.gated = True |
| elif v2: |
| model_core.trans.gated2 = True |
| model_core.trans.alpha = alpha |
| elif unpaired: |
| model_core.trans.gated2 = True |
| model_core.trans.alpha = alpha |
| for batch in tqdm(testing_data_loader): |
| with torch.no_grad(): |
| if norm_size: |
| input, name = batch[0], batch[1] |
| else: |
| input, name, h, w = batch[0], batch[1], batch[2], batch[3] |
| |
| input = input.cuda() |
| output = model(input**gamma) |
| |
| if not os.path.exists(output_folder): |
| os.mkdir(output_folder) |
| |
| output = torch.clamp(output.cuda(),0,1).cuda() |
| if not norm_size: |
| output = output[:, :, :h, :w] |
| |
| output_img = transforms.ToPILImage()(output.squeeze(0)) |
| output_img.save(output_folder + name[0]) |
| torch.cuda.empty_cache() |
| print('===> End evaluation') |
| if LOL: |
| model_core.trans.gated = False |
| elif v2: |
| model_core.trans.gated2 = False |
| torch.set_grad_enabled(True) |
| |
| if __name__ == '__main__': |
| |
| eval_parser = argparse.ArgumentParser(description='Eval') |
| eval_parser.add_argument('--perc', action='store_true', help='trained with perceptual loss') |
| eval_parser.add_argument('--lol', action='store_true', help='output lolv1 dataset') |
| eval_parser.add_argument('--lol_v2_real', action='store_true', help='output lol_v2_real dataset') |
| eval_parser.add_argument('--lol_v2_syn', action='store_true', help='output lol_v2_syn dataset') |
| eval_parser.add_argument('--SICE_grad', action='store_true', help='output SICE_grad dataset') |
| eval_parser.add_argument('--SICE_mix', action='store_true', help='output SICE_mix dataset') |
| eval_parser.add_argument('--fivek', action='store_true', help='output FiveK dataset') |
|
|
| eval_parser.add_argument('--best_GT_mean', action='store_true', help='output lol_v2_real dataset best_GT_mean') |
| eval_parser.add_argument('--best_PSNR', action='store_true', help='output lol_v2_real dataset best_PSNR') |
| eval_parser.add_argument('--best_SSIM', action='store_true', help='output lol_v2_real dataset best_SSIM') |
|
|
| eval_parser.add_argument('--custome', action='store_true', help='output custome dataset') |
| eval_parser.add_argument('--custome_path', type=str, default='./YOLO') |
| eval_parser.add_argument('--unpaired', action='store_true', help='output unpaired dataset') |
| eval_parser.add_argument('--DICM', action='store_true', help='output DICM dataset') |
| eval_parser.add_argument('--LIME', action='store_true', help='output LIME dataset') |
| eval_parser.add_argument('--MEF', action='store_true', help='output MEF dataset') |
| eval_parser.add_argument('--NPE', action='store_true', help='output NPE dataset') |
| eval_parser.add_argument('--VV', action='store_true', help='output VV dataset') |
| eval_parser.add_argument('--alpha', type=float, default=1.0) |
| eval_parser.add_argument('--gamma', type=float, default=1.0) |
| eval_parser.add_argument('--unpaired_weights', type=str, default='./weights/LOLv2_syn/w_perc.pth') |
|
|
| ep = eval_parser.parse_args() |
|
|
|
|
| cuda = True |
| if cuda and not torch.cuda.is_available(): |
| raise Exception("No GPU found, or need to change CUDA_VISIBLE_DEVICES number") |
| |
| if not os.path.exists('./output'): |
| os.mkdir('./output') |
| |
| norm_size = True |
| num_workers = 1 |
| alpha = None |
| if ep.lol: |
| eval_data = DataLoader(dataset=get_eval_set("./datasets/LOLdataset/eval15/low"), num_workers=num_workers, batch_size=1, shuffle=False) |
| output_folder = './output/LOLv1/' |
| if ep.perc: |
| weight_path = './weights/LOLv1/w_perc.pth' |
| else: |
| weight_path = './weights/LOLv1/wo_perc.pth' |
| |
| |
| elif ep.lol_v2_real: |
| eval_data = DataLoader(dataset=get_eval_set("./datasets/LOLv2/Real_captured/Test/Low"), num_workers=num_workers, batch_size=1, shuffle=False) |
| output_folder = './output/LOLv2_real/' |
| if ep.best_GT_mean: |
| weight_path = './weights/LOLv2_real/w_perc.pth' |
| alpha = 0.84 |
| elif ep.best_PSNR: |
| weight_path = './weights/LOLv2_real/best_PSNR.pth' |
| alpha = 0.8 |
| elif ep.best_SSIM: |
| weight_path = './weights/LOLv2_real/best_SSIM.pth' |
| alpha = 0.82 |
| |
| elif ep.lol_v2_syn: |
| eval_data = DataLoader(dataset=get_eval_set("./datasets/LOLv2/Synthetic/Test/Low"), num_workers=num_workers, batch_size=1, shuffle=False) |
| output_folder = './output/LOLv2_syn/' |
| if ep.perc: |
| weight_path = './weights/LOLv2_syn/w_perc.pth' |
| else: |
| weight_path = './weights/LOLv2_syn/wo_perc.pth' |
| |
| elif ep.SICE_grad: |
| eval_data = DataLoader(dataset=get_SICE_eval_set("./datasets/SICE/SICE_Grad"), num_workers=num_workers, batch_size=1, shuffle=False) |
| output_folder = './output/SICE_grad/' |
| weight_path = './weights/SICE.pth' |
| norm_size = False |
| |
| elif ep.SICE_mix: |
| eval_data = DataLoader(dataset=get_SICE_eval_set("./datasets/SICE/SICE_Mix"), num_workers=num_workers, batch_size=1, shuffle=False) |
| output_folder = './output/SICE_mix/' |
| weight_path = './weights/SICE.pth' |
| norm_size = False |
| |
| elif ep.fivek: |
| eval_data = DataLoader(dataset=get_SICE_eval_set("./datasets/FiveK/test/input"), num_workers=num_workers, batch_size=1, shuffle=False) |
| output_folder = './output/fivek/' |
| weight_path = './weights/fivek.pth' |
| norm_size = False |
| |
| elif ep.unpaired: |
| if ep.DICM: |
| eval_data = DataLoader(dataset=get_SICE_eval_set("./datasets/DICM"), num_workers=num_workers, batch_size=1, shuffle=False) |
| output_folder = './output/DICM/' |
| elif ep.LIME: |
| eval_data = DataLoader(dataset=get_SICE_eval_set("./datasets/LIME"), num_workers=num_workers, batch_size=1, shuffle=False) |
| output_folder = './output/LIME/' |
| elif ep.MEF: |
| eval_data = DataLoader(dataset=get_SICE_eval_set("./datasets/MEF"), num_workers=num_workers, batch_size=1, shuffle=False) |
| output_folder = './output/MEF/' |
| elif ep.NPE: |
| eval_data = DataLoader(dataset=get_SICE_eval_set("./datasets/NPE"), num_workers=num_workers, batch_size=1, shuffle=False) |
| output_folder = './output/NPE/' |
| elif ep.VV: |
| eval_data = DataLoader(dataset=get_SICE_eval_set("./datasets/VV"), num_workers=num_workers, batch_size=1, shuffle=False) |
| output_folder = './output/VV/' |
| elif ep.custome: |
| eval_data = DataLoader(dataset=get_SICE_eval_set(ep.custome_path), num_workers=num_workers, batch_size=1, shuffle=False) |
| output_folder = './output/custome/' |
| alpha = ep.alpha |
| norm_size = False |
| weight_path = ep.unpaired_weights |
| |
| eval_net = IADNet().cuda() |
| eval(eval_net, eval_data, weight_path, output_folder,norm_size=norm_size,LOL=ep.lol,v2=ep.lol_v2_real,unpaired=ep.unpaired,alpha=alpha,gamma=ep.gamma) |
|
|
|
|
