Add additive mode to rrdb

codes/models/archs/RRDBNet_arch.py

@@ -144,7 +144,9 @@ class RRDBNet(nn.Module):
                  growth_channels=32,
                  body_block=RRDB,
                  blocks_per_checkpoint=4,
-                 scale=4):
+                 scale=4,
+                 additive_mode="not_additive"  # Options: "not_additive", "additive", "additive_enforced"
+                 ):
         super(RRDBNet, self).__init__()
         self.num_blocks = num_blocks
         self.blocks_per_checkpoint = blocks_per_checkpoint
@@ -166,6 +168,10 @@ class RRDBNet(nn.Module):
         self.conv_hr = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
         self.conv_last = nn.Conv2d(mid_channels, out_channels, 3, 1, 1)
 
+        self.additive_mode = additive_mode
+        if additive_mode == "additive_enforced":
+            self.add_enforced_pool = nn.AvgPool2d(kernel_size=scale, stride=scale)
+
         self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
 
         for m in [
@@ -202,6 +208,14 @@ class RRDBNet(nn.Module):
         else:
             feat = self.lrelu(self.conv_up2(feat))
         out = self.conv_last(self.lrelu(self.conv_hr(feat)))
+        if "additive" in self.additive_mode:
+            x_interp = F.interpolate(x, scale_factor=self.scale, mode='bilinear')
+        if self.additive_mode == 'additive':
+            out = out + x_interp
+        elif self.additive_mode == 'additive_enforced':
+            out_pooled = self.add_enforced_pool(out)
+            out = out - F.interpolate(out_pooled, scale_factor=self.scale, mode='nearest')
+            out = out + x_interp
         return out
 
     def visual_dbg(self, step, path):

codes/models/networks.py

@@ -43,12 +43,15 @@ def define_G(opt, net_key='network_G', scale=None):
         netG = SRResNet_arch.MSRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'],
                                        nf=opt_net['nf'], nb=opt_net['nb'], upscale=opt_net['scale'])
     elif which_model == 'RRDBNet':
+        additive_mode = opt_net['additive_mode'] if 'additive_mode' in opt_net.keys() else 'not_additive'
         netG = RRDBNet_arch.RRDBNet(in_channels=opt_net['in_nc'], out_channels=opt_net['out_nc'],
-                                    mid_channels=opt_net['nf'], num_blocks=opt_net['nb'])
+                                    mid_channels=opt_net['nf'], num_blocks=opt_net['nb'], additive_mode=additive_mode)
     elif which_model == 'RRDBNetBypass':
+        additive_mode = opt_net['additive_mode'] if 'additive_mode' in opt_net.keys() else 'not_additive'
         netG = RRDBNet_arch.RRDBNet(in_channels=opt_net['in_nc'], out_channels=opt_net['out_nc'],
                                     mid_channels=opt_net['nf'], num_blocks=opt_net['nb'], body_block=RRDBNet_arch.RRDBWithBypass,
-                                    blocks_per_checkpoint=opt_net['blocks_per_checkpoint'], scale=opt_net['scale'])
+                                    blocks_per_checkpoint=opt_net['blocks_per_checkpoint'], scale=opt_net['scale'],
+                                    additive_mode=additive_mode)
     elif which_model == 'rcan':
         #args: n_resgroups, n_resblocks, res_scale, reduction, scale, n_feats
         opt_net['rgb_range'] = 255

codes/scripts/extract_square_images.py

@@ -19,11 +19,7 @@ def main():
     # compression time. If read raw images during training, use 0 for faster IO speed.
 
     opt['dest'] = 'file'
-    opt['input_folder'] = ['F:\\4k6k\\datasets\\images\\div2k\\DIV2K_train_HR',
+    opt['input_folder'] = ['F:\\4k6k\datasets\\images\\youtube\\videos\\4k_quote_unquote\\images']
-                           'F:\\4k6k\\datasets\\images\\flickr\\flickr2k\\Flickr2K_HR',
-                           'F:\\4k6k\\datasets\\images\\flickr\\flickr-scrape\\filtered',
-                           'F:\\4k6k\\datasets\\images\\goodeats\\hq\\new_season\\images',
-                           'F:\\4k6k\datasets\\images\\youtube\\images']
     opt['save_folder'] = 'F:\\4k6k\\datasets\\images\\ge_full_1024'
     opt['imgsize'] = 1024
 

codes/train.py

@@ -291,14 +291,14 @@ class Trainer:
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_stylegan2_celebA_separated_disc.yml')
+    parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_exd_mi1_rrdb4x_6bl_corrected_disc.yml')
     parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher')
     parser.add_argument('--local_rank', type=int, default=0)
     args = parser.parse_args()
     opt = option.parse(args.opt, is_train=True)
     trainer = Trainer()
 
-    #### distributed training settings
+#### distributed training settings
     if args.launcher == 'none':  # disabled distributed training
         opt['dist'] = False
         trainer.rank = -1