def main(): # ---------------------------------------- # Preparation # ---------------------------------------- noise_level_img = 0 / 255.0 # set AWGN noise level for LR image, default: 0 noise_level_model = noise_level_img # set noise level of model, default: 0 model_name = 'ircnn_color' # set denoiser, 'drunet_color' | 'ircnn_color' testset_name = 'Set18' # set testing set, 'set18' | 'set24' x8 = True # set PGSE to boost performance, default: True iter_num = 40 # set number of iterations, default: 40 for demosaicing modelSigma1 = 49 # set sigma_1, default: 49 modelSigma2 = max(0.6, noise_level_model * 255.) # set sigma_2, default matlab_init = True show_img = False # default: False save_L = True # save LR image save_E = True # save estimated image save_LEH = False # save zoomed LR, E and H images border = 10 # default 10 for demosaicing task_current = 'dm' # 'dm' for demosaicing n_channels = 3 # fixed model_zoo = 'model_zoo' # fixed testsets = 'testsets' # fixed results = 'results' # fixed result_name = testset_name + '_' + task_current + '_' + model_name model_path = os.path.join(model_zoo, model_name + '.pth') device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') torch.cuda.empty_cache() # ---------------------------------------- # L_path, E_path, H_path # ---------------------------------------- L_path = os.path.join(testsets, testset_name) # L_path, for Low-quality images E_path = os.path.join(results, result_name) # E_path, for Estimated images util.mkdir(E_path) logger_name = result_name utils_logger.logger_info(logger_name, log_path=os.path.join(E_path, logger_name + '.log')) logger = logging.getLogger(logger_name) # ---------------------------------------- # load model # ---------------------------------------- if 'drunet' in model_name: from models.network_unet import UNetRes as net model = net(in_nc=n_channels + 1, out_nc=n_channels, nc=[64, 128, 256, 512], nb=4, act_mode='R', downsample_mode="strideconv", upsample_mode="convtranspose") model.load_state_dict(torch.load(model_path), strict=True) model.eval() for _, v in model.named_parameters(): v.requires_grad = False model = model.to(device) elif 'ircnn' in model_name: from models.network_dncnn import IRCNN as net model = net(in_nc=n_channels, out_nc=n_channels, nc=64) model25 = torch.load(model_path) former_idx = 0 logger.info('model_name:{}, image sigma:{:.3f}, model sigma:{:.3f}'.format( model_name, noise_level_img, noise_level_model)) logger.info('Model path: {:s}'.format(model_path)) logger.info(L_path) L_paths = util.get_image_paths(L_path) test_results = OrderedDict() test_results['psnr'] = [] for idx, img in enumerate(L_paths): # -------------------------------- # (1) get img_H and img_L # -------------------------------- idx += 1 img_name, ext = os.path.splitext(os.path.basename(img)) img_H = util.imread_uint(img, n_channels=n_channels) CFA, CFA4, mosaic, mask = utils_mosaic.mosaic_CFA_Bayer(img_H) # -------------------------------- # (2) initialize x # -------------------------------- if matlab_init: # matlab demosaicing for initialization CFA4 = util.uint2tensor4(CFA4).to(device) x = utils_mosaic.dm_matlab(CFA4) else: x = cv2.cvtColor(CFA, cv2.COLOR_BAYER_BG2RGB_EA) x = util.uint2tensor4(x).to(device) img_L = util.tensor2uint(x) y = util.uint2tensor4(mosaic).to(device) util.imshow(img_L) if show_img else None mask = util.single2tensor4(mask.astype(np.float32)).to(device) # -------------------------------- # (3) get rhos and sigmas # -------------------------------- rhos, sigmas = pnp.get_rho_sigma(sigma=max(0.255 / 255., noise_level_img), iter_num=iter_num, modelSigma1=modelSigma1, modelSigma2=modelSigma2, w=1.0) rhos, sigmas = torch.tensor(rhos).to(device), torch.tensor(sigmas).to( device) # -------------------------------- # (4) main iterations # -------------------------------- for i in range(iter_num): # -------------------------------- # step 1, closed-form solution # -------------------------------- x = (y + rhos[i].float() * x).div(mask + rhos[i]) # -------------------------------- # step 2, denoiser # -------------------------------- if 'ircnn' in model_name: current_idx = np.int( np.ceil(sigmas[i].cpu().numpy() * 255. / 2.) - 1) if current_idx != former_idx: model.load_state_dict(model25[str(current_idx)], strict=True) model.eval() for _, v in model.named_parameters(): v.requires_grad = False model = model.to(device) former_idx = current_idx x = torch.clamp(x, 0, 1) if x8: x = util.augment_img_tensor4(x, i % 8) if 'drunet' in model_name: x = torch.cat((x, sigmas[i].float().repeat( 1, 1, x.shape[2], x.shape[3])), dim=1) x = utils_model.test_mode(model, x, mode=2, refield=32, min_size=256, modulo=16) # x = model(x) elif 'ircnn' in model_name: x = model(x) if x8: if i % 8 == 3 or i % 8 == 5: x = util.augment_img_tensor4(x, 8 - i % 8) else: x = util.augment_img_tensor4(x, i % 8) x[mask.to(torch.bool)] = y[mask.to(torch.bool)] # -------------------------------- # (4) img_E # -------------------------------- img_E = util.tensor2uint(x) psnr = util.calculate_psnr(img_E, img_H, border=border) test_results['psnr'].append(psnr) logger.info('{:->4d}--> {:>10s} -- PSNR: {:.2f}dB'.format( idx, img_name + ext, psnr)) if save_E: util.imsave( img_E, os.path.join(E_path, img_name + '_' + model_name + '.png')) if save_L: util.imsave(img_L, os.path.join(E_path, img_name + '_L.png')) if save_LEH: util.imsave( np.concatenate([img_L, img_E, img_H], axis=1), os.path.join(E_path, img_name + model_name + '_LEH.png')) ave_psnr = sum(test_results['psnr']) / len(test_results['psnr']) logger.info('------> Average PSNR(RGB) of ({}) is : {:.2f} dB'.format( testset_name, ave_psnr))
def main(): utils_logger.logger_info('AIM-track', log_path='AIM-track.log') logger = logging.getLogger('AIM-track') # -------------------------------- # basic settings # -------------------------------- testsets = 'DIV2K' testset_L = 'DIV2K_valid_LR_bicubic' #testset_L = 'DIV2K_test_LR_bicubic' torch.cuda.current_device() torch.cuda.empty_cache() #torch.backends.cudnn.benchmark = True device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # -------------------------------- # load model # -------------------------------- model_path = os.path.join('trained_model', 'RFDN_AIM.pth') model = RFDN() model.load_state_dict(torch.load(model_path), strict=True) model.eval() for k, v in model.named_parameters(): v.requires_grad = False model = model.to(device) # number of parameters number_parameters = sum(map(lambda x: x.numel(), model.parameters())) logger.info('Params number: {}'.format(number_parameters)) # -------------------------------- # read image # -------------------------------- L_folder = os.path.join(testsets, testset_L, 'X4') E_folder = os.path.join(testsets, testset_L+'_results') util.mkdir(E_folder) # record PSNR, runtime test_results = OrderedDict() test_results['runtime'] = [] logger.info(L_folder) logger.info(E_folder) idx = 0 start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) img_SR = [] for img in util.get_image_paths(L_folder): # -------------------------------- # (1) img_L # -------------------------------- idx += 1 img_name, ext = os.path.splitext(os.path.basename(img)) logger.info('{:->4d}--> {:>10s}'.format(idx, img_name+ext)) img_L = util.imread_uint(img, n_channels=3) img_L = util.uint2tensor4(img_L) img_L = img_L.to(device) start.record() img_E = model(img_L) end.record() torch.cuda.synchronize() test_results['runtime'].append(start.elapsed_time(end)) # milliseconds # -------------------------------- # (2) img_E # -------------------------------- img_E = util.tensor2uint(img_E) img_SR.append(img_E) # -------------------------------- # (3) save results # -------------------------------- #util.imsave(img_E, os.path.join(E_folder, img_name+ext)) ave_runtime = sum(test_results['runtime']) / len(test_results['runtime']) / 1000.0 logger.info('------> Average runtime of ({}) is : {:.6f} seconds'.format(L_folder, ave_runtime)) # -------------------------------- # (4) calculate psnr # -------------------------------- '''
def main(): utils_logger.logger_info('efficientsr_challenge', log_path='efficientsr_challenge.log') logger = logging.getLogger('efficientsr_challenge') # print(torch.__version__) # pytorch version # print(torch.version.cuda) # cuda version # print(torch.backends.cudnn.version()) # cudnn version # -------------------------------- # basic settings # -------------------------------- model_names = ['msrresnet', 'imdn'] model_id = 1 # set the model name model_name = model_names[model_id] logger.info('{:>16s} : {:s}'.format('Model Name', model_name)) testsets = 'testsets' # set path of testsets testset_L = 'DIV2K_valid_LR' # set current testing dataset; 'DIV2K_test_LR' testset_L = 'set12' save_results = True print_modelsummary = True # set False when calculating `Max Memery` and `Runtime` torch.cuda.set_device(0) # set GPU ID logger.info('{:>16s} : {:<d}'.format('GPU ID', torch.cuda.current_device())) torch.cuda.empty_cache() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # -------------------------------- # define network and load model # -------------------------------- if model_name == 'msrresnet': from models.network_msrresnet import MSRResNet1 as net model = net(in_nc=3, out_nc=3, nc=64, nb=16, upscale=4) # define network model_path = os.path.join('model_zoo', 'msrresnet_x4_psnr.pth') # set model path elif model_name == 'imdn': from models.network_imdn import IMDN as net model = net(in_nc=3, out_nc=3, nc=64, nb=8, upscale=4, act_mode='L', upsample_mode='pixelshuffle') # define network model_path = os.path.join('model_zoo', 'imdn_x4.pth') # set model path model.load_state_dict(torch.load(model_path), strict=True) model.eval() for k, v in model.named_parameters(): v.requires_grad = False model = model.to(device) # -------------------------------- # print model summary # -------------------------------- if print_modelsummary: from utils.utils_modelsummary import get_model_activation, get_model_flops input_dim = (3, 256, 256) # set the input dimension activations, num_conv2d = get_model_activation(model, input_dim) logger.info('{:>16s} : {:<.4f} [M]'.format('#Activations', activations / 10**6)) logger.info('{:>16s} : {:<d}'.format('#Conv2d', num_conv2d)) flops = get_model_flops(model, input_dim, False) logger.info('{:>16s} : {:<.4f} [G]'.format('FLOPs', flops / 10**9)) num_parameters = sum(map(lambda x: x.numel(), model.parameters())) logger.info('{:>16s} : {:<.4f} [M]'.format('#Params', num_parameters / 10**6)) # -------------------------------- # read image # -------------------------------- L_path = os.path.join(testsets, testset_L) E_path = os.path.join(testsets, testset_L + '_' + model_name) util.mkdir(E_path) # record runtime test_results = OrderedDict() test_results['runtime'] = [] logger.info('{:>16s} : {:s}'.format('Input Path', L_path)) logger.info('{:>16s} : {:s}'.format('Output Path', E_path)) idx = 0 start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) for img in util.get_image_paths(L_path): # -------------------------------- # (1) img_L # -------------------------------- idx += 1 img_name, ext = os.path.splitext(os.path.basename(img)) logger.info('{:->4d}--> {:>10s}'.format(idx, img_name + ext)) img_L = util.imread_uint(img, n_channels=3) img_L = util.uint2tensor4(img_L) torch.cuda.empty_cache() img_L = img_L.to(device) start.record() img_E = model(img_L) # logger.info('{:>16s} : {:<.3f} [M]'.format('Max Memery', torch.cuda.max_memory_allocated(torch.cuda.current_device())/1024**2)) # Memery end.record() torch.cuda.synchronize() test_results['runtime'].append(start.elapsed_time(end)) # milliseconds # torch.cuda.synchronize() # start = time.time() # img_E = model(img_L) # torch.cuda.synchronize() # end = time.time() # test_results['runtime'].append(end-start) # seconds # -------------------------------- # (2) img_E # -------------------------------- img_E = util.tensor2uint(img_E) if save_results: util.imsave(img_E, os.path.join(E_path, img_name + ext)) ave_runtime = sum(test_results['runtime']) / len( test_results['runtime']) / 1000.0 logger.info('------> Average runtime of ({}) is : {:.6f} seconds'.format( L_path, ave_runtime))
def main(): utils_logger.logger_info('blind_sr_log', log_path='blind_sr_log.log') logger = logging.getLogger('blind_sr_log') # print(torch.__version__) # pytorch version # print(torch.version.cuda) # cuda version # print(torch.backends.cudnn.version()) # cudnn version testsets = 'testsets' # fixed, set path of testsets testset_Ls = ['RealSRSet'] # ['RealSRSet','DPED'] model_names = ['RRDB','ESRGAN','FSSR_DPED','FSSR_JPEG','RealSR_DPED','RealSR_JPEG'] model_names = ['BSRGAN'] save_results = True sf = 4 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') for model_name in model_names: model_path = os.path.join('model_zoo', model_name+'.pth') # set model path logger.info('{:>16s} : {:s}'.format('Model Name', model_name)) # torch.cuda.set_device(0) # set GPU ID logger.info('{:>16s} : {:<d}'.format('GPU ID', torch.cuda.current_device())) torch.cuda.empty_cache() # -------------------------------- # define network and load model # -------------------------------- model = net(in_nc=3, out_nc=3, nf=64, nb=23, gc=32) # define network # model_old = torch.load(model_path) # state_dict = model.state_dict() # for ((key, param),(key2, param2)) in zip(model_old.items(), state_dict.items()): # state_dict[key2] = param # model.load_state_dict(state_dict, strict=True) model.load_state_dict(torch.load(model_path), strict=True) model.eval() for k, v in model.named_parameters(): v.requires_grad = False model = model.to(device) torch.cuda.empty_cache() for testset_L in testset_Ls: L_path = os.path.join(testsets, testset_L) #E_path = os.path.join(testsets, testset_L+'_'+model_name) E_path = os.path.join(testsets, testset_L+'_results_x'+str(sf)) util.mkdir(E_path) logger.info('{:>16s} : {:s}'.format('Input Path', L_path)) logger.info('{:>16s} : {:s}'.format('Output Path', E_path)) idx = 0 for img in util.get_image_paths(L_path): # -------------------------------- # (1) img_L # -------------------------------- idx += 1 img_name, ext = os.path.splitext(os.path.basename(img)) logger.info('{:->4d} --> {:<s} --> x{:<d}--> {:<s}'.format(idx, model_name, sf, img_name+ext)) img_L = util.imread_uint(img, n_channels=3) img_L = util.uint2tensor4(img_L) img_L = img_L.to(device) # -------------------------------- # (2) inference # -------------------------------- img_E = model(img_L) # -------------------------------- # (3) img_E # -------------------------------- img_E = util.tensor2uint(img_E) if save_results: util.imsave(img_E, os.path.join(E_path, img_name+'_'+model_name+'.png'))
def main(): # ---------------------------------------- # Preparation # ---------------------------------------- if noise_level_model == -1: model_name = 'srmdnf_x' + str(sf) else: model_name = 'srmd_x' + str(sf) model_path = os.path.join(model_pool, model_name+'.pth') in_nc = 18 if 'nf' in model_name else 19 # ---------------------------------------- # L_path, E_path, H_path # ---------------------------------------- L_path = sources # L_path, for Low-quality images E_path = results # E_path, for Estimated images if not os.path.splitext(E_path)[1]: util.mkdir(E_path) device = torch.device(using_device) # ---------------------------------------- # load model # ---------------------------------------- from utils.network_srmd import SRMD as net model = net(in_nc=in_nc, out_nc=n_channels, nc=nc, nb=nb, upscale=sf, act_mode='R', upsample_mode='pixelshuffle') model.load_state_dict(torch.load(model_path), strict=False) model.eval() for _, v in model.named_parameters(): v.requires_grad = False model = model.to(device) if os.path.isfile(L_path): L_paths = [L_path] else: L_paths = util.get_image_paths(L_path) # ---------------------------------------- # kernel and PCA reduced feature # ---------------------------------------- # Gaussian kernel, delta kernel 0.01 kernel = utils_deblur.fspecial('gaussian', 15, 0.01) P = loadmat(srmd_pca_path)['P'] degradation_vector = np.dot(P, np.reshape(kernel, (-1), order="F")) if 'nf' not in model_name: # noise-free SR degradation_vector = np.append( degradation_vector, noise_level_model/255.) degradation_vector = torch.from_numpy( degradation_vector).view(1, -1, 1, 1).float() for _, img in enumerate(L_paths): img_name, _ = os.path.splitext(os.path.basename(img)) try: # ------------------------------------ # (1) img_L # ------------------------------------ img_L, alpha = util.imread_uint_alpha(img, n_channels=n_channels) # Bicubic to handle alpha channel if the intended picture is supposed to have. if not alpha is None and picture_format == "png": alpha = util.uint2tensor4(alpha) alpha = torch.nn.functional.interpolate( alpha, scale_factor=sf, mode='bicubic', align_corners=False) alpha = alpha.to(device) alpha = torch.clamp(alpha, 0, 255) alpha = util.tensor2uint(alpha) img_L = util.uint2tensor4(img_L) degradation_map = degradation_vector.repeat( 1, 1, img_L.size(-2), img_L.size(-1)) img_L = torch.cat((img_L, degradation_map), dim=1) img_L = img_L.to(device) # ------------------------------------ # (2) img_E # ------------------------------------ if not x8: img_E = model(img_L) else: img_E = utils_model.test_mode(model, img_L, mode=3, sf=sf) img_E = util.tensor2uint(img_E) if not alpha is None and picture_format == "png": alpha = alpha.reshape((alpha.shape[0], alpha.shape[1], 1)) img_E = np.concatenate((img_E, alpha), axis=2) elif not alpha is None: print("Warning! You lost your alpha channel for this picture!") # ------------------------------------ # save results # ------------------------------------ if os.path.splitext(E_path)[1]: util.imsave(img_E, E_path) else: util.imsave(img_E, os.path.join( E_path, img_name+'.' + picture_format)) print(os.path.basename(img) + " successfully saved to disk!") except Exception: traceback.print_exc() print(os.path.basename(img) + " failed!")
def main(model=None, model_path=None): utils_logger.logger_info('AIM-track', log_path=os.path.join(model_path, 'AIM-track.log')) logger = logging.getLogger('AIM-track') # -------------------------------- # basic settings # -------------------------------- testsets = 'DIV2K' # DIV2K root path testset_L = 'DIV2K_test_LR_bicubic' # test image folder name torch.cuda.current_device() torch.cuda.empty_cache() torch.backends.cudnn.benchmark = True device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # -------------------------------- # load model # -------------------------------- # model_path = os.path.join('MSRResNetx4_model', 'MSRResNetx4.pth') if model is None: model_path = 'MSRResNetx4_model' model = MSRResNet(in_nc=3, out_nc=3, nf=64, nb=16, upscale=4) model.load_state_dict(torch.load(os.path.join(model_path, 'model.pth')), strict=True) model.eval() """ for k, v in model.named_parameters(): v.requires_grad = False model = model.to(device) """ # number of parameters number_parameters = sum(map(lambda x: x.numel(), model.parameters())) logger.info('Params number: {}'.format(number_parameters)) print('Params number: {}'.format(number_parameters)) # -------------------------------- # read image # -------------------------------- L_folder = os.path.join(testsets, testset_L) assert os.path.isdir(L_folder) # check the test images path E_folder = os.path.join(model_path, 'results') util.mkdir(E_folder) # record PSNR, runtime test_results = OrderedDict() test_results['runtime'] = [] logger.info(L_folder) logger.info(E_folder) idx = 0 start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) for img in util.get_image_paths(L_folder): # -------------------------------- # (1) img_L # -------------------------------- idx += 1 img_name, ext = os.path.splitext(os.path.basename(img)) logger.info('{:->4d}--> {:>10s}'.format(idx, img_name + ext)) img_L = util.imread_uint(img, n_channels=3) img_L = util.uint2tensor4(img_L) img_L = img_L.to(device) start.record() img_E = model(img_L) end.record() torch.cuda.synchronize() test_results['runtime'].append(start.elapsed_time(end)) # milliseconds # torch.cuda.synchronize() # start = time.time() # img_E = model(img_L) # torch.cuda.synchronize() # end = time.time() # test_results['runtime'].append(end-start) # seconds # -------------------------------- # (2) img_E # -------------------------------- img_E = util.tensor2uint(img_E) util.imsave(img_E, os.path.join(E_folder, img_name + ext)) ave_runtime = sum(test_results['runtime']) / len( test_results['runtime']) / 1000.0 logger.info('------> Average runtime of ({}) is : {:.6f} seconds'.format( L_folder, ave_runtime)) print('------> Average runtime of ({}) is : {:.6f} seconds'.format( L_folder, ave_runtime))