created the space

LICENSE

@@ -0,0 +1,26 @@
+Deep Halftoning with Reversible Binary Pattern
+
+Copyright (c) 2021 The Chinese University of Hong Kong
+
+Copyright and License Information: The source code, the binary executable, and all data files (hereafter, Software) are copyrighted by The Chinese University of Hong Kong and Tien-Tsin Wong (hereafter, Author), Copyright (c) 2021 The Chinese University of Hong Kong. All Rights Reserved.
+
+The Author grants to you ("Licensee") a non-exclusive license to use the Software for academic, research and commercial purposes, without fee. For commercial use, Licensee should submit a WRITTEN NOTICE to the Author. The notice should clearly identify the software package/system/hardware (name, version, and/or model number) using the Software. Licensee may distribute the Software to third parties provided that the copyright notice and this statement appears on all copies. Licensee agrees that the copyright notice and this statement will appear on all copies of the Software, or portions thereof. The Author retains exclusive ownership of the Software.
+
+Licensee may make derivatives of the Software, provided that such derivatives can only be used for the purposes specified in the license grant above.
+
+THE AUTHOR MAKES NO REPRESENTATIONS OR WARRANTIES ABOUT THE SUITABILITY OF THE SOFTWARE, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, OR NON-INFRINGEMENT. THE AUTHOR SHALL NOT BE LIABLE FOR ANY DAMAGES SUFFERED BY LICENSEE AS A RESULT OF USING, MODIFYING OR DISTRIBUTING THE SOFTWARE OR ITS DERIVATIVES.
+
+By using the source code, Licensee agrees to cite the following papers in
+Licensee's publication/work:
+
+  Menghan Xia, Wenbo Hu, Xueting Liu and Tien-Tsin Wong
+  "Deep Halftoning with Reversible Binary Pattern"
+  IEEE International Conference on Computer Vision (ICCV), 2021.
+
+
+By using or copying the Software, Licensee agrees to abide by the intellectual property laws, and all other applicable laws of the U.S., and the terms of this license.
+
+Author shall have the right to terminate this license immediately by written notice upon Licensee's breach of, or non-compliance with, any of its terms.
+Licensee may be held legally responsible for any infringement that is caused or encouraged by Licensee's failure to abide by the terms of this license.
+
+For more information or comments, send mail to: ttwong@acm.org

app.py

@@ -0,0 +1,80 @@
+import gradio as gr
+import os, requests
+import numpy as np
+import torch.nn.functional as F
+from model.model import ResHalf
+from inference import Inferencer
+from utils import util
+
+## local |  remote
+RUN_MODE = "remote"
+if RUN_MODE != "local":
+    os.system("wget https://huggingface.co/menghanxia/ReversibleHalftoning/resolve/main/model_best.pth.tar")
+    os.rename("model_best.pth.tar", "./checkpoints/model_best.pth.tar")
+    ## examples
+    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/girl.png")
+    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/wave.png")
+    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/painting.png")
+
+## step 1: set up model
+device = "cpu"
+checkpt_path = "checkpoints/model_best.pth.tar"
+invhalfer = Inferencer(checkpoint_path=checkpt_path, model=ResHalf(train=False), use_cuda=False)
+
+
+def prepare_data(input_img, decoding_only=False):
+    input_img = np.array(input_img / 255., np.float32)
+    if decoding_only:
+        input_img = input_img[:,:,:1]
+    input_img = util.img2tensor(input_img * 2. - 1.)
+    return input_img
+
+
+def run_invhalf(invhalfer, input_img, decoding_only, device="cuda"):
+    input_img = prepare_data(input_img, decoding_only)
+    input_img = input_img.to(device)
+    if decoding_only:
+        print('>>>:restoration mode')
+        resColor = invhalfer(input_img, decoding_only=decoding_only)
+        output = util.tensor2img(resColor / 2. + 0.5) * 255.
+    else:
+        print('>>>:halftoning mode')
+        resHalftone, resColor = invhalfer(input_img, decoding_only=decoding_only)
+        output = util.tensor2img(resHalftone / 2. + 0.5) * 255.
+    return (output+0.5).astype(np.uint8)
+
+
+def click_run(input_img, decoding_only):
+    output = run_invhalf(invhalfer, input_img, decoding_only, device)
+    return output
+
+## step 2: configure interface
+demo = gr.Blocks(title="ReversibleHalftoning")
+with demo:
+    gr.Markdown(value="""
+                    **Gradio demo for ReversibleHalftoning: Deep Halftoning with Reversible Binary Pattern**. Check our [github page](https://github.com/MenghanXia/ReversibleHalftoning) 😛.
+                    """)
+    with gr.Row():
+        with gr.Column():
+            Image_input = gr.Image(type="numpy", label="Input", interactive=True)
+            with gr.Row():
+                Radio_mode = gr.Radio(type="index", choices=["Halftoning (Photo2Halftone)", "Restoration (Halftone2Photo)"], \
+                                                label="Choose a running mode", value="Halftoning (Photo2Halftone)")
+                Button_run = gr.Button(value="Run")
+        with gr.Column():
+            Image_output = gr.Image(type="numpy", label="Output").style(height=480)
+
+    Button_run.click(fn=click_run, inputs=[Image_input, Radio_mode], outputs=Image_output)
+    
+    if RUN_MODE == "local":
+        gr.Examples(examples=[
+                    ['girl.png', "Halftoning (Photo2Halftone)"],
+                    ['wave.png', "Halftoning (Photo2Halftone)"],
+                    ['painting.png', "Restoration (Halftone2Photo)"],
+                    ], 
+                    inputs=[Image_input,Radio_mode], outputs=[Image_output], label="Examples")
+
+if RUN_MODE != "local":
+    demo.launch(server_name='9.134.253.83',server_port=7788)
+else:
+    demo.launch()

inference.py

@@ -0,0 +1,83 @@
+import numpy as np
+import cv2
+import os, argparse, json
+from os.path import join
+from glob import glob
+
+import torch
+import torch.nn.functional as F
+
+from model.model import ResHalf
+from model.model import Quantize
+from model.loss import l1_loss
+from utils import util
+from utils.dct import DCT_Lowfrequency
+from utils.filters_tensor import bgr2gray
+
+
+class Inferencer:
+    def __init__(self, checkpoint_path, model, use_cuda=True, multi_gpu=True):
+        self.checkpoint = torch.load(checkpoint_path)
+        self.use_cuda = use_cuda
+        self.model = model.eval()
+        if multi_gpu:
+            self.model = torch.nn.DataParallel(self.model)
+        if self.use_cuda:
+            self.model = self.model.cuda()
+        self.model.load_state_dict(self.checkpoint['state_dict'])
+
+    def __call__(self, input_img, decoding_only=False):
+        with torch.no_grad():
+            scale = 8
+            _, _, H, W = input_img.shape
+            if H % scale != 0 or W % scale != 0:
+                input_img = F.pad(input_img, [0, scale - W % scale, 0, scale - H % scale], mode='reflect')
+            if self.use_cuda:
+                input_img = input_img.cuda()
+            if decoding_only:
+                resColor = self.model(input_img, decoding_only)
+                if H % scale != 0 or W % scale != 0:
+                    resColor = resColor[:, :, :H, :W]
+                return resColor
+            else:
+                resHalftone, resColor = self.model(input_img, decoding_only)
+                resHalftone = Quantize.apply((resHalftone + 1.0) * 0.5) * 2.0 - 1.
+                if H % scale != 0 or W % scale != 0:
+                    resHalftone = resHalftone[:, :, :H, :W]
+                    resColor = resColor[:, :, :H, :W]
+                return resHalftone, resColor
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='invHalf')
+    parser.add_argument('--model', default=None, type=str,
+                        help='model weight file path')
+    parser.add_argument('--decoding', action='store_true', default=False, help='restoration from halftone input')
+    parser.add_argument('--data_dir', default=None, type=str, 
+                        help='where to load input data (RGB images)')
+    parser.add_argument('--save_dir', default=None, type=str, 
+                        help='where to save the result')
+    args = parser.parse_args()
+
+    invhalfer = Inferencer(
+        checkpoint_path=args.model,
+        model=ResHalf(train=False)        
+    )
+    save_dir = os.path.join(args.save_dir)
+    util.ensure_dir(save_dir)
+    test_imgs = glob(join(args.data_dir, '*.*g'))
+    print('------loaded %d images.' % len(test_imgs) )
+    for img in test_imgs:
+        print('[*] processing %s ...' % img)
+        if args.decoding:
+            input_img = cv2.imread(img, flags=cv2.IMREAD_GRAYSCALE) / 127.5 - 1.
+            c = invhalfer(util.img2tensor(input_img), decoding_only=True)
+            c = util.tensor2img(c / 2. + 0.5) * 255.
+            cv2.imwrite(join(save_dir, 'restored_' + img.split('/')[-1].split('.')[0] + '.png'), c)
+        else:
+            input_img = cv2.imread(img, flags=cv2.IMREAD_COLOR) / 127.5 - 1.
+            h, c = invhalfer(util.img2tensor(input_img), decoding_only=False)
+            h = util.tensor2img(h / 2. + 0.5) * 255.
+            c = util.tensor2img(c / 2. + 0.5) * 255.
+            cv2.imwrite(join(save_dir, 'halftone_' + img.split('/')[-1].split('.')[0] + '.png'), h)
+            cv2.imwrite(join(save_dir, 'restored_' + img.split('/')[-1].split('.')[0] + '.png'), c)

