Fix image saving from float32 to uint8

Author: adityac8

generate_raw_data.py

@@ -21,6 +21,7 @@
 from utils.noise_sampling import random_noise_levels_dnd, random_noise_levels_sidd, add_noise
 import utils
 import lycon
+from skimage import img_as_ubyte
 
 parser = argparse.ArgumentParser(description='RGB2RAW Network: From clean RGB images, generate {RAW_clean, RAW_noisy} pairs')
 parser.add_argument('--input_dir', default='./datasets/sample_rgb_images/',
@@ -34,7 +35,6 @@
 
 args = parser.parse_args()
 
-
 os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
 
@@ -45,8 +45,6 @@
 test_dataset = get_rgb_data(args.input_dir)
 test_loader = DataLoader(dataset=test_dataset, batch_size=4, shuffle=False, num_workers=2, drop_last=False)
 
-
-
 model_rgb2raw = Rgb2Raw()
 
 utils.load_checkpoint(model_rgb2raw,args.weights)
@@ -58,7 +56,6 @@
 
 model_rgb2raw.eval()
 
-
 with torch.no_grad():
     for ii, data in enumerate(tqdm(test_loader), 0):
         rgb_gt    = data[0].cuda()
@@ -68,8 +65,6 @@
         ## Convert clean rgb image to clean raw image
         raw_gt = model_rgb2raw(rgb_gt)       ## raw_gt is in RGGB format
         raw_gt = torch.clamp(raw_gt,0,1)
-        #raw_noisy = torch.zeros_like(raw_gt).cuda()
-        #variance = torch.zeros_like(raw_gt).cuda()
 
         ########## Add noise to clean raw images ##########
         for j in range(raw_gt.shape[0]):
@@ -79,29 +74,26 @@
             raw_noisy = add_noise(raw_gt[j], shot_noise, read_noise, use_cuda=True)
             raw_noisy = torch.clamp(raw_noisy,0,1)  ### CLIP NOISE
             variance = shot_noise * raw_noisy + read_noise
+
             #### Unpadding and saving
             clean_packed = raw_gt[j]
             clean_packed = clean_packed[:,padh[j]//2:-padh[j]//2,padw[j]//2:-padw[j]//2]   ## RGGB channels  (4 x H/2 x W/2)
-            clean_unpacked = utils.unpack_raw(clean_packed.unsqueeze(0))                                  ## Rearrange RGGB channels into Bayer pattern
+            clean_unpacked = utils.unpack_raw(clean_packed.unsqueeze(0))                   ## Rearrange RGGB channels into Bayer pattern
             clean_unpacked = clean_unpacked.squeeze().cpu().detach().numpy()
-            lycon.save(args.result_dir+'png/clean/'+filename[:-4]+'.png',(clean_unpacked*255).astype(np.uint8))
+            lycon.save(args.result_dir+'png/clean/'+filename[:-4]+'.png',img_as_ubyte(clean_unpacked))
 
             noisy_packed = raw_noisy
             noisy_packed = noisy_packed[:,padh[j]//2:-padh[j]//2,padw[j]//2:-padw[j]//2]   ## RGGB channels
-            noisy_unpacked = utils.unpack_raw(noisy_packed.unsqueeze(0))                                  ## Rearrange RGGB channels into Bayer pattern
+            noisy_unpacked = utils.unpack_raw(noisy_packed.unsqueeze(0))                   ## Rearrange RGGB channels into Bayer pattern
             noisy_unpacked = noisy_unpacked.squeeze().cpu().detach().numpy()
-            lycon.save(args.result_dir+'png/noisy/'+filename[:-4]+'.png',(noisy_unpacked*255).astype(np.uint8))
-
+            lycon.save(args.result_dir+'png/noisy/'+filename[:-4]+'.png',img_as_ubyte(noisy_unpacked))
 
             variance_packed = variance[:,padh[j]//2:-padh[j]//2,padw[j]//2:-padw[j]//2]   ## RGGB channels
 
-            #import pdb;pdb.set_trace()
             dict_ = {}
             dict_['clean'] = clean_packed.cpu().detach().numpy()       ## (4 x H/2 x W/2)
             dict_['noisy'] = noisy_packed.cpu().detach().numpy()       ## (4 x H/2 x W/2)
-            dict_['variance'] = variance_packed.cpu().detach().numpy()   ## (4 x H/2 x W/2)
+            dict_['variance'] = variance_packed.cpu().detach().numpy() ## (4 x H/2 x W/2)
             dict_['shot_noise'] = shot_noise.cpu().detach().numpy()
             dict_['read_noise'] = read_noise.cpu().detach().numpy()
             utils.save_dict(dict_, args.result_dir+'pkl/'+filename[:-4]+'.pkl')
-
-

generate_rgb_data.py

@@ -21,6 +21,7 @@
 from utils.noise_sampling import random_noise_levels_dnd, random_noise_levels_sidd, add_noise
 import utils
 import lycon
+from skimage import img_as_ubyte
 
 parser = argparse.ArgumentParser(description='From clean RGB images, generate {RGB_clean, RGB_noisy} pairs')
 parser.add_argument('--input_dir', default='./datasets/sample_rgb_images/',
@@ -35,7 +36,6 @@
 
 args = parser.parse_args()
 
-
 os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
 
@@ -45,16 +45,13 @@
 test_dataset = get_rgb_data(args.input_dir)
 test_loader = DataLoader(dataset=test_dataset, batch_size=4, shuffle=False, num_workers=2, drop_last=False)
 
-
-
 model_rgb2raw = Rgb2Raw()
 model_ccm     = CCM()
 model_raw2rgb = Raw2Rgb()
 
 utils.load_checkpoint(model_rgb2raw,args.weights_rgb2raw)
 utils.load_checkpoint(model_ccm,args.weights_ccm)
 utils.load_checkpoint(model_raw2rgb,args.weights_raw2rgb)
-#print("===>Testing using weights: ", args.weights)
 
 model_rgb2raw.cuda()
 model_ccm.cuda()
@@ -68,13 +65,13 @@
 model_ccm.eval()
 model_raw2rgb.eval()
 
-
 with torch.no_grad():
     for ii, data in enumerate(tqdm(test_loader), 0):
         rgb_gt    = data[0].cuda()
         filenames = data[1]
         padh = data[2]
         padw = data[3]
+
         ## Convert clean rgb image to clean raw image
         raw_gt = model_rgb2raw(rgb_gt)       ## raw_gt is in RGGB format
         raw_gt = torch.clamp(raw_gt,0,1)
@@ -87,7 +84,6 @@
             raw_noisy = add_noise(raw_gt[j], shot_noise, read_noise, use_cuda=True)
             raw_noisy = torch.clamp(raw_noisy,0,1)  ### CLIP NOISE
 
-
             #### Convert raw noisy to rgb noisy ####
             ccm_tensor = model_ccm(rgb_gt[j].unsqueeze(0))
             rgb_noisy = model_raw2rgb(raw_noisy.unsqueeze(0),ccm_tensor) 
@@ -96,13 +92,10 @@
             rgb_noisy = rgb_noisy.permute(0, 2, 3, 1).squeeze().cpu().detach().numpy()
 
             rgb_clean = rgb_gt[j].permute(1,2,0).cpu().detach().numpy()
+
             ## Unpadding
             rgb_clean = rgb_clean[padh[j]:-padh[j],padw[j]:-padw[j],:]   
             rgb_noisy = rgb_noisy[padh[j]:-padh[j],padw[j]:-padw[j],:]   
-            # import pdb;pdb.set_trace()
-
-            lycon.save(args.result_dir+'clean/'+filename[:-4]+'.png',(rgb_clean*255).astype(np.uint8))
-            lycon.save(args.result_dir+'noisy/'+filename[:-4]+'.png',(rgb_noisy*255).astype(np.uint8))
-
-           
 
+            lycon.save(args.result_dir+'clean/'+filename[:-4]+'.png',img_as_ubyte(rgb_clean))
+            lycon.save(args.result_dir+'noisy/'+filename[:-4]+'.png',img_as_ubyte(rgb_noisy))

test_dnd_raw.py

@@ -19,13 +19,11 @@
 import h5py
 import numpy as np
 from networks.denoising_raw import DenoiseNet
-# from dataloaders.data_raw import get_test_data 
 
 import utils
 import lycon
 from utils.bundle_submissions import bundle_submissions_raw
-
-
+from skimage import img_as_ubyte
 
 parser = argparse.ArgumentParser(description='RAW denoising evaluation on the DND dataset')
 parser.add_argument('--input_dir', default='./datasets/dnd/dnd_raw/',
@@ -83,7 +81,6 @@ def denoiser(raw_noisy,variance):
     img = h5py.File(filename, 'r')
     noisy = np.float32(np.array(img['Inoisy']).T)
 
-    #
     # Loads raw Bayer color pattern.
     bayer_pattern = np.asarray(info[info['camera'][0][i]]['pattern']).tolist()
 
@@ -150,9 +147,9 @@ def denoiser(raw_noisy,variance):
         sio.savemat(save_file, {'Idenoised_crop': Idenoised_crop})
 
         if args.save_images:
-            denoised_img = Idenoised_crop*255
+            denoised_img = img_as_ubyte(Idenoised_crop)
             save_file = os.path.join(args.result_dir+ 'png/', '%04d_%02d.png' % (i + 1, k + 1))
-            lycon.save(save_file, denoised_img.astype(np.uint8))
+            lycon.save(save_file, denoised_img)
 
 bundle_submissions_raw(args.result_dir+'matfile/', 'raw_results_for_server_submission/')
 os.system("rm {}".format(args.result_dir+'matfile/*.mat'))

test_dnd_rgb.py

@@ -21,7 +21,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')
 parser.add_argument('--input_dir', default='./datasets/dnd/dnd_rgb/',
@@ -45,8 +45,6 @@
 test_dataset = get_test_data(args.input_dir)
 test_loader = DataLoader(dataset=test_dataset, batch_size=32, shuffle=False, num_workers=8, drop_last=False)
 
-
-
 model_restoration = DenoiseNet()
 
 utils.load_checkpoint(model_restoration,args.weights)
@@ -58,7 +56,6 @@
 
 model_restoration.eval()
 
-
 with torch.no_grad():
     psnr_val_rgb = []
     for ii, data_test in enumerate(tqdm(test_loader), 0):
@@ -72,13 +69,10 @@
 
         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])})
 
-  
-
 bundle_submissions_srgb_v1(args.result_dir+'matfile/', 'srgb_results_for_server_submission/')
 os.system("rm {}".format(args.result_dir+'matfile/*.mat'))

test_sidd_raw.py

@@ -21,6 +21,7 @@
 
 import utils
 import lycon
+from skimage import img_as_ubyte
 
 parser = argparse.ArgumentParser(description='RAW denoising evaluation on the validation set of SIDD')
 parser.add_argument('--input_dir', default='./datasets/sidd/sidd_raw/',
@@ -34,7 +35,6 @@
 
 args = parser.parse_args()
 
-
 os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
 
@@ -43,11 +43,8 @@
 test_dataset = get_validation_data(args.input_dir)
 test_loader = DataLoader(dataset=test_dataset, batch_size=16, shuffle=False, num_workers=8, drop_last=False)
 
-
-
 model_restoration = DenoiseNet()
 
-
 utils.load_checkpoint(model_restoration,args.weights)
 print("===>Testing using weights: ", args.weights)
 
@@ -71,9 +68,9 @@
         if args.save_images:
             for batch in range(len(raw_gt)):
                 denoised_img = utils.unpack_raw(raw_restored[batch,:,:,:].unsqueeze(0))
-                denoised_img = denoised_img.permute(0, 2, 3, 1).cpu().detach().numpy()[0] *255
+                denoised_img = denoised_img.permute(0, 2, 3, 1).cpu().detach().numpy()[0]
                 denoised_img = np.squeeze(np.stack((denoised_img,) * 3, -1))
-                lycon.save(args.result_dir + filenames[batch][:-4] + '.png', denoised_img.astype(np.uint8))
+                lycon.save(args.result_dir + filenames[batch][:-4] + '.png', img_as_ubyte(denoised_img))
 
 
 psnr_val_raw = sum(psnr_val_raw)/len(psnr_val_raw)

test_sidd_rgb.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='RGB denoising evaluation on the validation set of SIDD')
 parser.add_argument('--input_dir', default='./datasets/sidd/sidd_rgb/',
@@ -73,10 +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))
+                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))
-