added robust support

Author: kocchop

train.py

@@ -72,6 +72,7 @@ def getOpt():
     parser.add_argument("--model", type=str, default="nyu_modelA", required = True, help="name of the model (nyu_modelA | nyu_modelB)")
     parser.add_argument("--dataset_path", type=str, default="/home/mdl/mzk591/dataset/data.nyuv2/disk3/", help="path to the dataset")
     parser.add_argument("--batch_size", type=int, default=16, help="size of the batches")
+    parser.add_argument('--robust', '-r', action='store_true', help="flag to enable robust training")
     parser.add_argument("--save_size", type=int, default=8, help="batch size for saved outputs")
     parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
     parser.add_argument("--b1", type=float, default=0.9, help="adam: decay of first order momentum of gradient")
@@ -217,6 +218,14 @@ def main():
     # snapshot_interval = round(total_train_batches/2)
     snapshot_interval = 30
 
+
+    if opt.robust:
+        # Finding noisy batches
+        train_rgb_noise, train_sparse_noise = send_noisy_batches(total_train_batches, train_flag=True)
+
+        logger.info("RGB noisy batches for training are: {}".format(train_rgb_noise))
+        logger.info("Sparse noisy batches for training are: {}".format(train_sparse_noise))
+
     # ----------
     #  Training
     # ----------
@@ -238,6 +247,19 @@ def main():
             gt_temp = torch.unsqueeze(imgs["gt"], 1)
             rgb_temp = imgs["rgb"]
 
+            if opt.robust:
+                # do not want to start the training during warm up
+                rstart = False
+                if batches_done >= opt.warmup_batches:
+                    rstart = True 
+                    
+                if (i in train_rgb_noise) and rstart:
+                    rgb_temp = torch.zeros(rgb_temp.size()) # it can be any other noise
+                    logger.info("Current batch {} is a noisy RGB sample!".format(batches_done))
+                elif (i in train_sparse_noise) and rstart:
+                    sparse_temp = torch.zeros(sparse_temp.size()) # it can be any other form of noise
+                    logger.info("Current batch {} is a noisy sparse sample!".format(batches_done))
+            
             # Configure model input
             sparse_depth = Variable(sparse_temp.type(Tensor))
             gt_depth = Variable(gt_temp.type(Tensor))

utils.py

@@ -6,6 +6,7 @@
 import os
 from torchvision.utils import make_grid
 from torchvision import transforms
+from numpy.random import default_rng
 
 import logging
 import cv2
@@ -36,7 +37,12 @@ def generate_depth_cmap(in_tensor):
     for img in range(depth_tensor.shape[0]):
         min_val = np.amin(depth_tensor[img])
         max_val = np.amax(depth_tensor[img])
-        gray = (depth_tensor[img]-min_val)/(max_val-min_val)
+
+        if (max_val - min_val) < 1e-4: # when they are pretty close, no normalization
+            gray = depth_tensor[img]
+        else:
+            gray = (depth_tensor[img]-min_val)/(max_val-min_val)
+        
         # gray = depth_tensor[img]/255.0
         gray = np.clip(gray,0,1)
         heatmap = np.round(colormap(gray) * 255).astype(np.uint8)[:,:,:3]
@@ -121,18 +127,44 @@ def save_my_image(image_array, fp) -> None:
     # print('image_array',image_array.shape)
     cv2.imwrite(fp, image_array)
 
-def save_sample_images(gt_depth, imgs_rgb, sparse_depth, gen_depth, image_save_path, image_id) -> None:
+def save_sample_images(gt_depth, imgs_rgb, sparse_depth, gen_depth, image_save_path, image_id, sn_flag=False, rn_flag=False) -> None:
 
     denorm_gt = denormalize_dense(gt_depth)
-    denorm_sparse = denormalize_sparse(sparse_depth)
+
+    if sn_flag:
+        denorm_sparse = sparse_depth
+    else:
+        denorm_sparse = denormalize_sparse(sparse_depth)
+    
     denorm_pred = denormalize_dense(gen_depth)
 
     gt_depth = generate_depth_cmap(denorm_gt)
     sparse_depth = generate_depth_cmap(denorm_sparse)
     gen_depth = generate_depth_cmap(denorm_pred)
 