model/base_module.py

@@ -0,0 +1,81 @@
+import torch.nn as nn
+from torch.nn import functional as F
+import torch
+import numpy as np
+
+def tensor2array(tensors):
+    arrays = tensors.detach().to("cpu").numpy()
+    return np.transpose(arrays, (0, 2, 3, 1))
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels):
+        super(ResidualBlock, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(channels, channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+        )
+
+    def forward(self, x):
+        residual = self.conv(x)
+        return x + residual
+
+
+class DownsampleBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, withConvRelu=True):
+        super(DownsampleBlock, self).__init__()
+        if withConvRelu:
+            self.conv = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, stride=2),
+                nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+                nn.ReLU(inplace=True)
+            )
+        else:
+            self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, stride=2)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, inChannels, outChannels, convNum):
+        super(ConvBlock, self).__init__()
+        self.inConv = nn.Sequential(
+            nn.Conv2d(inChannels, outChannels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True)
+        )
+        layers = []
+        for _ in range(convNum - 1):
+            layers.append(nn.Conv2d(outChannels, outChannels, kernel_size=3, padding=1))
+            layers.append(nn.ReLU(inplace=True))
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.inConv(x)
+        x = self.conv(x)
+        return x
+
+
+class UpsampleBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(UpsampleBlock, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, stride=1),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        return self.conv(x)
+
+
+class SkipConnection(nn.Module):
+    def __init__(self, channels):
+        super(SkipConnection, self).__init__()
+        self.conv = nn.Conv2d(2 * channels, channels, 1, bias=False)
+
+    def forward(self, x, y):
+        x = torch.cat((x, y), 1)
+        return self.conv(x)

model/hourglass.py

