def test_ffdnet(**args): r"""Denoises an input image with FFDNet """ # Init logger logger = init_logger_ipol() # Check if input exists and if it is RGB try: rgb_den = is_rgb(args['input']) except: raise Exception('Could not open the input image') # Measure runtime start_t = time.time() # Open image as a CxHxW torch.Tensor if rgb_den: in_ch = 3 model_fn = 'net_rgb.pth' imorig = Image.open(args['input']) imorig = np.array(imorig, dtype=np.float32).transpose(2, 0, 1) else: # from HxWxC to CxHxW grayscale image (C=1) in_ch = 1 model_fn = 'models/net_gray.pth' # imorig = cv2.imread(args['input'], cv2.IMREAD_GRAYSCALE) imorig = np.expand_dims(imorig, 0) imorig = np.expand_dims(imorig, 0) # Handle odd sizes expanded_h = False expanded_w = False sh_im = imorig.shape if sh_im[2] % 2 == 1: expanded_h = True imorig = np.concatenate((imorig, \ imorig[:, :, -1, :][:, :, np.newaxis, :]), axis=2) if sh_im[3] % 2 == 1: expanded_w = True imorig = np.concatenate((imorig, \ imorig[:, :, :, -1][:, :, :, np.newaxis]), axis=3) imorig = normalize(imorig) imorig = torch.Tensor(imorig) # Absolute path to model file model_fn = os.path.join(os.path.abspath(os.path.dirname(__file__)), \ model_fn) # Create model print('Loading model ...\n') net = FFDNet(num_input_channels=in_ch) # Load saved weights if args['cuda']: state_dict = torch.load(model_fn) #device_ids = [0,1,2,3] #model = nn.DataParallel(net, device_ids=device_ids).cuda() #state_dict = remove_dataparallel_wrapper(state_dict) model = net else: state_dict = torch.load(model_fn, map_location='cpu') # CPU mode: remove the DataParallel wrapper state_dict = remove_dataparallel_wrapper(state_dict) model = net model.load_state_dict(state_dict) # Sets the model in evaluation mode (e.g. it removes BN) model.eval() # Sets data type according to CPU or GPU modes if args['cuda']: dtype = torch.cuda.FloatTensor else: dtype = torch.FloatTensor # Test mode with torch.no_grad(): # PyTorch v0.4.0 imorig = Variable(imorig.type(dtype)) nsigma = Variable(torch.FloatTensor([args['noise_sigma']]).type(dtype)) # # Measure runtime # start_t = time.time() # Estimate noise and subtract it to the input image im_noise_estim = model(imorig, nsigma) stop_t = time.time() # log time if rgb_den: print("### RGB denoising ###") else: print("### Grayscale denoising ###") print("\tRuntime {0:0.4f}s".format(stop_t - start_t)) # Save noises noise = variable_to_numpy(imorig.to(3) - im_noise_estim).transpose(1, 2, 0) filename = args['input'].split('/')[-1].split('.')[0] if args['save_path']: sio.savemat( './output_noise/' + args['save_path'] + '/' + filename + '.mat', {'Noisex': noise}) else: sio.savemat('./output_noise/' + filename + '.mat', {'Noisex': noise})
def test_ffdnet(**args): r"""Denoises an input image with FFDNet """ # Init logger logger = init_logger_ipol() # Check if input exists and if it is RGB try: rgb_den = is_rgb(args['input']) except: raise Exception('Could not open the input image') # Open image as a CxHxW torch.Tensor if rgb_den: in_ch = 3 model_fn = 'models/net_rgb.pth' imorig = cv2.imread(args['input']) # from HxWxC to CxHxW, RGB image imorig = (cv2.cvtColor(imorig, cv2.COLOR_BGR2RGB)).transpose(2, 0, 1) else: # from HxWxC to CxHxW grayscale image (C=1) in_ch = 1 model_fn = 'logs/net.pth' imorig = cv2.imread(args['input'], cv2.IMREAD_GRAYSCALE) imorig_copy = imorig.copy() imorig = np.expand_dims(imorig, 0) imorig = np.expand_dims(imorig, 0) # Handle odd sizes expanded_h = False expanded_w = False sh_im = imorig.shape if sh_im[2] % 2 == 1: expanded_h = True imorig = np.concatenate((imorig, \ imorig[:, :, -1, :][:, :, np.newaxis, :]), axis=2) if sh_im[3] % 2 == 1: expanded_w = True imorig = np.concatenate((imorig, \ imorig[:, :, :, -1][:, :, :, np.newaxis]), axis=3) imorig = normalize(imorig) imorig = torch.Tensor(imorig) # Absolute path to model file model_fn = os.path.join(os.path.abspath(os.path.dirname(__file__)), \ model_fn) # Create model print('Loading model ...\n') net = FFDNet(num_input_channels=in_ch) # Load saved weights if args['cuda']: state_dict = torch.load(model_fn) device_ids = [0] model = nn.DataParallel(net, device_ids=device_ids).cuda() else: state_dict = torch.load(model_fn, map_location='cpu') # CPU mode: remove the DataParallel wrapper state_dict = remove_dataparallel_wrapper(state_dict) model = net model.load_state_dict(state_dict) # Sets the model in evaluation mode (e.g. it removes BN) model.eval() # Sets data type according to CPU or GPU modes if args['cuda']: dtype = torch.cuda.FloatTensor else: dtype = torch.FloatTensor # Add noise if args['add_noise']: noise = torch.FloatTensor(imorig.size()).\ normal_(mean=0, std=args['noise_sigma']) imnoisy = imorig + noise else: imnoisy = imorig.clone() # Test mode with torch.no_grad(): # PyTorch v0.4.0 imorig, imnoisy = Variable(imorig.type(dtype)), \ Variable(imnoisy.type(dtype)) nsigma = Variable(torch.FloatTensor([args['noise_sigma']]).type(dtype)) # Measure runtime start_t = time.time() # Estimate noise and subtract it to the input image im_noise_estim = model(imnoisy, nsigma) outim = torch.clamp(imnoisy - im_noise_estim, 0., 1.) stop_t = time.time() if expanded_h: imorig = imorig[:, :, :-1, :] outim = outim[:, :, :-1, :] imnoisy = imnoisy[:, :, :-1, :] if expanded_w: imorig = imorig[:, :, :, :-1] outim = outim[:, :, :, :-1] imnoisy = imnoisy[:, :, :, :-1] # Compute PSNR and log it if rgb_den: print("### RGB denoising ###") else: print("### Grayscale denoising ###") if args['add_noise']: psnr = batch_psnr(outim, imorig, 1.) psnr_noisy = batch_psnr(imnoisy, imorig, 1.) print("----------PSNR noisy {0:0.2f}dB".format(psnr_noisy)) print("----------PSNR denoised {0:0.2f}dB".format(psnr)) else: logger.info("\tNo noise was added, cannot compute PSNR") print("----------Runtime {0:0.4f}s".format(stop_t - start_t)) # Compute difference diffout = 2 * (outim - imorig) + .5 diffnoise = 2 * (imnoisy - imorig) + .5 # Save images if not args['dont_save_results']: noisyimg = variable_to_cv2_image(imnoisy) outimg = variable_to_cv2_image(outim) cv2.imwrite( "bfffd/noisy-" + str(int(args['noise_sigma'] * 255)) + '-' + args['input'], noisyimg) cv2.imwrite( "bfffd/ffdnet-" + str(int(args['noise_sigma'] * 255)) + '-' + args['input'], outimg) if args['add_noise']: cv2.imwrite("noisy_diff.png", variable_to_cv2_image(diffnoise)) cv2.imwrite("ffdnet_diff.png", variable_to_cv2_image(diffout)) (score, diff) = compare_ssim(noisyimg, imorig_copy, full=True) (score2, diff) = compare_ssim(outimg, imorig_copy, full=True) print("----------Noisy ssim: {0:0.4f}".format(score)) print("----------Denoisy ssim: {0:0.4f}".format(score2))
def clean_noise(request): input = '//172.30.20.157/235-project' + request.args.get('filepath') try: rgb_den = is_rgb(input) except: raise Exception('Could not open the input image') # Open image as a CxHxW torch.Tensor if rgb_den: imorig = cv2.imdecode(np.fromfile(input, dtype=np.uint8), -1) # from HxWxC to CxHxW, RGB image imorig = (cv2.cvtColor(imorig, cv2.COLOR_BGR2RGB)).transpose(2, 0, 1) else: # from HxWxC to CxHxW grayscale image (C=1) imorig = np.array(Image.open(input).convert('RGB')).transpose(2, 0, 1) img_shape = imorig.shape h = 200 num_h = img_shape[1] // h ori_h = h * num_h crop_h = img_shape[1] - ori_h w = 200 num_w = img_shape[2] // w ori_w = w * num_w crop_w = img_shape[2] - ori_w images = [] for i in range(num_h + 1): for j in range(num_w + 1): if i == num_h and j == num_w: img1 = imorig[:, i * int(h):i * int(h) + crop_h, j * int(w):j * int(w) + crop_w] if i == num_h: img1 = imorig[:, i * int(h):i * int(h) + crop_h, j * int(w):(j + 1) * int(w)] if j == num_w: img1 = imorig[:, i * int(h):(i + 1) * int(h), j * int(w):j * int(w) + crop_w] img1 = imorig[:, i * int(h):(i + 1) * int(h), j * int(w):(j + 1) * int(w)] images.append([i, j, img1]) denoise_images = [] for index, i in enumerate(images): denoise_images.append([i[0], i[1], noise(i[2])]) new_img = np.zeros((img_shape[1], img_shape[2], 3)) for img in denoise_images: i, j = img[:2] if i == num_h and j == num_w: img[2] = cv2.resize( img[2], (new_img[int(i) * h:(int(i) + 1) * h + crop_h, int(j) * w:(int(j) + 1) * w + crop_w].shape[-2::-1])) new_img[i * int(h):i * int(h) + crop_h, j * int(w):j * int(w) + crop_w] = img[2] if i == num_h: img[2] = cv2.resize( img[2], (new_img[i * int(h):i * int(h) + crop_h, j * int(w):(j + 1) * int(w)].shape[-2::-1])) new_img[i * int(h):i * int(h) + crop_h, j * int(w):(j + 1) * int(w)] = img[2] if j == num_w: img[2] = cv2.resize( img[2], (new_img[i * int(h):(i + 1) * int(h), j * int(w):j * int(w) + crop_w].shape[-2::-1])) new_img[i * int(h):(i + 1) * int(h), j * int(w):j * int(w) + crop_w] = img[2] new_img[int(i) * h:(int(i) + 1) * h, int(j) * w:(int(j) + 1) * w] = img[2] new_img = evaluate(new_img) image_path = request.args.get('filepath').split('/') image_path[2] = 'clean' new_image_path = '//172.30.20.157/project' + '/'.join(image_path[:-1]) if not os.path.exists(new_image_path): os.makedirs(new_image_path) new_image_path = new_image_path + '/' + image_path[-1] Image.fromarray(new_img.astype(np.uint8)).save(new_image_path, dpi=(300., 300.), quality=60) return response.json({input: new_image_path})