Fix image saving from float32 to uint8

Iceage7 · Oct 29, 2020 · 58eaf63 · 58eaf63
1 parent a668d27
commit 58eaf63
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 20 deletions.
diff --git a/test_denoising_dnd.py b/test_denoising_dnd.py
@@ -22,6 +22,7 @@
 import utils
 import lycon
 from utils.bundle_submissions import bundle_submissions_srgb_v1
+from skimage import img_as_ubyte
 
 
 parser = argparse.ArgumentParser(description='RGB denoising evaluation on DND dataset')
@@ -74,9 +75,8 @@
 
         if args.save_images:
             for batch in range(len(rgb_noisy)):
-                #temp = np.concatenate((rgb_noisy[batch]*255, rgb_restored[batch]*255),axis=1)
-                denoised_img = rgb_restored[batch]*255
-                lycon.save(args.result_dir + 'png/'+ filenames[batch][:-4] + '.png', denoised_img.astype(np.uint8))
+                denoised_img = img_as_ubyte(rgb_restored[batch])
+                lycon.save(args.result_dir + 'png/'+ filenames[batch][:-4] + '.png', denoised_img)
                 save_file = os.path.join(args.result_dir+ 'matfile/', filenames[batch][:-4] +'.mat')
                 sio.savemat(save_file, {'Idenoised_crop': np.float32(rgb_restored[batch])})
 

diff --git a/test_denoising_sidd.py b/test_denoising_sidd.py
@@ -23,6 +23,7 @@
 from dataloaders.data_rgb import get_validation_data
 import utils
 import lycon
+from skimage import img_as_ubyte
 
 parser = argparse.ArgumentParser(description='RGB denoising evaluation on the validation set of SIDD')
 parser.add_argument('--input_dir', default='./datasets/sidd/',
@@ -77,9 +78,8 @@
 
         if args.save_images:
             for batch in range(len(rgb_gt)):
-                #temp = np.concatenate((rgb_noisy[batch]*255, rgb_restored[batch]*255, rgb_gt[batch]*255),axis=1)
-                denoised_img = rgb_restored[batch]*255
-                lycon.save(args.result_dir + filenames[batch][:-4] + '.png', denoised_img.astype(np.uint8))
+                denoised_img = img_as_ubyte(rgb_restored[batch])
+                lycon.save(args.result_dir + filenames[batch][:-4] + '.png', denoised_img)
 
 psnr_val_rgb = sum(psnr_val_rgb)/len(psnr_val_rgb)
 print("PSNR: %.2f " %(psnr_val_rgb))

diff --git a/test_enhancement.py b/test_enhancement.py
@@ -20,6 +20,7 @@
 from dataloaders.data_rgb import get_validation_data
 import utils
 import lycon
+from skimage import img_as_ubyte
 
 parser = argparse.ArgumentParser(description='Image Enhancement using MIRNet')
 # parser.add_argument('--input_dir', default='./datasets/lol/', type=str, help='Directory of validation images')
@@ -45,8 +46,6 @@
 test_dataset = get_validation_data(args.input_dir)
 test_loader = DataLoader(dataset=test_dataset, batch_size=args.bs, shuffle=False, num_workers=8, drop_last=False)
 
-
-
 model_restoration = MIRNet()
 
 utils.load_checkpoint(model_restoration,args.weights)
@@ -58,7 +57,6 @@
 
 model_restoration.eval()
 
-
 with torch.no_grad():
     psnr_val_rgb = []
     for ii, data_test in enumerate(tqdm(test_loader), 0):
@@ -76,9 +74,8 @@
 
         if args.save_images:
             for batch in range(len(rgb_gt)):
-                #temp = np.concatenate((rgb_noisy[batch]*255, rgb_restored[batch]*255, rgb_gt[batch]*255),axis=1)
-                denoised_img = rgb_restored[batch]*255
-                lycon.save(args.result_dir + filenames[batch][:-4] + '.png', denoised_img.astype(np.uint8))
+                enhanced_img = img_as_ubyte(rgb_restored[batch])
+                lycon.save(args.result_dir + filenames[batch][:-4] + '.png', enhanced_img)
 
 psnr_val_rgb = sum(psnr_val_rgb)/len(psnr_val_rgb)
 print("PSNR: %.2f " %(psnr_val_rgb))

diff --git a/test_super_resolution.py b/test_super_resolution.py
@@ -20,7 +20,7 @@
 from dataloaders.data_rgb import get_test_data_SR
 import utils
 import lycon
-
+from skimage import img_as_ubyte
 
 parser = argparse.ArgumentParser(description='Super-resolve images of RealSR dataset')
 parser.add_argument('--input_dir', default='./datasets/realSR/x',
@@ -36,7 +36,6 @@
 
 args = parser.parse_args()
 
-
 os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
 
@@ -48,8 +47,6 @@
 
 test_loader = DataLoader(dataset=test_dataset, batch_size=args.bs, shuffle=False, num_workers=8, drop_last=False)
 
-
-
 model_restoration = MIRNet()
 
 weights = args.weights+args.scale+'.pth'
@@ -62,7 +59,6 @@
 
 model_restoration.eval()
 
-
 with torch.no_grad():
     for ii, data_test in enumerate(tqdm(test_loader), 0):
         LR_img = data_test[0].cuda()
@@ -75,6 +71,5 @@
 
         if args.save_images:
             for batch in range(len(LR_img)):
-                #temp = np.concatenate((LR_img[batch]*255, rgb_restored[batch]*255),axis=1)
-                denoised_img = rgb_restored[batch]*255
-                lycon.save(os.path.join(output_dir, filenames[batch][:-4]+'.png'), denoised_img.astype(np.uint8))
+                sr_img = img_as_ubyte(rgb_restored[batch])
+                lycon.save(os.path.join(output_dir, filenames[batch][:-4]+'.png'), sr_img)