@@ -0,0 +1,70 @@
+import torch.nn as nn
+from .base_module import ConvBlock, DownsampleBlock, ResidualBlock, SkipConnection, UpsampleBlock
+
+
+class HourGlass(nn.Module):
+    def __init__(self, convNum=4, resNum=4, inChannel=6, outChannel=3):
+        super(HourGlass, self).__init__()
+        self.inConv = ConvBlock(inChannel, 64, convNum=2)
+        self.down1 = nn.Sequential(*[DownsampleBlock(64, 128, withConvRelu=False), ConvBlock(128, 128, convNum=2)])
+        self.down2 = nn.Sequential(
+            *[DownsampleBlock(128, 256, withConvRelu=False), ConvBlock(256, 256, convNum=convNum)])
+        self.down3 = nn.Sequential(
+            *[DownsampleBlock(256, 512, withConvRelu=False), ConvBlock(512, 512, convNum=convNum)])
+        self.residual = nn.Sequential(*[ResidualBlock(512) for _ in range(resNum)])
+        self.up3 = nn.Sequential(*[UpsampleBlock(512, 256), ConvBlock(256, 256, convNum=convNum)])
+        self.skip3 = SkipConnection(256)
+        self.up2 = nn.Sequential(*[UpsampleBlock(256, 128), ConvBlock(128, 128, convNum=2)])
+        self.skip2 = SkipConnection(128)
+        self.up1 = nn.Sequential(*[UpsampleBlock(128, 64), ConvBlock(64, 64, convNum=2)])
+        self.skip1 = SkipConnection(64)
+        self.outConv = nn.Sequential(
+            nn.Conv2d(64, 64, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, outChannel, kernel_size=1, padding=0)
+        )
+
+    def forward(self, x):
+        f1 = self.inConv(x)
+        f2 = self.down1(f1)
+        f3 = self.down2(f2)
+        f4 = self.down3(f3)
+        r4 = self.residual(f4)
+        r3 = self.skip3(self.up3(r4), f3)
+        r2 = self.skip2(self.up2(r3), f2)
+        r1 = self.skip1(self.up1(r2), f1)
+        y = self.outConv(r1)
+        return y
+
+
+class ResidualHourGlass(nn.Module):
+    def __init__(self, resNum=4, inChannel=6, outChannel=3):
+        super(ResidualHourGlass, self).__init__()
+        self.inConv = nn.Conv2d(inChannel, 64, kernel_size=3, padding=1)
+        self.residualBefore = nn.Sequential(*[ResidualBlock(64) for _ in range(2)])
+        self.down1 = nn.Sequential(
+            *[DownsampleBlock(64, 128, withConvRelu=False), ConvBlock(128, 128, convNum=2)])
+        self.down2 = nn.Sequential(
+            *[DownsampleBlock(128, 256, withConvRelu=False), ConvBlock(256, 256, convNum=2)])
+        self.residual = nn.Sequential(*[ResidualBlock(256) for _ in range(resNum)])
+        self.up2 = nn.Sequential(*[UpsampleBlock(256, 128), ConvBlock(128, 128, convNum=2)])
+        self.skip2 = SkipConnection(128)
+        self.up1 = nn.Sequential(*[UpsampleBlock(128, 64), ConvBlock(64, 64, convNum=2)])
+        self.skip1 = SkipConnection(64)
+        self.residualAfter = nn.Sequential(*[ResidualBlock(64) for _ in range(2)])
+        self.outConv = nn.Sequential(
+            nn.Conv2d(64, outChannel, kernel_size=3, padding=1),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        f1 = self.inConv(x)
+        f1 = self.residualBefore(f1)
+        f2 = self.down1(f1)
+        f3 = self.down2(f2)
+        r3 = self.residual(f3)
+        r2 = self.skip2(self.up2(r3), f2)
+        r1 = self.skip1(self.up1(r2), f1)
+        y = self.residualAfter(r1)
+        y = self.outConv(y)
+        return y

model/loss.py

@@ -0,0 +1,93 @@
+import os
+import torch.nn.functional as F
+import torch
+from utils.filters_tensor import GaussianSmoothing, bgr2gray
+from utils import pytorch_ssim
+from torch import nn
+from .hourglass import HourGlass
+from torchvision.models.vgg import vgg19
+
+
+def l2_loss(y_input, y_target):
+    return F.mse_loss(y_input, y_target)
+
+
+def l1_loss(y_input, y_target):
+    return F.l1_loss(y_input, y_target)
+
+
+def gaussianL2(yInput, yTarget):
+    # data range [-1,1]
+    smoother = GaussianSmoothing(channels=1, kernel_size=11, sigma=2.0)
+    gaussianInput = smoother(yInput)
+    gaussianTarget = smoother(bgr2gray(yTarget))
+    return F.mse_loss(gaussianInput, gaussianTarget)
+
+
+def binL1(yInput):
+    # data range is [-1,1]
+    return (yInput.abs() - 1.0).abs().mean()
+
+
+def ssimLoss(yInput, yTarget):
+    # data range is [-1,1]
+    ssim = pytorch_ssim.ssim(yInput / 2. + 0.5, bgr2gray(yTarget / 2. + 0.5), window_size=11)
+    return 1. - ssim
+
+
+class InverseHalf(nn.Module):
+    def __init__(self):
+        super(InverseHalf, self).__init__()
+        self.net = HourGlass(inChannel=1, outChannel=1)
+
+    def forward(self, x):
+        grayscale = self.net(x)
+        return grayscale
+
+
+class FeatureLoss:
+    def __init__(self, pretrainedPath, requireGrad=False, multiGpu=True):
+        self.featureExactor = InverseHalf()
+        if multiGpu:
+            self.featureExactor = torch.nn.DataParallel(self.featureExactor).cuda()
+        print("-loading feature extractor: {} ...".format(pretrainedPath))
+        checkpoint = torch.load(pretrainedPath)
+        self.featureExactor.load_state_dict(checkpoint['state_dict'])
+        print("-feature network loaded")
+        if not requireGrad:
+            for param in self.featureExactor.parameters():
+                param.requires_grad = False
+
+    def __call__(self, yInput, yTarget):
+        inFeature = self.featureExactor(yInput)
+        return l2_loss(inFeature, yTarget)
+
+
+class Vgg19Loss:
+    def __init__(self, multiGpu=True):
+        os.environ['TORCH_HOME']='~/bigdata/0ProgramS/checkpoints'
+        # data in BGR format, [0,1] range
+        self.mean = [0.485, 0.456, 0.406]
+        self.mean.reverse()
+        self.std = [0.229, 0.224, 0.225]
+        self.std.reverse()
+        vgg = vgg19(pretrained=True)
+        # maxpoll after conv4_4
+        self.featureExactor = nn.Sequential(*list(vgg.features)[:28]).eval()
+        for param in self.featureExactor.parameters():
+            param.requires_grad = False
+        if multiGpu:
+            self.featureExactor = torch.nn.DataParallel(self.featureExactor).cuda()
+        print('[*] Vgg19Loss init!')
+
+    def normalize(self, tensor):
+        tensor = tensor.clone()
+        mean = torch.as_tensor(self.mean, dtype=torch.float32, device=tensor.device)
+        std = torch.as_tensor(self.std, dtype=torch.float32, device=tensor.device)
+        tensor.sub_(mean[None, :, None, None]).div_(std[None, :, None, None])
+        return tensor
+
+    def __call__(self, yInput, yTarget):
+        inFeature = self.featureExactor(self.normalize(yInput).flip(1))
+        targetFeature = self.featureExactor(self.normalize(yTarget).flip(1))
+        return l2_loss(inFeature, targetFeature)

model/model.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Function
+
+from .hourglass import HourGlass
+from utils.dct import DCT_Lowfrequency
+from utils.filters_tensor import bgr2gray
+
+from collections import OrderedDict
+import numpy as np
+
+
+class Quantize(Function):
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        y = x.round()
+        return y
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        inputX = ctx.saved_tensors
+        return grad_output
+
+
+class ResHalf(nn.Module):
+    def __init__(self, train=True, warm_stage=False):
+        super(ResHalf, self).__init__()
+        self.encoder = HourGlass(inChannel=4, outChannel=1, resNum=4, convNum=4)
+        self.decoder = HourGlass(inChannel=1, outChannel=3, resNum=4, convNum=4)
+        self.dcter = DCT_Lowfrequency(size=256, fLimit=50)
+        # quantize [-1,1] data to be {-1,1}
+        self.quantizer = lambda x: Quantize.apply(0.5 * (x + 1.)) * 2. - 1.
+        self.isTrain = train
+        if warm_stage:
+            for name, param in self.decoder.named_parameters():
+                param.requires_grad = False
+
+    def add_impluse_noise(self, input_halfs, p=0.0):
+        N,C,H,W = input_halfs.shape
+        SNR = 1-p
+        np_input_halfs = input_halfs.detach().to("cpu").numpy()
+        np_input_halfs = np.transpose(np_input_halfs, (0, 2, 3, 1))
+        for i in range(N):
+            mask = np.random.choice((0, 1, 2), size=(H, W, 1), p=[SNR, (1 - SNR) / 2., (1 - SNR) / 2.])
+            np_input_halfs[i, mask==1] = 1.0
+            np_input_halfs[i, mask==2] = -1.0
+        return torch.from_numpy(np_input_halfs.transpose((0, 3, 1, 2))).to(input_halfs.device)
+
+    def forward(self, input_img, decoding_only=False):
+        if decoding_only:
+            halfResQ = self.quantizer(input_img)
+            restored = self.decoder(halfResQ)
+            return restored
+            
+        noise = torch.randn_like(input_img) * 0.3
+        halfRes = self.encoder(torch.cat((input_img, noise[:,:1,:,:]), dim=1))
+        halfResQ = self.quantizer(halfRes)
+        restored = self.decoder(halfResQ)
+        if self.isTrain:
+            halfDCT = self.dcter(halfRes / 2. + 0.5)
+            refDCT = self.dcter(bgr2gray(input_img / 2. + 0.5))
+            return halfRes, halfDCT, refDCT, restored
+        else:
+            return halfRes, restored

requirements.txt

@@ -0,0 +1,17 @@
+addict
+future
+numpy
+opencv-python
+pandas
+Pillow
+pyyaml
+requests
+scikit-image
+scikit-learn
+scipy
+torch>=1.8.0
+torchvision
+tensorboardx>=2.4
+tqdm
+yapf
+lpips

scripts/invhalf_full.json

@@ -0,0 +1,42 @@
+{
+    "name": "invhalf_full",
+    "initial_ckpt": "checkpoints/model_warm.pth.tar",
+    "model": "ResHalf",
+	"data_dir": "dataset/",
+	"save_dir": "./",
+    "trainer": {
+        "epochs": 1000,
+        "save_epochs": 5
+    },
+    "data_loader": {
+        "dataset": "HalftoneVOC2012.json",
+        "special_set": "special_color.json",
+        "batch_size": 1,
+        "shuffle": true,
+        "num_workers": 32
+    },
+    "quantizeLoss": "binL1",
+    "quantizeLossWeight": 0.1,
+    "toneLoss": "gaussianL2",
+    "toneLossWeight": 0.6,
+    "structureLoss": "ssimLoss",
+    "structureLossWeight": 0.0,
+    "restoreLoss": "l2_loss",
+    "restoreLossWeight": 1.0,
+    "blueNoiseLossWeight": 0.3,
+    "vggLossWeight": 0.0002,
+    "cuda": true,
+    "multi-gpus": true,
+    "optimizer_type": "Adam",
+    "optimizer": {
+        "lr": 0.0001,
+        "weight_decay": 0
+    },
+    "lr_sheduler": {
+        "factor": 0.5,
+        "patience": 3,
+        "threshold": 1e-05,
+        "cooldown": 0
+    },
+    "seed": 131
+}

scripts/invhalf_warm.json

@@ -0,0 +1,39 @@
+{
+    "name": "invhalf_warmup",
+    "model": "ResHalf",
+	"data_dir": "dataset/",
+	"save_dir": "./",
+    "trainer": {
+        "epochs": 1000,
+        "save_epochs": 5
+    },
+    "data_loader": {
+        "dataset": "HalftoneVOC2012.json",
+        "special_set": "special_color.json",
+        "batch_size": 8,
+        "shuffle": true,
+        "num_workers": 32
+    },	
+    "quantizeLoss": "binL1",
+    "quantizeLossWeight": 0.2,
+    "toneLoss": "gaussianL2",
+    "toneLossWeight": 0.6,
+    "structureLoss": "ssimLoss",
+    "structureLossWeight": 0.0,
+    "blueNoiseLossWeight": 0.3,
+    "featureLossWeight": 1.0,
+    "cuda": true,
+    "multi-gpus": true,
+    "optimizer_type": "Adam",
+    "optimizer": {
+        "lr": 0.0001,
+        "weight_decay": 0
+    },
+    "lr_sheduler": {
+        "factor": 0.5,
+        "patience": 3,
+        "threshold": 1e-05,
+        "cooldown": 0
+    },
+    "seed": 131
+}

train.py

@@ -0,0 +1,272 @@
+import os, glob, datetime, time
+import argparse, json
+
+import torch
+import torch.optim as optim
+from torch.autograd import Variable
+import torchvision
+from torch.utils.data import DataLoader
+from torch.backends import cudnn
+
+from model.base_module import tensor2array
+from model.model import ResHalf
+from model.loss import *
+from utils.dataset import HalftoneVOC2012 as Dataset
+from utils.util import ensure_dir, save_list, save_images_from_batch
+
+
+class Trainer():
+    def __init__(self, config, resume):
+        self.config = config
+        self.name = config['name']
+        self.resume_path = resume
+        self.n_epochs = config['trainer']['epochs']
+        self.with_cuda = config['cuda'] and torch.cuda.is_available()
+        self.seed = config['seed']
+        self.start_epoch = 0
+        self.save_freq = config['trainer']['save_epochs']
+        self.checkpoint_dir = os.path.join(config['save_dir'], self.name)
+        ensure_dir(self.checkpoint_dir)
+        json.dump(config, open(os.path.join(self.checkpoint_dir, 'config.json'), 'w'),
+                  indent=4, sort_keys=False)
+        print("@Workspace: %s *************"%self.checkpoint_dir)
+        self.cache = os.path.join(self.checkpoint_dir, 'train_cache')
+        self.val_halftone = os.path.join(self.cache, 'halftone')
+        self.val_restored = os.path.join(self.cache, 'restored')
+        ensure_dir(self.val_halftone)
+        ensure_dir(self.val_restored)
+
+        ## model        
+        self.model = eval(config['model'])()
+        if self.config['multi-gpus']:
+            self.model = torch.nn.DataParallel(self.model).cuda()
+        elif self.with_cuda:
+            self.model = self.model.cuda()
+
+        ## optimizer
+        self.optimizer = getattr(optim, config['optimizer_type'])(self.model.parameters(), **config['optimizer'])
+        self.lr_sheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, **config['lr_sheduler'])
+        
+        ## dataset loader
+        with open(os.path.join(config['data_dir'], config['data_loader']['dataset'])) as f:
+            dataset = json.load(f)
+        train_set = Dataset(dataset['train'])
+        self.train_data_loader = DataLoader(train_set, batch_size=config['data_loader']['batch_size'],
+                                   shuffle=config['data_loader']['shuffle'],
+                                   num_workers=config['data_loader']['num_workers'])
+        val_set = Dataset(dataset['val'])
+        self.valid_data_loader = DataLoader(val_set, batch_size=config['data_loader']['batch_size'],
+                                   shuffle=False,
+                                   num_workers=config['data_loader']['num_workers'])
+        # special dataloader: constant color images
+        with open(os.path.join(config['data_dir'], config['data_loader']['special_set'])) as f:
+            dataset = json.load(f)
+        specialSet = Dataset(dataset['train'])
+        self.specialDataloader = DataLoader(specialSet, batch_size=config['data_loader']['batch_size'],
+                                    shuffle=config['data_loader']['shuffle'],
+                                    num_workers=config['data_loader']['num_workers'])
+
+        ## loss function
+        self.quantizeLoss = eval(config['quantizeLoss'])
+        self.quantizeLossWeight = config['quantizeLossWeight']
+        self.toneLoss = eval(config['toneLoss'])
+        self.toneLossWeight = config['toneLossWeight']
+        self.structureLoss = eval(config['structureLoss'])
+        self.structureLossWeight = config['structureLossWeight']
+        self.restoreLoss = eval(config['restoreLoss'])
+        self.restoreLossWeight = config['restoreLossWeight']
+        # quantize [-1,1] data to be {-1,1}
+        self.quantizer = lambda x: Quantize.apply(0.5 * (x + 1.)) * 2. - 1.
+        self.blueNoiseLossWeight = config['blueNoiseLossWeight']
+        self.vggloss = Vgg19Loss()
+        self.vggLossWeight = config['vggLossWeight']
+
+        # resume checkpoint or load warm-up checkpoint
+        checkpt_path = self.config['initial_ckpt']
+        if self.resume_path:
+            checkpt_path = self.resume_path
+        assert os.path.exists(checkpt_path), 'Invalid checkpoint Path: %s' % checkpt_path
+        self.load_checkpoint(checkpt_path)
+        
+    
+    def _train(self):
+        torch.manual_seed(self.config['seed'])
+        torch.cuda.manual_seed(self.config['seed'])
+        cudnn.benchmark = True
+        
+        start_time = time.time()
+        self.monitor_best = 999.
+        for epoch in range(self.start_epoch, self.n_epochs + 1):
+            ep_st = time.time()
+            epoch_loss = self._train_epoch(epoch)
+            # perform lr_sheduler
+            self.lr_sheduler.step(epoch_loss['total_loss'])
+            epoch_lr = self.optimizer.state_dict()['param_groups'][0]['lr']
+            epoch_metric = self._valid_epoch(epoch)
+            print("[*] --- epoch: %d/%d | loss: %4.4f | metric: %4.4f | time-consumed: %4.2f ---" % \
+                (epoch+1, self.n_epochs, epoch_loss['total_loss'], epoch_metric, (time.time()-ep_st)))
+
+            # save losses and learning rate
+            epoch_loss['metric'] = epoch_metric
+            epoch_loss['lr'] = epoch_lr
+            self.save_loss(epoch_loss, epoch)
+            if ((epoch+1) % self.save_freq == 0 or epoch == (self.n_epochs-1)):
+                print('---------- saving model ...')
+                self.save_checkpoint(epoch)
+            if self.monitor_best > epoch_metric:
+                self.monitor_best = epoch_metric
+                self.save_checkpoint(epoch, save_best=True)
+
+        print("Training finished! consumed %f sec" % (time.time() - start_time))
+
+
+    def _to_variable(self, data, target):
+        data, target = Variable(data), Variable(target)
+        if self.with_cuda:
+            data, target = data.cuda(), target.cuda()
+        return data, target
+
+
+    def _train_epoch(self, epoch):
+        self.model.train()
+        total_loss, quantize_loss, restore_loss = 0, 0, 0
+        tone_loss, structure_loss, blue_noise_loss = 0, 0, 0
+
+        specialIter = iter(self.specialDataloader)
+        time_stamp = time.time()
+        for batch_idx, (color, halftone) in enumerate(self.train_data_loader):
+            color, halftone = self._to_variable(color, halftone)
+            # special data
+            try:
+                specialColor, specialHalftone = next(specialIter)
+            except StopIteration:
+                # reinitialize data loader
+                specialIter = iter(self.specialDataloader)
+                specialColor, specialHalftone = next(specialIter)
+            specialColor, specialHalftone = self._to_variable(specialColor, specialHalftone)
+            self.optimizer.zero_grad()
+            output = self.model(color, halftone)
+            quantizeLoss = self.quantizeLoss(output[0])
+            toneLoss = self.toneLoss(output[0], color)
+            structureLoss = self.structureLoss(output[0], color)
+
+            # restore
+            restoredColor = output[-1]
+            restoreLoss = self.restoreLoss(restoredColor, color)
+            vggLoss = self.vggloss(restoredColor / 2. + 0.5, color / 2. + 0.5)
+
+            # special data
+            output = self.model(specialColor, specialHalftone)
+            toneLossSpecial = self.toneLoss(output[0], specialColor)
+            blueNoiseLoss = l1_loss(output[1], output[2])
+            quantizeLossSpecial = self.quantizeLoss(output[0])
+            loss = (self.toneLossWeight * toneLoss + self.blueNoiseLossWeight*toneLossSpecial) \
+                   + self.quantizeLossWeight * (0.5*quantizeLoss + 0.5*quantizeLossSpecial) \
+                   + self.structureLossWeight * structureLoss \
+                   + self.blueNoiseLossWeight * blueNoiseLoss \
+                   + self.vggLossWeight * vggLoss \
+                   + self.restoreLossWeight * restoreLoss
+                   
+            loss.backward()
+            # apply grad clip to make training roboust
+            # torch.nn.utils.clip_grad_norm_(self.model.parameters(), 0.0001)
+            self.optimizer.step()
+
+            total_loss += loss.item()
+            quantize_loss += quantizeLoss.item()
+            restore_loss += (self.restoreLossWeight*restoreLoss + self.vggLossWeight*vggLoss).item()
+            tone_loss += toneLoss.item()
+            structure_loss += structureLoss.item()
+            blue_noise_loss += blueNoiseLoss.item()
+            if batch_idx % 100 == 0:
+                tm = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')
+                print("%s >> [%d/%d] iter:%d loss:%4.4f "%(tm, epoch+1, self.n_epochs, batch_idx+1, loss.item()))
+
+        epoch_loss = dict()
+        epoch_loss['total_loss'] = total_loss / (batch_idx+1)
+        epoch_loss['quantize_loss'] = quantize_loss / (batch_idx+1)
+        epoch_loss['tone_loss'] = tone_loss / (batch_idx+1)
+        epoch_loss['structure_loss'] = structure_loss / (batch_idx+1)
+        epoch_loss['bluenoise_loss'] = blue_noise_loss / (batch_idx+1)
+        epoch_loss['restore_loss'] = restore_loss / (batch_idx+1)
+
+        return epoch_loss
+
+
+    def _valid_epoch(self, epoch):
+        self.model.eval()
+        total_loss = 0
+        with torch.no_grad():
+            for batch_idx, (color, halftone) in enumerate(self.valid_data_loader):
+                color, halftone = self._to_variable(color, halftone)
+                output = self.model(color, halftone)
+                quantizeLoss = self.quantizeLoss(output[0])
+                toneLoss = self.toneLoss(output[0], color)
+                structureLoss = self.structureLoss(output[0], color)
+                # restore
+                restoredColor = output[-1]
+                restoreLoss = self.restoreLoss(restoredColor, color)
+                vggLoss = self.vggloss(restoredColor / 2. + 0.5, color / 2. + 0.5)
+
+                loss = self.toneLossWeight * toneLoss \
+                       + self.quantizeLossWeight * quantizeLoss \
+                       + self.structureLossWeight * structureLoss \
+                       + self.vggLossWeight * vggLoss \
+                       + self.restoreLossWeight * restoreLoss
+
+                total_loss += loss.item()
+                #! save intermediate images
+                gray_imgs = tensor2array(output[0])
+                color_imgs = tensor2array(output[-1])
+                save_images_from_batch(gray_imgs, self.val_halftone, None, batch_idx)
+                save_images_from_batch(color_imgs, self.val_restored, None, batch_idx)
+            
+            return total_loss
+
+
+    def save_loss(self, epoch_loss, epoch):
+        if epoch == 0:
+            for key in epoch_loss:
+                save_list(os.path.join(self.cache, key), [epoch_loss[key]], append_mode=False)
+        else:
+            for key in epoch_loss:
+                save_list(os.path.join(self.cache, key), [epoch_loss[key]], append_mode=True)
+
+
+    def load_checkpoint(self, checkpt_path):
+        print("-loading checkpoint from: {} ...".format(checkpt_path))
+        if self.resume_path:
+            checkpoint = torch.load(checkpt_path)
+            self.start_epoch = checkpoint['epoch'] + 1
+            self.monitor_best = checkpoint['monitor_best']
+            self.model.load_state_dict(checkpoint['state_dict'])
+            self.optimizer.load_state_dict(checkpoint['optimizer'])
+        else:
+            checkpoint = torch.load(checkpt_path)
+            self.model.load_state_dict(checkpoint['state_dict'], strict=False)
+        print("-pretrained checkpoint loaded.")
+
+    
+    def save_checkpoint(self, epoch, save_best=False):
+        state = {
+            'epoch': epoch,
+            'state_dict': self.model.state_dict(),
+            'optimizer': self.optimizer.state_dict(),
+            'monitor_best': self.monitor_best
+        }
+        save_path = os.path.join(self.checkpoint_dir, 'model_last.pth.tar')
+        if save_best:
+            save_path = os.path.join(self.checkpoint_dir, 'model_best.pth.tar')
+        torch.save(state, save_path)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='InvHalf')
+    parser.add_argument('-c', '--config', default=None, type=str,
+                        help='config file path (default: None)')
+    parser.add_argument('-r', '--resume', default=None, type=str,
+                        help='path to latest checkpoint (default: None)')
+    args = parser.parse_args()
+    config_dict = json.load(open(args.config))
+    node = Trainer(config_dict, resume=args.resume)
+    node._train()