-    imgs_rgb = denormalize_rgb(imgs_rgb).permute(0,2,3,1).to('cpu').detach().numpy()
+    if rn_flag:
+        imgs_rgb = imgs_rgb.permute(0,2,3,1).to('cpu').detach().numpy()
+    else:
+        imgs_rgb = denormalize_rgb(imgs_rgb).permute(0,2,3,1).to('cpu').detach().numpy()
 
     img_grid = np.concatenate((gt_depth, imgs_rgb, sparse_depth, gen_depth), axis=2)
     saved_image_file = os.path.join(image_save_path,"%04d.png"%image_id)
-    save_my_image(img_grid, saved_image_file)
+    save_my_image(img_grid, saved_image_file)
+
+def send_noisy_batches(batches, train_flag=False, ratio=0.2):
+
+    n_noisy_batches = round(batches * ratio)
+    n_rgb = round(n_noisy_batches * 0.5)
+
+    rng = default_rng()
+    
+    if train_flag:
+        allowed = np.arange(batches + 1)
+        selection = rng.choice(allowed, size=n_noisy_batches, replace=False)
+    else:
+        selection = rng.choice(batches+1, size=n_noisy_batches, replace=False)
+
+    rgb_batches = selection[:n_rgb]
+    sparse_batches = selection[n_rgb:]
+
+    return sorted(rgb_batches), sorted(sparse_batches)

validate.py

@@ -63,6 +63,7 @@ def getOpt():
     parser.add_argument("--model", type=str, default="nyu_modelA", required = True, help="name of the model (nyu_modelA | nyu_modelB)")
     parser.add_argument("--dataset_path", type=str, default="/home/mdl/mzk591/dataset/data.nyuv2/disk3/", help="path to the dataset")
     parser.add_argument("--batch_size", type=int, default=4, help="size of the batches")
+    parser.add_argument('--robust', '-r', action='store_true', help="flag to enable robust training")
     parser.add_argument("--save_size", type=int, default=8, help="batch size for saved outputs")
     parser.add_argument("--n_cpu", type=int, default=16, help="number of cpu threads to use during batch generation")
     parser.add_argument("--channels", type=int, default=1, help="number of image channels")
@@ -79,9 +80,21 @@ def getOpt():
 def validate(generator, discriminator, opt, Tensor, val_dataloader, criterion_GAN, criterion_content, criterion_pixel, logger, val_image_save_path, writer, batches_done=0):
 
     total_val_batches = len(val_dataloader)
-    
-    batch_to_be_saved = np.random.randint(total_val_batches, size=5)
-    # batch_to_be_saved = [1, 2, 3, 4] #it can be any numbers
+
+    if opt.robust:
+        # Finding noisy batches
+        val_rgb_noise, val_sparse_noise = send_noisy_batches(total_val_batches, train_flag=False)
+
+        logger.info("RGB noisy batches for validation are: {}".format(val_rgb_noise))
+        logger.info("Sparse noisy batches for validation are: {}".format(val_sparse_noise))
+
+        batch_to_be_saved = np.random.randint(total_val_batches, size=3)
+        batch_to_be_saved = set(batch_to_be_saved)
+        batch_to_be_saved.add(val_rgb_noise[0])
+        batch_to_be_saved.add(val_sparse_noise[0])
+    else:  
+        batch_to_be_saved = np.random.randint(total_val_batches, size=5)
+        # batch_to_be_saved = [1, 2, 3, 4] #it can be any numbers
 
     val_sample_path = os.path.join(val_image_save_path,"%06d"%batches_done)
     os.makedirs(val_sample_path, exist_ok=True)
@@ -94,6 +107,14 @@ def validate(generator, discriminator, opt, Tensor, val_dataloader, criterion_GA
         sparse_temp = torch.unsqueeze(imgs["sparse"], 1)
         gt_temp = torch.unsqueeze(imgs["gt"], 1)
         rgb_temp = imgs["rgb"]
+
+        if opt.robust:
+            if (i in val_rgb_noise):
+                rgb_temp = torch.zeros(rgb_temp.size()) # it can be any other noise
+                logger.info("Current batch {} is a noisy RGB sample!".format(i+1))
+            elif (i in val_sparse_noise):
+                sparse_temp = torch.zeros(sparse_temp.size()) # it can be any other form of noise
+                logger.info("Current batch {} is a noisy sparse sample!".format(i+1))
 
         # Configure model input
         sparse_depth = Variable(sparse_temp.type(Tensor))