diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e3fbdf1 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/generate.py b/generate.py new file mode 100644 index 0000000..966355d --- /dev/null +++ b/generate.py @@ -0,0 +1,131 @@ +import argparse +import os + +import cv2 +import numpy as np +import torch +import sys + +import utils.imgops as ops +import utils.architecture.architecture as arch + +parser = argparse.ArgumentParser() +parser.add_argument('--input', default='input', help='Input folder') +parser.add_argument('--output', default='output', help='Output folder') +parser.add_argument('--reverse', help='Reverse Order', action="store_true") +parser.add_argument('--tile_size', default=512, + help='Tile size for splitting', type=int) +parser.add_argument('--seamless', action='store_true', + help='Seamless upscaling') +parser.add_argument('--mirror', action='store_true', + help='Mirrored seamless upscaling') +parser.add_argument('--replicate', action='store_true', + help='Replicate edge pixels for padding') +parser.add_argument('--cpu', action='store_true', + help='Use CPU instead of CUDA') +args = parser.parse_args() + +if not os.path.exists(args.input): + print('Error: Folder [{:s}] does not exist.'.format(args.input)) + sys.exit(1) +elif os.path.isfile(args.input): + print('Error: Folder [{:s}] is a file.'.format(args.input)) + sys.exit(1) +elif os.path.isfile(args.output): + print('Error: Folder [{:s}] is a file.'.format(args.output)) + sys.exit(1) +elif not os.path.exists(args.output): + os.mkdir(args.output) + +device = torch.device('cpu' if args.cpu else 'cuda') + +input_folder = os.path.normpath(args.input) +output_folder = os.path.normpath(args.output) + +NORMAL_MAP_MODEL = 'utils/models/1x_NormalMapGenerator-CX-Lite_200000_G.pth' +OTHER_MAP_MODEL = 'utils/models/1x_FrankenMapGenerator-CX-Lite_215000_G.pth' + +def process(img, model): + img = img * 1. / np.iinfo(img.dtype).max + img = img[:, :, [2, 1, 0]] + img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float() + img_LR = img.unsqueeze(0) + img_LR = img_LR.to(device) + + output = model(img_LR).data.squeeze( + 0).float().cpu().clamp_(0, 1).numpy() + output = output[[2, 1, 0], :, :] + output = np.transpose(output, (1, 2, 0)) + output = (output * 255.).round() + return output + +def load_model(model_path): + global device + state_dict = torch.load(model_path) + model = arch.RRDB_Net(3, 3, 32, 12, gc=32, upscale=1, norm_type=None, act_type='leakyrelu', + mode='CNA', res_scale=1, upsample_mode='upconv') + model.load_state_dict(state_dict, strict=True) + del state_dict + model.eval() + for k, v in model.named_parameters(): + v.requires_grad = False + return model.to(device) + +images=[] +for root, _, files in os.walk(input_folder): + for file in sorted(files, reverse=args.reverse): + if file.split('.')[-1].lower() in ['png', 'jpg', 'jpeg', 'gif', 'bmp', 'tiff', 'tga']: + images.append(os.path.join(root, file)) +for idx, path in enumerate(images, 1): + base = os.path.splitext(os.path.relpath(path, input_folder))[0] + output_dir = os.path.dirname(os.path.join(output_folder, base)) + os.makedirs(output_dir, exist_ok=True) + print(idx, base) + # read image + img = cv2.imread(path, cv2.cv2.IMREAD_COLOR) + + # Seamless modes + if args.seamless: + img = cv2.copyMakeBorder(img, 16, 16, 16, 16, cv2.BORDER_WRAP) + elif args.mirror: + img = cv2.copyMakeBorder(img, 16, 16, 16, 16, cv2.BORDER_REFLECT_101) + elif args.replicate: + img = cv2.copyMakeBorder(img, 16, 16, 16, 16, cv2.BORDER_REPLICATE) + + img_height, img_width = img.shape[:2] + + # Whether or not to perform the split/merge action + do_split = img_height > args.tile_size or img_width > args.tile_size + + # NORMAL MAP + model = load_model(NORMAL_MAP_MODEL) + + if do_split: + rlt = ops.esrgan_launcher_split_merge(img, process, 1, args.tile_size, model) + else: + rlt = process(img, model) + + if args.seamless or args.mirror or args.replicate: + rlt = ops.crop_seamless(rlt) + + rlt = rlt.astype('uint8') + + cv2.imwrite(os.path.join(output_folder, '{:s}_Normal.png'.format(base)), rlt) + + # ROUGHNESS/DISPLACEMENT MAPS + model = load_model(OTHER_MAP_MODEL) + + if do_split: + rlt = ops.esrgan_launcher_split_merge(img, process, 1, args.tile_size, model) + else: + rlt = process(img, model) + + if args.seamless or args.mirror or args.replicate: + rlt = ops.crop_seamless(rlt) + + rlt = rlt.astype('uint8') + roughness = rlt[:, :, 1] + displacement = rlt[:, :, 0] + + cv2.imwrite(os.path.join(output_folder, '{:s}_Roughness.png'.format(base)), roughness) + cv2.imwrite(os.path.join(output_folder, '{:s}_Displacement.png'.format(base)), displacement) diff --git a/input/example.png b/input/example.png new file mode 100644 index 0000000..2243a3b Binary files /dev/null and b/input/example.png differ diff --git a/output/example_Displacement.png b/output/example_Displacement.png new file mode 100644 index 0000000..3ff6d65 Binary files /dev/null and b/output/example_Displacement.png differ diff --git a/output/example_Normal.png b/output/example_Normal.png new file mode 100644 index 0000000..c5467ea Binary files /dev/null and b/output/example_Normal.png differ diff --git a/output/example_Roughness.png b/output/example_Roughness.png new file mode 100644 index 0000000..6cc58ca Binary files /dev/null and b/output/example_Roughness.png differ diff --git a/utils/architecture/architecture.py b/utils/architecture/architecture.py new file mode 100644 index 0000000..b8740d6 --- /dev/null +++ b/utils/architecture/architecture.py @@ -0,0 +1,40 @@ +import math +import torch.nn as nn +import utils.architecture.block as B + +#################### +# Generator +#################### + +class RRDB_Net(nn.Module): + def __init__(self, in_nc, out_nc, nf, nb, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', \ + mode='CNA', res_scale=1, upsample_mode='upconv'): + super(RRDB_Net, self).__init__() + n_upscale = int(math.log(upscale, 2)) + if upscale == 3: + n_upscale = 1 + + fea_conv = B.conv_block(in_nc, nf, kernel_size=3, norm_type=None, act_type=None) + rb_blocks = [B.RRDB(nf, kernel_size=3, gc=32, stride=1, bias=True, pad_type='zero', \ + norm_type=norm_type, act_type=act_type, mode='CNA') for _ in range(nb)] + LR_conv = B.conv_block(nf, nf, kernel_size=3, norm_type=norm_type, act_type=None, mode=mode) + + if upsample_mode == 'upconv': + upsample_block = B.upconv_blcok + elif upsample_mode == 'pixelshuffle': + upsample_block = B.pixelshuffle_block + else: + raise NotImplementedError('upsample mode [%s] is not found' % upsample_mode) + if upscale == 3: + upsampler = upsample_block(nf, nf, 3, act_type=act_type) + else: + upsampler = [upsample_block(nf, nf, act_type=act_type) for _ in range(n_upscale)] + HR_conv0 = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=act_type) + HR_conv1 = B.conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None) + + self.model = B.sequential(fea_conv, B.ShortcutBlock(B.sequential(*rb_blocks, LR_conv)),\ + *upsampler, HR_conv0, HR_conv1) + + def forward(self, x): + x = self.model(x) + return x \ No newline at end of file diff --git a/utils/architecture/block.py b/utils/architecture/block.py new file mode 100644 index 0000000..4d49b86 --- /dev/null +++ b/utils/architecture/block.py @@ -0,0 +1,278 @@ +from collections import OrderedDict +import torch +import torch.nn as nn + +#################### +# Basic blocks +#################### + + +def act(act_type, inplace=True, neg_slope=0.2, n_prelu=1): + # helper selecting activation + # neg_slope: for leakyrelu and init of prelu + # n_prelu: for p_relu num_parameters + act_type = act_type.lower() + if act_type == 'relu': + layer = nn.ReLU(inplace) + elif act_type == 'leakyrelu': + layer = nn.LeakyReLU(neg_slope, inplace) + elif act_type == 'prelu': + layer = nn.PReLU(num_parameters=n_prelu, init=neg_slope) + else: + raise NotImplementedError('activation layer [{:s}] is not found'.format(act_type)) + return layer + + +def norm(norm_type, nc): + # helper selecting normalization layer + norm_type = norm_type.lower() + if norm_type == 'batch': + layer = nn.BatchNorm2d(nc, affine=True) + elif norm_type == 'instance': + layer = nn.InstanceNorm2d(nc, affine=False) + else: + raise NotImplementedError('normalization layer [{:s}] is not found'.format(norm_type)) + return layer + + +def pad(pad_type, padding): + # helper selecting padding layer + # if padding is 'zero', do by conv layers + pad_type = pad_type.lower() + if padding == 0: + return None + if pad_type == 'reflect': + layer = nn.ReflectionPad2d(padding) + elif pad_type == 'replicate': + layer = nn.ReplicationPad2d(padding) + else: + raise NotImplementedError('padding layer [{:s}] is not implemented'.format(pad_type)) + return layer + + +def get_valid_padding(kernel_size, dilation): + kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1) + padding = (kernel_size - 1) // 2 + return padding + + +class ConcatBlock(nn.Module): + # Concat the output of a submodule to its input + def __init__(self, submodule): + super(ConcatBlock, self).__init__() + self.sub = submodule + + def forward(self, x): + output = torch.cat((x, self.sub(x)), dim=1) + return output + + def __repr__(self): + tmpstr = 'Identity .. \n|' + modstr = self.sub.__repr__().replace('\n', '\n|') + tmpstr = tmpstr + modstr + return tmpstr + + +class ShortcutBlock(nn.Module): + #Elementwise sum the output of a submodule to its input + def __init__(self, submodule): + super(ShortcutBlock, self).__init__() + self.sub = submodule + + def forward(self, x): + output = x + self.sub(x) + return output + + def __repr__(self): + tmpstr = 'Identity + \n|' + modstr = self.sub.__repr__().replace('\n', '\n|') + tmpstr = tmpstr + modstr + return tmpstr + + +class ShortcutBlockSPSR(nn.Module): + #Elementwise sum the output of a submodule to its input + def __init__(self, submodule): + super(ShortcutBlockSPSR, self).__init__() + self.sub = submodule + + def forward(self, x): + return x, self.sub + + def __repr__(self): + tmpstr = 'Identity + \n|' + modstr = self.sub.__repr__().replace('\n', '\n|') + tmpstr = tmpstr + modstr + return tmpstr + + +def sequential(*args): + # Flatten Sequential. It unwraps nn.Sequential. + if len(args) == 1: + if isinstance(args[0], OrderedDict): + raise NotImplementedError('sequential does not support OrderedDict input.') + return args[0] # No sequential is needed. + modules = [] + for module in args: + if isinstance(module, nn.Sequential): + for submodule in module.children(): + modules.append(submodule) + elif isinstance(module, nn.Module): + modules.append(module) + return nn.Sequential(*modules) + + +def conv_block(in_nc, out_nc, kernel_size, stride=1, dilation=1, groups=1, bias=True, \ + pad_type='zero', norm_type=None, act_type='relu', mode='CNA'): + ''' + Conv layer with padding, normalization, activation + mode: CNA --> Conv -> Norm -> Act + NAC --> Norm -> Act --> Conv (Identity Mappings in Deep Residual Networks, ECCV16) + ''' + assert mode in ['CNA', 'NAC', 'CNAC'], 'Wrong conv mode [{:s}]'.format(mode) + padding = get_valid_padding(kernel_size, dilation) + p = pad(pad_type, padding) if pad_type and pad_type != 'zero' else None + padding = padding if pad_type == 'zero' else 0 + + c = nn.Conv2d(in_nc, out_nc, kernel_size=kernel_size, stride=stride, padding=padding, \ + dilation=dilation, bias=bias, groups=groups) + a = act(act_type) if act_type else None + if 'CNA' in mode: + n = norm(norm_type, out_nc) if norm_type else None + return sequential(p, c, n, a) + elif mode == 'NAC': + if norm_type is None and act_type is not None: + a = act(act_type, inplace=False) + # Important! + # input----ReLU(inplace)----Conv--+----output + # |________________________| + # inplace ReLU will modify the input, therefore wrong output + n = norm(norm_type, in_nc) if norm_type else None + return sequential(n, a, p, c) + + +#################### +# Useful blocks +#################### + + +class ResNetBlock(nn.Module): + ''' + ResNet Block, 3-3 style + with extra residual scaling used in EDSR + (Enhanced Deep Residual Networks for Single Image Super-Resolution, CVPRW 17) + ''' + + def __init__(self, in_nc, mid_nc, out_nc, kernel_size=3, stride=1, dilation=1, groups=1, \ + bias=True, pad_type='zero', norm_type=None, act_type='relu', mode='CNA', res_scale=1): + super(ResNetBlock, self).__init__() + conv0 = conv_block(in_nc, mid_nc, kernel_size, stride, dilation, groups, bias, pad_type, \ + norm_type, act_type, mode) + if mode == 'CNA': + act_type = None + if mode == 'CNAC': # Residual path: |-CNAC-| + act_type = None + norm_type = None + conv1 = conv_block(mid_nc, out_nc, kernel_size, stride, dilation, groups, bias, pad_type, \ + norm_type, act_type, mode) + # if in_nc != out_nc: + # self.project = conv_block(in_nc, out_nc, 1, stride, dilation, 1, bias, pad_type, \ + # None, None) + # print('Need a projecter in ResNetBlock.') + # else: + # self.project = lambda x:x + self.res = sequential(conv0, conv1) + self.res_scale = res_scale + + def forward(self, x): + res = self.res(x).mul(self.res_scale) + return x + res + + +class ResidualDenseBlock_5C(nn.Module): + ''' + Residual Dense Block + style: 5 convs + The core module of paper: (Residual Dense Network for Image Super-Resolution, CVPR 18) + ''' + + def __init__(self, nc, kernel_size=3, gc=32, stride=1, bias=True, pad_type='zero', \ + norm_type=None, act_type='leakyrelu', mode='CNA'): + super(ResidualDenseBlock_5C, self).__init__() + # gc: growth channel, i.e. intermediate channels + self.conv1 = conv_block(nc, gc, kernel_size, stride, bias=bias, pad_type=pad_type, \ + norm_type=norm_type, act_type=act_type, mode=mode) + self.conv2 = conv_block(nc+gc, gc, kernel_size, stride, bias=bias, pad_type=pad_type, \ + norm_type=norm_type, act_type=act_type, mode=mode) + self.conv3 = conv_block(nc+2*gc, gc, kernel_size, stride, bias=bias, pad_type=pad_type, \ + norm_type=norm_type, act_type=act_type, mode=mode) + self.conv4 = conv_block(nc+3*gc, gc, kernel_size, stride, bias=bias, pad_type=pad_type, \ + norm_type=norm_type, act_type=act_type, mode=mode) + if mode == 'CNA': + last_act = None + else: + last_act = act_type + self.conv5 = conv_block(nc+4*gc, nc, 3, stride, bias=bias, pad_type=pad_type, \ + norm_type=norm_type, act_type=last_act, mode=mode) + + def forward(self, x): + x1 = self.conv1(x) + x2 = self.conv2(torch.cat((x, x1), 1)) + x3 = self.conv3(torch.cat((x, x1, x2), 1)) + x4 = self.conv4(torch.cat((x, x1, x2, x3), 1)) + x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1)) + return x5.mul(0.2) + x + + +class RRDB(nn.Module): + ''' + Residual in Residual Dense Block + (ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks) + ''' + + def __init__(self, nc, kernel_size=3, gc=32, stride=1, bias=True, pad_type='zero', \ + norm_type=None, act_type='leakyrelu', mode='CNA'): + super(RRDB, self).__init__() + self.RDB1 = ResidualDenseBlock_5C(nc, kernel_size, gc, stride, bias, pad_type, \ + norm_type, act_type, mode) + self.RDB2 = ResidualDenseBlock_5C(nc, kernel_size, gc, stride, bias, pad_type, \ + norm_type, act_type, mode) + self.RDB3 = ResidualDenseBlock_5C(nc, kernel_size, gc, stride, bias, pad_type, \ + norm_type, act_type, mode) + + def forward(self, x): + out = self.RDB1(x) + out = self.RDB2(out) + out = self.RDB3(out) + return out.mul(0.2) + x + + +#################### +# Upsampler +#################### + + +def pixelshuffle_block(in_nc, out_nc, upscale_factor=2, kernel_size=3, stride=1, bias=True, \ + pad_type='zero', norm_type=None, act_type='relu'): + ''' + Pixel shuffle layer + (Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional + Neural Network, CVPR17) + ''' + conv = conv_block(in_nc, out_nc * (upscale_factor ** 2), kernel_size, stride, bias=bias, \ + pad_type=pad_type, norm_type=None, act_type=None) + pixel_shuffle = nn.PixelShuffle(upscale_factor) + + n = norm(norm_type, out_nc) if norm_type else None + a = act(act_type) if act_type else None + return sequential(conv, pixel_shuffle, n, a) + + +def upconv_blcok(in_nc, out_nc, upscale_factor=2, kernel_size=3, stride=1, bias=True, \ + pad_type='zero', norm_type=None, act_type='relu', mode='nearest'): + # Up conv + # described in https://distill.pub/2016/deconv-checkerboard/ + upsample = nn.Upsample(scale_factor=upscale_factor, mode=mode) + conv = conv_block(in_nc, out_nc, kernel_size, stride, bias=bias, \ + pad_type=pad_type, norm_type=norm_type, act_type=act_type) + return sequential(upsample, conv) diff --git a/utils/imgops.py b/utils/imgops.py new file mode 100644 index 0000000..df727ae --- /dev/null +++ b/utils/imgops.py @@ -0,0 +1,74 @@ +import numpy as np +import math + +def crop_seamless(img): + img_height, img_width = img.shape[:2] + y, x = 16, 16 + h, w = img_height - 32, img_width - 32 + img = img[y:y+h, x:x+w] + return img + +# from https://github.com/ata4/esrgan-launcher/blob/master/upscale.py +def esrgan_launcher_split_merge(input_image, upscale_function, model, scale_factor=4, tile_size=512, tile_padding=0.125): + width, height, depth = input_image.shape + output_width = width * scale_factor + output_height = height * scale_factor + output_shape = (output_width, output_height, depth) + + # start with black image + output_image = np.zeros(output_shape, np.uint8) + + tile_padding = math.ceil(tile_size * tile_padding) + tile_size = math.ceil(tile_size / scale_factor) + + tiles_x = math.ceil(width / tile_size) + tiles_y = math.ceil(height / tile_size) + + for y in range(tiles_y): + for x in range(tiles_x): + # extract tile from input image + ofs_x = x * tile_size + ofs_y = y * tile_size + + # input tile area on total image + input_start_x = ofs_x + input_end_x = min(ofs_x + tile_size, width) + + input_start_y = ofs_y + input_end_y = min(ofs_y + tile_size, height) + + # input tile area on total image with padding + input_start_x_pad = max(input_start_x - tile_padding, 0) + input_end_x_pad = min(input_end_x + tile_padding, width) + + input_start_y_pad = max(input_start_y - tile_padding, 0) + input_end_y_pad = min(input_end_y + tile_padding, height) + + # input tile dimensions + input_tile_width = input_end_x - input_start_x + input_tile_height = input_end_y - input_start_y + + input_tile = input_image[input_start_x_pad:input_end_x_pad, input_start_y_pad:input_end_y_pad] + + # upscale tile + output_tile = upscale_function(input_tile, model) + + # output tile area on total image + output_start_x = input_start_x * scale_factor + output_end_x = input_end_x * scale_factor + + output_start_y = input_start_y * scale_factor + output_end_y = input_end_y * scale_factor + + # output tile area without padding + output_start_x_tile = (input_start_x - input_start_x_pad) * scale_factor + output_end_x_tile = output_start_x_tile + input_tile_width * scale_factor + + output_start_y_tile = (input_start_y - input_start_y_pad) * scale_factor + output_end_y_tile = output_start_y_tile + input_tile_height * scale_factor + + # put tile into output image + output_image[output_start_x:output_end_x, output_start_y:output_end_y] = \ + output_tile[output_start_x_tile:output_end_x_tile, output_start_y_tile:output_end_y_tile] + + return output_image \ No newline at end of file diff --git a/utils/models/1x_FrankenMapGenerator-CX-Lite_215000_G.pth b/utils/models/1x_FrankenMapGenerator-CX-Lite_215000_G.pth new file mode 100644 index 0000000..86c82bb Binary files /dev/null and b/utils/models/1x_FrankenMapGenerator-CX-Lite_215000_G.pth differ diff --git a/utils/models/1x_NormalMapGenerator-CX-Lite_200000_G.pth b/utils/models/1x_NormalMapGenerator-CX-Lite_200000_G.pth new file mode 100644 index 0000000..60b0d8c Binary files /dev/null and b/utils/models/1x_NormalMapGenerator-CX-Lite_200000_G.pth differ