train_warm.py

@@ -0,0 +1,256 @@
+import os, glob, datetime, time
+import argparse, json
+
+import torch
+import torch.optim as optim
+from torch.autograd import Variable
+import torchvision
+from torch.utils.data import DataLoader
+from torch.backends import cudnn
+
+from model.base_module import tensor2array
+from model.model import ResHalf
+from model.loss import *
+from utils.dataset import HalftoneVOC2012 as Dataset
+from utils.util import ensure_dir, save_list, save_images_from_batch
+
+
+class Trainer():
+    def __init__(self, config, resume):
+        self.config = config
+        self.name = config['name']
+        self.resume_path = resume
+        self.n_epochs = config['trainer']['epochs']
+        self.with_cuda = config['cuda'] and torch.cuda.is_available()
+        self.seed = config['seed']
+        self.start_epoch = 0
+        self.save_freq = config['trainer']['save_epochs']
+        self.checkpoint_dir = os.path.join(config['save_dir'], self.name)
+        ensure_dir(self.checkpoint_dir)
+        json.dump(config, open(os.path.join(self.checkpoint_dir, 'config.json'), 'w'),
+                  indent=4, sort_keys=False)
+        print("@Workspace: %s *************"%self.checkpoint_dir)
+        self.cache = os.path.join(self.checkpoint_dir, 'train_cache')
+        self.val_halftone = os.path.join(self.cache, 'halftone')
+        self.val_restored = os.path.join(self.cache, 'restored')
+        ensure_dir(self.val_halftone)
+        ensure_dir(self.val_restored)
+
+        ## model        
+        self.model = eval(config['model'])(train=True, warm_stage=True)
+        if self.config['multi-gpus']:
+            self.model = torch.nn.DataParallel(self.model).cuda()
+        elif self.with_cuda:
+            self.model = self.model.cuda()
+
+        ## optimizer
+        self.optimizer = getattr(optim, config['optimizer_type'])(self.model.parameters(), **config['optimizer'])
+        self.lr_sheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, **config['lr_sheduler'])
+        
+        ## dataset loader
+        with open(os.path.join(config['data_dir'], config['data_loader']['dataset'])) as f:
+            dataset = json.load(f)
+        train_set = Dataset(dataset['train'])
+        self.train_data_loader = DataLoader(train_set, batch_size=config['data_loader']['batch_size'],
+                                   shuffle=config['data_loader']['shuffle'],
+                                   num_workers=config['data_loader']['num_workers'])
+        val_set = Dataset(dataset['val'])
+        self.valid_data_loader = DataLoader(val_set, batch_size=config['data_loader']['batch_size'],
+                                   shuffle=False,
+                                   num_workers=config['data_loader']['num_workers'])
+        # special dataloader: constant color images
+        with open(os.path.join(config['data_dir'], config['data_loader']['special_set'])) as f:
+            dataset = json.load(f)
+        specialSet = Dataset(dataset['train'])
+        self.specialDataloader = DataLoader(specialSet, batch_size=config['data_loader']['batch_size'],
+                                    shuffle=config['data_loader']['shuffle'],
+                                    num_workers=config['data_loader']['num_workers'])
+
+        ## loss function
+        self.quantizeLoss = eval(config['quantizeLoss'])
+        self.quantizeLossWeight = config['quantizeLossWeight']
+        self.toneLoss = eval(config['toneLoss'])
+        self.toneLossWeight = config['toneLossWeight']
+        self.structureLoss = eval(config['structureLoss'])
+        self.structureLossWeight = config['structureLossWeight']
+        # quantize [-1,1] data to be {-1,1}
+        self.quantizer = lambda x: Quantize.apply(0.5 * (x + 1.)) * 2. - 1.
+        self.blueNoiseLossWeight = config['blueNoiseLossWeight']
+        self.featureLoss = FeatureLoss(
+             requireGrad=False, pretrainedPath='checkpoints/invhalftone_checkpt/model_best.pth.tar')
+        self.featureLossWeight = config['featureLossWeight']
+
+        # resume checkpoint
+        if self.resume_path:
+            assert os.path.exists(resume_path), 'Invalid checkpoint Path: %s' % resume_path
+            self.load_checkpoint(self.resume_path)
+        
+    
+    def _train(self):
+        torch.manual_seed(self.config['seed'])
+        torch.cuda.manual_seed(self.config['seed'])
+        cudnn.benchmark = True
+        
+        start_time = time.time()
+        self.monitor_best = 999.
+        for epoch in range(self.start_epoch, self.n_epochs + 1):
+            ep_st = time.time()
+            epoch_loss = self._train_epoch(epoch)
+            # perform lr_sheduler
+            self.lr_sheduler.step(epoch_loss['total_loss'])
+            epoch_lr = self.optimizer.state_dict()['param_groups'][0]['lr']
+            epoch_metric = self._valid_epoch(epoch)
+            print("[*] --- epoch: %d/%d | loss: %4.4f | metric: %4.4f | time-consumed: %4.2f ---" % \
+                (epoch+1, self.n_epochs, epoch_loss['total_loss'], epoch_metric, (time.time()-ep_st)))
+
+            # save losses and learning rate
+            epoch_loss['metric'] = epoch_metric
+            epoch_loss['lr'] = epoch_lr
+            self.save_loss(epoch_loss, epoch)
+            if ((epoch+1) % self.save_freq == 0 or epoch == (self.n_epochs-1)):
+                print('---------- saving model ...')
+                self.save_checkpoint(epoch)
+            if self.monitor_best > epoch_metric:
+                self.monitor_best = epoch_metric
+                self.save_checkpoint(epoch, save_best=True)
+
+        print("Training finished! consumed %f sec" % (time.time() - start_time))
+
+
+    def _to_variable(self, data, target):
+        data, target = Variable(data), Variable(target)
+        if self.with_cuda:
+            data, target = data.cuda(), target.cuda()
+        return data, target
+
+
+    def _train_epoch(self, epoch):
+        self.model.train()
+        total_loss, quantize_loss, feature_loss = 0, 0, 0
+        tone_loss, structure_loss, blue_noise_loss = 0, 0, 0
+
+        specialIter = iter(self.specialDataloader)
+        time_stamp = time.time()
+        for batch_idx, (color, halftone) in enumerate(self.train_data_loader):
+            color, halftone = self._to_variable(color, halftone)
+            # special data
+            try:
+                specialColor, specialHalftone = next(specialIter)
+            except StopIteration:
+                # reinitialize data loader
+                specialIter = iter(self.specialDataloader)
+                specialColor, specialHalftone = next(specialIter)
+            specialColor, specialHalftone = self._to_variable(specialColor, specialHalftone)
+            self.optimizer.zero_grad()
+            output = self.model(color, halftone)
+            quantizeLoss = self.quantizeLoss(output[0])
+            toneLoss = self.toneLoss(output[0], color)
+            structureLoss = self.structureLoss(output[0], color)
+            featureLoss = self.featureLoss(output[0], bgr2gray(color))
+
+            # special data
+            output = self.model(specialColor, specialHalftone)
+            toneLossSpecial = self.toneLoss(output[0], specialColor)
+            blueNoiseLoss = l1_loss(output[1], output[2])
+            quantizeLossSpecial = self.quantizeLoss(output[0])
+            loss = (self.toneLossWeight * toneLoss + self.blueNoiseLossWeight*toneLossSpecial) \
+                   + self.quantizeLossWeight * (0.5*quantizeLoss + 0.5*quantizeLossSpecial) \
+                   + self.structureLossWeight * structureLoss \
+                   + self.blueNoiseLossWeight * blueNoiseLoss \
+                   + self.featureLossWeight * featureLoss
+                   
+            loss.backward()
+            # apply grad clip to make training roboust
+            # torch.nn.utils.clip_grad_norm_(self.model.parameters(), 0.0001)
+            self.optimizer.step()
+
+            total_loss += loss.item()
+            quantize_loss += quantizeLoss.item()
+            feature_loss += featureLoss.item()
+            tone_loss += toneLoss.item()
+            structure_loss += structureLoss.item()
+            blue_noise_loss += blueNoiseLoss.item()
+            if batch_idx % 100 == 0:
+                tm = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')
+                print("%s >> [%d/%d] iter:%d loss:%4.4f "%(tm, epoch+1, self.n_epochs, batch_idx+1, loss.item()))
+
+        epoch_loss = dict()
+        epoch_loss['total_loss'] = total_loss / (batch_idx+1)
+        epoch_loss['quantize_loss'] = quantize_loss / (batch_idx+1)
+        epoch_loss['tone_loss'] = tone_loss / (batch_idx+1)
+        epoch_loss['structure_loss'] = structure_loss / (batch_idx+1)
+        epoch_loss['bluenoise_loss'] = blue_noise_loss / (batch_idx+1)
+        epoch_loss['feature_loss'] = feature_loss / (batch_idx+1)
+
+        return epoch_loss
+
+
+    def _valid_epoch(self, epoch):
+        self.model.eval()
+        total_loss = 0
+        with torch.no_grad():
+            for batch_idx, (color, halftone) in enumerate(self.valid_data_loader):
+                color, halftone = self._to_variable(color, halftone)
+                output = self.model(color, halftone)
+                quantizeLoss = self.quantizeLoss(output[0])
+                toneLoss = self.toneLoss(output[0], color)
+                structureLoss = self.structureLoss(output[0], color)
+                featureLoss = self.featureLoss(output[0], bgr2gray(color))
+
+                loss = self.toneLossWeight * toneLoss \
+                       + self.quantizeLossWeight * quantizeLoss \
+                       + self.structureLossWeight * structureLoss \
+                       + self.featureLossWeight * featureLoss
+
+                total_loss += loss.item()
+                #! save intermediate images
+                gray_imgs = tensor2array(output[0])
+                color_imgs = tensor2array(output[-1])
+                save_images_from_batch(gray_imgs, self.val_halftone, None, batch_idx)
+                save_images_from_batch(color_imgs, self.val_restored, None, batch_idx)
+            
+            return total_loss
+
+
+    def save_loss(self, epoch_loss, epoch):
+        if epoch == 0:
+            for key in epoch_loss:
+                save_list(os.path.join(self.cache, key), [epoch_loss[key]], append_mode=False)
+        else:
+            for key in epoch_loss:
+                save_list(os.path.join(self.cache, key), [epoch_loss[key]], append_mode=True)
+
+
+    def load_checkpoint(self, checkpt_path):
+        print("-loading checkpoint from: {} ...".format(checkpt_path))
+        checkpoint = torch.load(checkpt_path)
+        self.start_epoch = checkpoint['epoch'] + 1
+        self.monitor_best = checkpoint['monitor_best']
+        self.model.load_state_dict(checkpoint['state_dict'])
+        self.optimizer.load_state_dict(checkpoint['optimizer'])
+        print("-pretrained checkpoint loaded.")
+
+    
+    def save_checkpoint(self, epoch, save_best=False):
+        state = {
+            'epoch': epoch,
+            'state_dict': self.model.state_dict(),
+            'optimizer': self.optimizer.state_dict(),
+            'monitor_best': self.monitor_best
+        }
+        save_path = os.path.join(self.checkpoint_dir, 'model_last.pth.tar')
+        if save_best:
+            save_path = os.path.join(self.checkpoint_dir, 'model_best.pth.tar')
+        torch.save(state, save_path)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='InvHalf')
+    parser.add_argument('-c', '--config', default=None, type=str,
+                        help='config file path (default: None)')
+    parser.add_argument('-r', '--resume', default=None, type=str,
+                        help='path to latest checkpoint (default: None)')
+    args = parser.parse_args()
+    config_dict = json.load(open(args.config))
+    node = Trainer(config_dict, resume=args.resume)
+    node._train()

