feat: diff augment

models/base_model.py

@@ -11,6 +11,7 @@
 from .modules.fid.pytorch_fid.fid_score import _compute_statistics_of_path,calculate_frechet_distance
 from util.util import save_image,tensor2im
 import numpy as np
+from util.diff_aug import diff_augment
 
 class BaseModel(ABC):
     """This class is an abstract base class (ABC) for models.
@@ -93,6 +94,8 @@ def __init__(self, opt,rank):
             self.fidA=0
             self.fidB=0
 
+        self.diff_aug_policy = self.opt.diff_aug_policy
+
     @staticmethod
     def modify_commandline_options(parser, is_train):
         """Add new model-specific options, and rewrite default values for existing options.
@@ -419,7 +422,10 @@ def compute_D_loss_generic(self,netD,domain_img,loss,real_name=None,fake_name=No
             real = getattr(self,"real_"+domain_img+noisy)
         else:
             real = getattr(self,real_name)
-
+
+        real = diff_augment(real,self.diff_aug_policy)
+        fake = diff_augment(fake,self.diff_aug_policy)
+
         loss = loss.compute_loss_D(netD, real, fake)
         return loss
 
@@ -432,6 +438,9 @@ def compute_G_loss_GAN_generic(self,netD,domain_img,loss,real_name=None,fake_nam
             real = getattr(self,"real_"+domain_img)
         else:
             real = getattr(self,real_name)
+
+        real = diff_augment(real,self.diff_aug_policy)
+        fake = diff_augment(fake,self.diff_aug_policy)
 
         loss = loss.compute_loss_G(netD, real, fake)
         return loss

options/train_options.py

@@ -46,5 +46,8 @@ def initialize(self, parser):
         parser.add_argument('--use_contrastive_loss_D', action='store_true')
         parser.add_argument('--ddp_port', type=str, default='12355')
 
+        #imgaug options
+        parser.add_argument('--diff_aug_policy',type=str, default='',help='choose the augmentation policy : color translation and cutout. If you want more than one, please write them separated by a comma with no space (e.g. color,translation)')
+
         self.isTrain = True
         return parser

util/diff_aug.py

@@ -0,0 +1,72 @@
+import torch
+import torch.nn.functional as F
+
+
+def diff_augment(x, policy='', channels_first=True):
+    if policy:
+        if not channels_first:
+            x = x.permute(0, 3, 1, 2)
+        for p in policy.split(','):
+            for f in AUGMENT_FNS[p]:
+                x = f(x)
+        if not channels_first:
+            x = x.permute(0, 2, 3, 1)
+        x = x.contiguous()
+    return x
+
+
+def rand_brightness(x):
+    x = x + (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) - 0.5)
+    return x
+
+
+def rand_saturation(x):
+    x_mean = x.mean(dim=1, keepdim=True)
+    x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) * 2) + x_mean
+    return x
+
+
+def rand_contrast(x):
+    x_mean = x.mean(dim=[1, 2, 3], keepdim=True)
+    x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) + 0.5) + x_mean
+    return x
+
+
+def rand_translation(x, ratio=0.125):
+    shift_x, shift_y = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
+    translation_x = torch.randint(-shift_x, shift_x + 1, size=[x.size(0), 1, 1], device=x.device)
+    translation_y = torch.randint(-shift_y, shift_y + 1, size=[x.size(0), 1, 1], device=x.device)
+    grid_batch, grid_x, grid_y = torch.meshgrid(
+        torch.arange(x.size(0), dtype=torch.long, device=x.device),
+        torch.arange(x.size(2), dtype=torch.long, device=x.device),
+        torch.arange(x.size(3), dtype=torch.long, device=x.device),
+    )
+    grid_x = torch.clamp(grid_x + translation_x + 1, 0, x.size(2) + 1)
+    grid_y = torch.clamp(grid_y + translation_y + 1, 0, x.size(3) + 1)
+    x_pad = F.pad(x, [1, 1, 1, 1, 0, 0, 0, 0])
+    x = x_pad.permute(0, 2, 3, 1).contiguous()[grid_batch, grid_x, grid_y].permute(0, 3, 1, 2)
+    return x
+
+
+def rand_cutout(x, ratio=0.5):
+    cutout_size = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
+    offset_x = torch.randint(0, x.size(2) + (1 - cutout_size[0] % 2), size=[x.size(0), 1, 1], device=x.device)
+    offset_y = torch.randint(0, x.size(3) + (1 - cutout_size[1] % 2), size=[x.size(0), 1, 1], device=x.device)
+    grid_batch, grid_x, grid_y = torch.meshgrid(
+        torch.arange(x.size(0), dtype=torch.long, device=x.device),
+        torch.arange(cutout_size[0], dtype=torch.long, device=x.device),
+        torch.arange(cutout_size[1], dtype=torch.long, device=x.device),
+    )
+    grid_x = torch.clamp(grid_x + offset_x - cutout_size[0] // 2, min=0, max=x.size(2) - 1)
+    grid_y = torch.clamp(grid_y + offset_y - cutout_size[1] // 2, min=0, max=x.size(3) - 1)
+    mask = torch.ones(x.size(0), x.size(2), x.size(3), dtype=x.dtype, device=x.device)
+    mask[grid_batch, grid_x, grid_y] = 0
+    x = x * mask.unsqueeze(1)
+    return x
+
+
+AUGMENT_FNS = {
+    'color': [rand_brightness, rand_saturation, rand_contrast],
+    'translation': [rand_translation],
+    'cutout': [rand_cutout],
+}