diff --git a/codes/models/archs/RRDBNet_arch.py b/codes/models/archs/RRDBNet_arch.py index 5118f345..d86493f5 100644 --- a/codes/models/archs/RRDBNet_arch.py +++ b/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): diff --git a/codes/models/networks.py b/codes/models/networks.py index dffecec2..117e2806 100644 --- a/codes/models/networks.py +++ b/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 diff --git a/codes/scripts/extract_square_images.py b/codes/scripts/extract_square_images.py index 4a6d5869..1a32d912 100644 --- a/codes/scripts/extract_square_images.py +++ b/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', - '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['input_folder'] = ['F:\\4k6k\datasets\\images\\youtube\\videos\\4k_quote_unquote\\images'] opt['save_folder'] = 'F:\\4k6k\\datasets\\images\\ge_full_1024' opt['imgsize'] = 1024 diff --git a/codes/train.py b/codes/train.py index 2aff1dc9..08005c7d 100644 --- a/codes/train.py +++ b/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