utils/__init__.py

@@ -0,0 +1 @@
+from .util import *

utils/_dct.py

@@ -0,0 +1,268 @@
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+
+def dct1(x):
+    """
+    Discrete Cosine Transform, Type I
+
+    :param x: the input signal
+    :return: the DCT-I of the signal over the last dimension
+    """
+    x_shape = x.shape
+    x = x.view(-1, x_shape[-1])
+
+    #return torch.rfft(torch.cat([x, x.flip([1])[:, 1:-1]], dim=1), 1)[:, :, 0].view(*x_shape)
+    return torch.fft.fft(torch.cat([x, x.flip([1])[:, 1:-1]], dim=1), 1)[:, :, 0].view(*x_shape)
+
+
+def idct1(X):
+    """
+    The inverse of DCT-I, which is just a scaled DCT-I
+
+    Our definition if idct1 is such that idct1(dct1(x)) == x
+
+    :param X: the input signal
+    :return: the inverse DCT-I of the signal over the last dimension
+    """
+    n = X.shape[-1]
+    return dct1(X) / (2 * (n - 1))
+
+
+def dct(x, norm=None):
+    """
+    Discrete Cosine Transform, Type II (a.k.a. the DCT)
+
+    For the meaning of the parameter `norm`, see:
+    https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.fftpack.dct.html
+
+    :param x: the input signal
+    :param norm: the normalization, None or 'ortho'
+    :return: the DCT-II of the signal over the last dimension
+    """
+    x_shape = x.shape
+    N = x_shape[-1]
+    x = x.contiguous().view(-1, N)
+
+    v = torch.cat([x[:, ::2], x[:, 1::2].flip([1])], dim=1)
+
+    #Vc = torch.fft.rfft(v, 1, onesided=False)
+    Vc = torch.view_as_real(torch.fft.fft(v, dim=1))
+
+    k = - torch.arange(N, dtype=x.dtype, device=x.device)[None, :] * np.pi / (2 * N)
+    W_r = torch.cos(k)
+    W_i = torch.sin(k)
+
+    V = Vc[:, :, 0] * W_r - Vc[:, :, 1] * W_i
+
+    if norm == 'ortho':
+        V[:, 0] /= np.sqrt(N) * 2
+        V[:, 1:] /= np.sqrt(N / 2) * 2
+
+    V = 2 * V.view(*x_shape)
+
+    return V
+
+
+def idct(X, norm=None):
+    """
+    The inverse to DCT-II, which is a scaled Discrete Cosine Transform, Type III
+
+    Our definition of idct is that idct(dct(x)) == x
+
+    For the meaning of the parameter `norm`, see:
+    https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.fftpack.dct.html
+
+    :param X: the input signal
+    :param norm: the normalization, None or 'ortho'
+    :return: the inverse DCT-II of the signal over the last dimension
+    """
+
+    x_shape = X.shape
+    N = x_shape[-1]
+
+    X_v = X.contiguous().view(-1, x_shape[-1]) / 2
+
+    if norm == 'ortho':
+        X_v[:, 0] *= np.sqrt(N) * 2
+        X_v[:, 1:] *= np.sqrt(N / 2) * 2
+
+    k = torch.arange(x_shape[-1], dtype=X.dtype, device=X.device)[None, :] * np.pi / (2 * N)
+    W_r = torch.cos(k)
+    W_i = torch.sin(k)
+
+    V_t_r = X_v
+    V_t_i = torch.cat([X_v[:, :1] * 0, -X_v.flip([1])[:, :-1]], dim=1)
+
+    V_r = V_t_r * W_r - V_t_i * W_i
+    V_i = V_t_r * W_i + V_t_i * W_r
+
+    V = torch.cat([V_r.unsqueeze(2), V_i.unsqueeze(2)], dim=2)
+
+    #v = torch.irfft(V, 1, onesided=False)
+    v = torch.fft.irfft(torch.view_as_complex(V), n=V.shape[1], dim=1)
+    x = v.new_zeros(v.shape)
+    x[:, ::2] += v[:, :N - (N // 2)]
+    x[:, 1::2] += v.flip([1])[:, :N // 2]
+
+    return x.view(*x_shape)
+
+
+def dct_2d(x, norm=None):
+    """
+    2-dimentional Discrete Cosine Transform, Type II (a.k.a. the DCT)
+
+    For the meaning of the parameter `norm`, see:
+    https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.fftpack.dct.html
+
+    :param x: the input signal
+    :param norm: the normalization, None or 'ortho'
+    :return: the DCT-II of the signal over the last 2 dimensions
+    """
+    X1 = dct(x, norm=norm)
+    X2 = dct(X1.transpose(-1, -2), norm=norm)
+    return X2.transpose(-1, -2)
+
+
+def idct_2d(X, norm=None):
+    """
+    The inverse to 2D DCT-II, which is a scaled Discrete Cosine Transform, Type III
+
+    Our definition of idct is that idct_2d(dct_2d(x)) == x
+
+    For the meaning of the parameter `norm`, see:
+    https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.fftpack.dct.html
+
+    :param X: the input signal
+    :param norm: the normalization, None or 'ortho'
+    :return: the DCT-II of the signal over the last 2 dimensions
+    """
+    x1 = idct(X, norm=norm)
+    x2 = idct(x1.transpose(-1, -2), norm=norm)
+    return x2.transpose(-1, -2)
+
+
+def dct_3d(x, norm=None):
+    """
+    3-dimentional Discrete Cosine Transform, Type II (a.k.a. the DCT)
+
+    For the meaning of the parameter `norm`, see:
+    https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.fftpack.dct.html
+
+    :param x: the input signal
+    :param norm: the normalization, None or 'ortho'
+    :return: the DCT-II of the signal over the last 3 dimensions
+    """
+    X1 = dct(x, norm=norm)
+    X2 = dct(X1.transpose(-1, -2), norm=norm)
+    X3 = dct(X2.transpose(-1, -3), norm=norm)
+    return X3.transpose(-1, -3).transpose(-1, -2)
+
+
+def idct_3d(X, norm=None):
+    """
+    The inverse to 3D DCT-II, which is a scaled Discrete Cosine Transform, Type III
+
+    Our definition of idct is that idct_3d(dct_3d(x)) == x
+
+    For the meaning of the parameter `norm`, see:
+    https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.fftpack.dct.html
+
+    :param X: the input signal
+    :param norm: the normalization, None or 'ortho'
+    :return: the DCT-II of the signal over the last 3 dimensions
+    """
+    x1 = idct(X, norm=norm)
+    x2 = idct(x1.transpose(-1, -2), norm=norm)
+    x3 = idct(x2.transpose(-1, -3), norm=norm)
+    return x3.transpose(-1, -3).transpose(-1, -2)
+
+
+# class LinearDCT(nn.Linear):
+#     """Implement any DCT as a linear layer; in practice this executes around
+#     50x faster on GPU. Unfortunately, the DCT matrix is stored, which will
+#     increase memory usage.
+#     :param in_features: size of expected input
+#     :param type: which dct function in this file to use"""
+#
+#     def __init__(self, in_features, type, norm=None, bias=False):
+#         self.type = type
+#         self.N = in_features
+#         self.norm = norm
+#         super(LinearDCT, self).__init__(in_features, in_features, bias=bias)
+#
+#     def reset_parameters(self):
+#         # initialise using dct function
+#         I = torch.eye(self.N)
+#         if self.type == 'dct1':
+#             self.weight.data = dct1(I).data.t()
+#         elif self.type == 'idct1':
+#             self.weight.data = idct1(I).data.t()
+#         elif self.type == 'dct':
+#             self.weight.data = dct(I, norm=self.norm).data.t()
+#         elif self.type == 'idct':
+#             self.weight.data = idct(I, norm=self.norm).data.t()
+#         self.weight.require_grad = False  # don't learn this!
+
+class LinearDCT(nn.Module):
+    """Implement any DCT as a linear layer; in practice this executes around
+    50x faster on GPU. Unfortunately, the DCT matrix is stored, which will
+    increase memory usage.
+    :param in_features: size of expected input
+    :param type: which dct function in this file to use"""
+
+    def __init__(self, in_features, type, norm=None):
+        super(LinearDCT, self).__init__()
+        self.type = type
+        self.N = in_features
+        self.norm = norm
+        I = torch.eye(self.N)
+        if self.type == 'dct1':
+            self.weight = dct1(I).data.t()
+        elif self.type == 'idct1':
+            self.weight = idct1(I).data.t()
+        elif self.type == 'dct':
+            self.weight = dct(I, norm=self.norm).data.t()
+        elif self.type == 'idct':
+            self.weight = idct(I, norm=self.norm).data.t()
+        # self.register_buffer('weight', kernel)
+        # self.weight = kernel
+
+    def forward(self, x):
+        return F.linear(x, weight=self.weight.cuda(x.get_device()))
+
+
+def apply_linear_2d(x, linear_layer):
+    """Can be used with a LinearDCT layer to do a 2D DCT.
+    :param x: the input signal
+    :param linear_layer: any PyTorch Linear layer
+    :return: result of linear layer applied to last 2 dimensions
+    """
+    X1 = linear_layer(x)
+    X2 = linear_layer(X1.transpose(-1, -2))
+    return X2.transpose(-1, -2)
+
+
+def apply_linear_3d(x, linear_layer):
+    """Can be used with a LinearDCT layer to do a 3D DCT.
+    :param x: the input signal
+    :param linear_layer: any PyTorch Linear layer
+    :return: result of linear layer applied to last 3 dimensions
+    """
+    X1 = linear_layer(x)
+    X2 = linear_layer(X1.transpose(-1, -2))
+    X3 = linear_layer(X2.transpose(-1, -3))
+    return X3.transpose(-1, -3).transpose(-1, -2)
+
+
+if __name__ == '__main__':
+    x = torch.Tensor(1000, 4096)
+    x.normal_(0, 1)
+    linear_dct = LinearDCT(4096, 'dct')
+    error = torch.abs(dct(x) - linear_dct(x))
+    assert error.max() < 1e-3, (error, error.max())
+    linear_idct = LinearDCT(4096, 'idct')
+    error = torch.abs(idct(x) - linear_idct(x))
+    assert error.max() < 1e-3, (error, error.max())

utils/dataset.py

@@ -0,0 +1,39 @@
+import torch
+import torch.utils.data as data
+import cv2
+import numpy as np
+from os.path import join
+
+
+class HalftoneVOC2012(data.Dataset):
+    # data range is [-1,1], color image is in BGR format
+    def __init__(self, data_list):
+        super(HalftoneVOC2012, self).__init__()
+        self.inputs = [join('Data', x) for x in data_list['inputs']]
+        self.labels = [join('Data', x) for x in data_list['labels']]
+
+    @staticmethod
+    def load_input(name):
+        img = cv2.imread(name, flags=cv2.IMREAD_COLOR)
+        # transpose data
+        img = img.transpose((2, 0, 1))
+        # to Tensor
+        img = torch.from_numpy(img.astype(np.float32) / 127.5 - 1.0)
+        return img
+
+    @staticmethod
+    def load_label(name):
+        img = cv2.imread(name, flags=cv2.IMREAD_GRAYSCALE)
+        # transpose data
+        img = img[np.newaxis, :, :]
+        # to Tensor
+        img = torch.from_numpy(img.astype(np.float32) / 127.5 - 1.0)
+        return img
+
+    def __getitem__(self, index):
+        input_data = self.load_input(self.inputs[index])
+        label_data = self.load_label(self.labels[index])
+        return input_data, label_data
+
+    def __len__(self):
+        return len(self.inputs)

utils/dct.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python
+"""
+-------------------------------------------------
+   File Name：   dct
+   Author :      wenbo
+   date：         12/4/2019
+   Description :
+-------------------------------------------------
+   Change Activity:
+                   12/4/2019:
+-------------------------------------------------
+"""
+__author__ = 'wenbo'
+
+from torch import nn
+from ._dct import LinearDCT, apply_linear_2d
+
+
+class DCT_Lowfrequency(nn.Module):
+    def __init__(self, size=256, fLimit=50):
+        super(DCT_Lowfrequency, self).__init__()
+        self.fLimit = fLimit
+        self.dct = LinearDCT(size, type='dct', norm='ortho')
+        self.dctTransformer = lambda x: apply_linear_2d(x, self.dct)
+
+    def forward(self, x):
+        x = self.dctTransformer(x)
+        x = x[:, :, :self.fLimit, :self.fLimit]
+        return x

utils/filters_tensor.py

@@ -0,0 +1,81 @@
+import math
+import numbers
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class GaussianSmoothing(nn.Module):
+    """
+    Apply gaussian smoothing on a
+    1d, 2d or 3d tensor. Filtering is performed seperately for each channel
+    in the input using a depthwise convolution.
+    Arguments:
+        channels (int, sequence): Number of channels of the input tensors. Output will
+            have this number of channels as well.
+        kernel_size (int, sequence): Size of the gaussian kernel.
+        sigma (float, sequence): Standard deviation of the gaussian kernel.
+        dim (int, optional): The number of dimensions of the data.
+            Default value is 2 (spatial).
+    """
+
+    def __init__(self, channels, kernel_size, sigma, dim=2, cuda=True):
+        super(GaussianSmoothing, self).__init__()
+        if isinstance(kernel_size, numbers.Number):
+            kernel_size = [kernel_size] * dim
+        if isinstance(sigma, numbers.Number):
+            sigma = [sigma] * dim
+
+        # The gaussian kernel is the product of the
+        # gaussian function of each dimension.
+        kernel = 1
+        meshgrids = torch.meshgrid(
+            [
+                torch.arange(size, dtype=torch.float32)
+                for size in kernel_size
+            ]
+        )
+        for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
+            mean = (size - 1) / 2
+            kernel *= 1 / (std * math.sqrt(2 * math.pi)) * torch.exp(-((mgrid - mean) / std) ** 2 / 2)
+
+        # Make sure sum of values in gaussian kernel equals 1.
+        kernel = kernel / torch.sum(kernel)
+
+        # Reshape to depthwise convolutional weight
+        kernel = kernel.view(1, 1, *kernel.size())
+        kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1))
+
+        # if cuda:
+        #     kernel = kernel.cuda()
+        # self.register_buffer('weight', kernel)
+        self.weight = kernel
+        self.groups = channels
+
+        if dim == 1:
+            self.conv = F.conv1d
+        elif dim == 2:
+            self.conv = F.conv2d
+        elif dim == 3:
+            self.conv = F.conv3d
+        else:
+            raise RuntimeError(
+                'Only 1, 2 and 3 dimensions are supported. Received {}.'.format(dim)
+            )
+
+    def forward(self, input):
+        """
+        Apply gaussian filter to input.
+        Arguments:
+            input (torch.Tensor): Input to apply gaussian filter on.
+        Returns:
+            filtered (torch.Tensor): Filtered output.
+        """
+        return self.conv(input, weight=self.weight.cuda(input.get_device()), groups=self.groups)
+
+
+def bgr2gray(color):
+    # gray = 0.299⋅R+0.587⋅G+0.114⋅B
+    gray = color[:, 0, ...] * 0.114 + color[:, 1, ...] * 0.587 + color[:, 2, ...] * 0.299
+    gray = gray.unsqueeze_(1)
+    return gray

utils/pytorch_ssim.py

@@ -0,0 +1,77 @@
+import torch
+import torch.nn.functional as F
+from torch.autograd import Variable
+import numpy as np
+from math import exp
+
+
+def gaussian(window_size, sigma):
+    gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])
+    return gauss / gauss.sum()
+
+
+def create_window(window_size, channel):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
+    return window
+
+
+def _ssim(img1, img2, window, window_size, channel, size_average=True):
+    mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)
+    mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+
+    sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq
+    sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq
+    sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2
+
+    C1 = 0.01 ** 2
+    C2 = 0.03 ** 2
+
+    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
+
+    if size_average:
+        return ssim_map.mean()
+    else:
+        return ssim_map.mean(1).mean(1).mean(1)
+
+
+class SSIM(torch.nn.Module):
+    def __init__(self, window_size=11, size_average=True):
+        super(SSIM, self).__init__()
+        self.window_size = window_size
+        self.size_average = size_average
+        self.channel = 1
+        self.window = create_window(window_size, self.channel)
+
+    def forward(self, img1, img2):
+        (_, channel, _, _) = img1.size()
+
+        if channel == self.channel and self.window.data.type() == img1.data.type():
+            window = self.window
+        else:
+            window = create_window(self.window_size, channel)
+
+            if img1.is_cuda:
+                window = window.cuda(img1.get_device())
+            window = window.type_as(img1)
+
+            self.window = window
+            self.channel = channel
+
+        return _ssim(img1, img2, window, self.window_size, channel, self.size_average)
+
+
+def ssim(img1, img2, window_size=11, size_average=True):
+    (_, channel, _, _) = img1.size()
+    window = create_window(window_size, channel)
+
+    if img1.is_cuda:
+        window = window.cuda(img1.get_device())
+    window = window.type_as(img1)
+
+    return _ssim(img1, img2, window, window_size, channel, size_average)

utils/util.py

@@ -0,0 +1,78 @@
+import os
+import numpy as np
+import cv2
+import torch
+
+def ensure_dir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+def get_filelist(data_dir):
+    file_list = glob.glob(os.path.join(data_dir, '*.*'))
+    file_list.sort()
+    return file_list
+    
+
+def collect_filenames(data_dir):
+    file_list = get_filelist(data_dir)
+    name_list = []
+    for file_path in file_list:
+        _, file_name = os.path.split(file_path)
+        name_list.append(file_name)
+    name_list.sort()
+    return name_list
+
+
+def save_list(save_path, data_list, append_mode=False):
+    n = len(data_list)
+    if append_mode:
+        with open(save_path, 'a') as f:
+            f.writelines([str(data_list[i]) + '\n' for i in range(n-1,n)])
+    else:
+        with open(save_path, 'w') as f:
+            f.writelines([str(data_list[i]) + '\n' for i in range(n)])
+    return None
+
+
+def save_images_from_batch(img_batch, save_dir, filename_list, batch_no=-1):
+    N,H,W,C = img_batch.shape
+    if C == 3:
+        #! rgb color image
+        for i in range(N):
+            # [-1,1] >>> [0,255]
+            img_batch_i = np.clip(img_batch[i,:,:,:]*0.5+0.5, 0, 1)
+            image = (255.0*img_batch_i).astype(np.uint8)
+            save_name = filename_list[i] if batch_no==-1 else '%05d.png' % (batch_no*N+i)
+            cv2.imwrite(os.path.join(save_dir, save_name), image)
+    elif C == 1:
+        #! single-channel gray image
+        for i in range(N):
+            # [-1,1] >>> [0,255]
+            img_batch_i = np.clip(img_batch[i,:,:,0]*0.5+0.5, 0, 1)
+            image = (255.0*img_batch_i).astype(np.uint8)
+            save_name = filename_list[i] if batch_no==-1 else '%05d.png' % (batch_no*img_batch.shape[0]+i)
+            cv2.imwrite(os.path.join(save_dir, save_name), image)
+    return None
+
+
+def imagesc(nd_array):
+    plt.imshow(nd_array)
+    plt.colorbar()
+    plt.show()
+
+
+def img2tensor(img):
+    if len(img.shape) == 2:
+        img = img[..., np.newaxis]
+    img_t = np.expand_dims(img.transpose(2, 0, 1), axis=0)
+    img_t = torch.from_numpy(img_t.astype(np.float32))
+    return img_t
+
+
+def tensor2img(img_t):
+    img = img_t[0].detach().to("cpu").numpy()
+    img = np.transpose(img, (1, 2, 0))
+    if img.shape[-1] == 1:
+        img = img[..., 0]
+    return img