First commit

.gitignore

@@ -0,0 +1,129 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 CVLab@StonyBrook
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,13 +0,0 @@
----
-title: PaperEdgeDemo
-emoji: 🌖
-colorFrom: blue
-colorTo: green
-sdk: gradio
-sdk_version: 3.6
-app_file: app.py
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,114 @@
+# -*- encoding: utf-8 -*-
+import copy
+import os
+
+os.system('pip install -r requirements.txt')
+
+import time
+from pathlib import Path
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from networks.paperedge_cpu import GlobalWarper, LocalWarper, WarperUtil
+import gradio as gr
+
+
+cv2.setNumThreads(0)
+cv2.ocl.setUseOpenCL(False)
+
+
+class PaperEdge(object):
+    def __init__(self, enet_path, tnet_path, device, dst_dir) -> None:
+        self.device = device
+        self.dst_dir = dst_dir
+
+        self.netG = GlobalWarper().to(device)
+        netG_state = torch.load(enet_path, map_location=device)['G']
+        self.netG.load_state_dict(netG_state)
+        self.netG.eval()
+
+        self.netL = LocalWarper().to(device)
+        netL_state = torch.load(tnet_path, map_location=device)['L']
+        self.netL.load_state_dict(netL_state)
+        self.netL.eval()
+
+        self.warpUtil = WarperUtil(64).to(device)
+
+    @staticmethod
+    def load_img(img_path):
+        im = cv2.imread(img_path).astype(np.float32) / 255.0
+        im = im[:, :, (2, 1, 0)]
+        im = cv2.resize(im, (256, 256), interpolation=cv2.INTER_AREA)
+        im = torch.from_numpy(np.transpose(im, (2, 0, 1)))
+        return im
+
+    def __call__(self, img_path):
+        time_stamp = time.strftime('%Y-%m-%d-%H-%M-%S',
+                                   time.localtime(time.time()))
+
+        gs_d, ls_d = None, None
+        with torch.no_grad():
+            x = self.load_img(img_path)
+            x = x.unsqueeze(0).to(self.device)
+
+            d = self.netG(x)
+
+            d = self.warpUtil.global_post_warp(d, 64)
+
+            gs_d = copy.deepcopy(d)
+
+            d = F.interpolate(d, size=256, mode='bilinear', align_corners=True)
+            y0 = F.grid_sample(x, d.permute(0, 2, 3, 1), align_corners=True)
+            ls_d = self.netL(y0)
+
+            ls_d = F.interpolate(ls_d, size=256, mode='bilinear', align_corners=True)
+            ls_d = ls_d.clamp(-1.0, 1.0)
+
+        im = cv2.imread(img_path).astype(np.float32) / 255.0
+        im = torch.from_numpy(np.transpose(im, (2, 0, 1)))
+        im = im.to(self.device).unsqueeze(0)
+
+        gs_d = F.interpolate(gs_d, (im.size(2), im.size(3)), mode='bilinear', align_corners=True)
+        gs_y = F.grid_sample(im, gs_d.permute(0, 2, 3, 1), align_corners=True).detach()
+
+        ls_d = F.interpolate(ls_d, (im.size(2), im.size(3)), mode='bilinear', align_corners=True)
+        ls_y = F.grid_sample(gs_y, ls_d.permute(0, 2, 3, 1), align_corners=True).detach()
+        ls_y = ls_y.squeeze().permute(1, 2, 0).cpu().numpy()
+
+        save_path = f'{self.dst_dir}/{time_stamp}.png'
+        cv2.imwrite(save_path, ls_y * 255.)
+        return save_path
+
+
+def inference(img):
+    img_path = img.name
+    save_img_path = paper_edge(img_path)
+    return save_img_path
+
+
+enet_path = 'models/G_w_checkpoint_13820.pt'
+tnet_path = 'models/L_w_checkpoint_27640.pt'
+device = torch.device('cpu')
+
+dst_dir = Path('inference/')
+if not dst_dir.exists():
+    dst_dir.mkdir(parents=True, exist_ok=True)
+
+paper_edge = PaperEdge(enet_path, tnet_path, device, dst_dir)
+
+title = 'PaperEdge Demo'
+description = 'This is the demo for the paper "Learning From Documents in the Wild to Improve Document Unwarping" (SIGGRAPH 2022). Github repo: https://github.com/cvlab-stonybrook/PaperEdge'
+css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}"
+
+gr.Interface(
+    inference,
+    inputs=gr.inputs.Image(type='file', label='Input'),
+    outputs=[
+        gr.outputs.Image(type='file', label='Output_image'),
+    ],
+    title=title,
+    description=description,
+    css=css,
+    allow_flagging='never',
+    ).launch(debug=True, enable_queue=True)

demo_cpu.py

@@ -0,0 +1,98 @@
+# -*- encoding: utf-8 -*-
+import argparse
+import copy
+import time
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from networks.paperedge_cpu import GlobalWarper, LocalWarper, WarperUtil
+
+cv2.setNumThreads(0)
+cv2.ocl.setUseOpenCL(False)
+
+
+class PaperEdge(object):
+    def __init__(self, enet_path, tnet_path, device) -> None:
+        self.device = device
+
+        self.netG = GlobalWarper().to(device)
+        netG_state = torch.load(enet_path, map_location=device)['G']
+        self.netG.load_state_dict(netG_state)
+        self.netG.eval()
+
+        self.netL = LocalWarper().to(device)
+        netL_state = torch.load(tnet_path, map_location=device)['L']
+        self.netL.load_state_dict(netL_state)
+        self.netL.eval()
+
+        self.warpUtil = WarperUtil(64).to(device)
+
+    @staticmethod
+    def load_img(img_path):
+        im = cv2.imread(img_path).astype(np.float32) / 255.0
+        im = im[:, :, (2, 1, 0)]
+        im = cv2.resize(im, (256, 256), interpolation=cv2.INTER_AREA)
+        im = torch.from_numpy(np.transpose(im, (2, 0, 1)))
+        return im
+
+    def infer(self, img_path):
+        gs_d, ls_d = None, None
+        with torch.no_grad():
+            x = self.load_img(img_path)
+            x = x.unsqueeze(0).to(self.device)
+
+            d = self.netG(x)
+
+            d = self.warpUtil.global_post_warp(d, 64)
+
+            gs_d = copy.deepcopy(d)
+
+            d = F.interpolate(d, size=256, mode='bilinear', align_corners=True)
+            y0 = F.grid_sample(x, d.permute(0, 2, 3, 1), align_corners=True)
+            ls_d = self.netL(y0)
+
+            ls_d = F.interpolate(ls_d, size=256, mode='bilinear', align_corners=True)
+            ls_d = ls_d.clamp(-1.0, 1.0)
+
+        im = cv2.imread(img_path).astype(np.float32) / 255.0
+        im = torch.from_numpy(np.transpose(im, (2, 0, 1)))
+        im = im.to(self.device).unsqueeze(0)
+
+        gs_d = F.interpolate(gs_d, (im.size(2), im.size(3)), mode='bilinear', align_corners=True)
+        gs_y = F.grid_sample(im, gs_d.permute(0, 2, 3, 1), align_corners=True).detach()
+
+        ls_d = F.interpolate(ls_d, (im.size(2), im.size(3)), mode='bilinear', align_corners=True)
+        ls_y = F.grid_sample(gs_y, ls_d.permute(0, 2, 3, 1), align_corners=True).detach()
+        ls_y = ls_y.squeeze().permute(1, 2, 0).cpu().numpy()
+
+        save_path = f'{dst_dir}/result_ls.png'
+        cv2.imwrite(save_path, ls_y * 255.)
+        return save_path
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--Enet_ckpt', type=str,
+                        default='models/G_w_checkpoint_13820.pt')
+    parser.add_argument('--Tnet_ckpt', type=str,
+                        default='models/L_w_checkpoint_27640.pt')
+    parser.add_argument('--img_path', type=str, default='images/3.jpg')
+    parser.add_argument('--out_dir', type=str, default='output')
+    parser.add_argument('--device', type=str, default='cpu')
+    args = parser.parse_args()
+
+    if args.device == 'cuda' and torch.cuda.is_available():
+        device = torch.device('cuda:0')
+    else:
+        device = torch.device('cpu')
+
+    dst_dir = args.out_dir
+    Path(dst_dir).mkdir(parents=True, exist_ok=True)
+
+    paper_edge = PaperEdge(args.Enet_ckpt, args.Tnet_ckpt, args.device)
+
+    paper_edge.inder(args.img_path)
+    print('ok')

models/G_w_checkpoint_13820.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:49489f5490f6786e4cc6896b47c3ff6511a92438fec1b3be7b1c53580975236f
+size 146530295

models/L_w_checkpoint_27640.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eed0653d9c9fa2bd1b21f469336fc534bfaebfe20422624581141b47ddab7b6f
+size 146530295

networks/paperedge.py

@@ -0,0 +1,553 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+# from torch.nn.utils import spectral_norm as SN
+# from torchvision.models.densenet import _DenseBlock
+from .tps_warp import TpsWarp, PspWarp
+from functools import partial
+# import plotly.graph_objects as go
+import random
+import numpy as np
+import cv2
+
+torch.autograd.set_detect_anomaly(True)
+# torch.manual_seed(0)
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.actv = nn.ReLU()
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.actv(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.actv(out)
+
+        return out
+
+def _make_layer(block, inplanes, planes, blocks, stride=1, dilate=False):
+        norm_layer = nn.BatchNorm2d
+        downsample = None
+
+        if stride != 1 or inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(inplanes, planes * block.expansion, 1, stride, bias=False),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(inplanes, planes, stride, downsample, norm_layer=norm_layer))
+        for _ in range(1, blocks):
+            layers.append(block(planes, planes,
+                                norm_layer=norm_layer))
+
+        return nn.Sequential(*layers)
+
+class Interpolate(nn.Module):
+    def __init__(self, size, mode):
+        super(Interpolate, self).__init__()
+        self.interp = nn.functional.interpolate
+        self.size = size
+        self.mode = mode
+
+    def forward(self, x):
+        x = self.interp(x, size=self.size, mode=self.mode, align_corners=True)
+        return x
+
+class GlobalWarper(nn.Module):
+    def __init__(self):
+        super(GlobalWarper, self).__init__()
+        modules = [
+            nn.Conv2d(5, 64, kernel_size=7, stride=2, padding=3, bias=False),
+            nn.BatchNorm2d(64),
+            nn.ReLU()
+        ]
+
+        # encoder
+        planes = [64, 128, 256, 256, 512, 512]
+        strides = [2, 2, 2, 2, 2]
+        blocks = [1, 1, 1, 1, 1]
+        for k in range(len(planes) - 1):
+            modules.append(_make_layer(BasicBlock, planes[k], planes[k + 1], blocks[k], strides[k]))
+        self.encoder = nn.Sequential(*modules)
+
+        # decoder
+        modules = []
+        planes = [512, 512, 256, 128, 64]
+        strides = [2, 2, 2, 2]
+        # tsizes = [3, 5, 9, 17, 33]
+        blocks = [1, 1, 1, 1]
+        for k in range(len(planes) - 1):
+            # modules += [nn.Sequential(Interpolate(size=tsizes[k], mode='bilinear'),
+            #             _make_layer(BasicBlock, planes[k], planes[k + 1], blocks[k], 1))]
+            modules += [nn.Sequential(nn.Upsample(scale_factor=strides[k], mode='bilinear', align_corners=True),
+                        _make_layer(BasicBlock, planes[k], planes[k + 1], blocks[k], 1))]
+        # self.decoder = nn.ModuleList(modules)
+        self.decoder = nn.Sequential(*modules)
+
+        self.to_warp = nn.Sequential(nn.Conv2d(64, 2, 1))
+        self.to_warp[0].weight.data.fill_(0.0)
+        self.to_warp[0].bias.data.fill_(0.0)
+
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, 256), torch.linspace(-1, 1, 256))
+        self.coord = torch.stack((ix, iy), dim=0).unsqueeze(0).to('cuda')
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, 64), torch.linspace(-1, 1, 64))
+        ### note we mulitply a 0.9 so the network is initialized closer to GT. This is different from localwarper net
+        self.basegrid = torch.stack((ix * 0.9, iy * 0.9), dim=0).unsqueeze(0).to('cuda')
+
+        # # box filter
+        # ksize = 7
+        # p = int((ksize - 1) / 2)
+        # self.pad_replct = partial(F.pad, pad=(p, p, p, p), mode='replicate')
+        # bw = torch.ones(1, 1, ksize, ksize, device='cuda') / ksize / ksize
+        # self.box_filter = partial(F.conv2d, weight=bw)
+
+
+
+    def forward(self, im):
+        # print(self.to_warp[0].weight.data)
+        # coordconv
+        B = im.size(0)
+        c = self.coord.expand(B, -1, -1, -1).detach()
+        t = torch.cat((im, c), dim=1)
+
+        t = self.encoder(t)
+        t = self.decoder(t)
+        t = self.to_warp(t)
+
+        gs = t + self.basegrid
+
+        return gs
+
+class LocalWarper(nn.Module):
+    def __init__(self):
+        super().__init__()
+        modules = [
+            nn.Conv2d(5, 64, kernel_size=7, stride=2, padding=3, bias=False),
+            nn.BatchNorm2d(64),
+            nn.ReLU()
+        ]
+        # encoder
+        planes = [64, 128, 256, 256, 512, 512]
+        strides = [2, 2, 2, 2, 2]
+        blocks = [1, 1, 1, 1, 1]
+        for k in range(len(planes) - 1):
+            modules.append(_make_layer(BasicBlock, planes[k], planes[k + 1], blocks[k], strides[k]))
+        self.encoder = nn.Sequential(*modules)
+
+        # decoder
+        modules = []
+        planes = [512, 512, 256, 128, 64]
+        strides = [2, 2, 2, 2]
+        # tsizes = [3, 5, 9, 17, 33]
+        blocks = [1, 1, 1, 1]
+        for k in range(len(planes) - 1):
+            modules += [nn.Sequential(nn.Upsample(scale_factor=strides[k], mode='bilinear', align_corners=True),
+                        _make_layer(BasicBlock, planes[k], planes[k + 1], blocks[k], 1))]
+        self.decoder = nn.Sequential(*modules)
+
+        self.to_warp = nn.Sequential(nn.Conv2d(64, 2, 1))
+        self.to_warp[0].weight.data.fill_(0.0)
+        self.to_warp[0].bias.data.fill_(0.0)
+
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, 256), torch.linspace(-1, 1, 256))
+        self.coord = torch.stack((ix, iy), dim=0).unsqueeze(0).to('cuda')
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, 64), torch.linspace(-1, 1, 64))
+        self.basegrid = torch.stack((ix, iy), dim=0).unsqueeze(0).to('cuda')
+
+        # box filter
+        ksize = 5
+        p = int((ksize - 1) / 2)
+        self.pad_replct = partial(F.pad, pad=(p, p, p, p), mode='replicate')
+        bw = torch.ones(1, 1, ksize, ksize, device='cuda') / ksize / ksize
+        self.box_filter = partial(F.conv2d, weight=bw)
+
+    def forward(self, im):
+        c = self.coord.expand(im.size(0), -1, -1, -1).detach()
+        t = torch.cat((im, c), dim=1)
+
+        # encoder
+        t = self.encoder(t)
+        t = self.decoder(t)
+        t = self.to_warp(t)
+
+        # # filter
+        # t = self.pad_replct(t)
+        # tx = self.box_filter(t[:, 0 : 1, ...])
+        # ty = self.box_filter(t[:, 1 : 2, ...])
+        # t = torch.cat((tx, ty), dim=1)
+
+        # bd condition
+        t[..., 1, 0, :] = 0
+        t[..., 1, -1, :] = 0
+        t[..., 0, :, 0] = 0
+        t[..., 0, :, -1] = 0
+
+        gs = t + self.basegrid
+        return gs
+
+def gs_to_bd(gs):
+    # gs: B 2 H W
+    t = torch.cat([gs[..., 0, :], gs[..., -1, :], gs[..., 1 : -1, 0], gs[..., 1 : -1, -1]], dim=2).permute(0, 2, 1)
+    # t: B 2(W + H - 1) 2
+    return t
+
+class MaskLoss(nn.Module):
+    def __init__(self, gsize):
+        super().__init__()
+        self.tpswarper = TpsWarp(gsize)
+        self.pspwarper = PspWarp()
+        # self.imsize = imsize
+        self.msk = torch.ones(1, 1, gsize, gsize, device='cuda')
+        self.cn = torch.tensor([[-1, -1], [1, -1], [1, 1], [-1, 1]], dtype=torch.float, device='cuda').unsqueeze(0)
+
+    def forward(self, gs, y, s):
+        # resize gs to s*s
+        B, _, s0, _ = gs.size()
+        tgs = F.interpolate(gs, s, mode='bilinear', align_corners=True)
+
+        # use only the boundary points
+        srcpts = gs_to_bd(tgs)
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, s), torch.linspace(-1, 1, s))
+        t = torch.stack((ix, iy), dim=0).unsqueeze(0).to('cuda').expand_as(tgs)
+        dstpts = gs_to_bd(t)
+
+        tgs_f = self.tpswarper(srcpts, dstpts.detach())
+        ym = self.msk.expand_as(y)
+        yh = F.grid_sample(ym, tgs_f.permute(0, 2, 3, 1), align_corners=True)
+        loss_f = F.l1_loss(yh, y)
+
+        # forward/backward consistency loss
+        tgs_b = self.tpswarper(dstpts.detach(), srcpts)
+        # tgs_b = F.interpolate(tgs, s0, mode='bilinear', align_corners=True)
+        yy = F.grid_sample(y, tgs_b.permute(0, 2, 3, 1), align_corners=True)
+        loss_b = F.l1_loss(yy, ym)
+
+        return loss_f + loss_b, tgs_f
+
+    def _dist(self, x):
+        # adjacent point distance
+        # B, 2, n
+        x = torch.cat([x[..., 0 : 1].detach(), x[..., 1 : -1], x[..., -1 : ].detach()], dim=2)
+        d = x[..., 1:] - x[..., :-1]
+        return torch.norm(d, dim=1)
+
+# class TVLoss(nn.Module):
+#     def __init__(self):
+#         super(TVLoss, self).__init__()
+
+#     def forward(self, gs):
+#         loss = self._dist(gs[..., 1:], gs[..., :-1]) + self._dist(gs[..., 1:, :], gs[..., :-1, :])
+#         return loss
+
+#     def _dist(self, x1, x0):
+#         d = torch.norm(x1 - x0, dim=1, keepdim=True)
+#         d = torch.abs(d - torch.mean(d, dim=(2, 3), keepdim=True)).mean()
+#         return d
+
+class WarperUtil(nn.Module):
+    def __init__(self, imsize):
+        super().__init__()
+        self.tpswarper = TpsWarp(imsize)
+        self.pspwarper = PspWarp()
+        self.s = imsize
+
+    def global_post_warp(self, gs, s):
+        # B, _, s0, _ = gs.size()
+        gs = F.interpolate(gs, s, mode='bilinear', align_corners=True)
+        # gs = F.interpolate(gs, s0, mode='bilinear', align_corners=True)
+        # extract info
+        m1 = gs[..., 0, :]
+        m2 = gs[..., -1, :]
+        n1 = gs[..., 0]
+        n2 = gs[..., -1]
+        # for x
+        m1x_interval_ratio = m1[:, 0, 1:] - m1[:, 0, :-1]
+        m1x_interval_ratio /= m1x_interval_ratio.sum(dim=1, keepdim=True)
+        m2x_interval_ratio = m2[:, 0, 1:] - m2[:, 0, :-1]
+        m2x_interval_ratio /= m2x_interval_ratio.sum(dim=1, keepdim=True)
+        # interpolate all x ratio
+        t = torch.stack([m1x_interval_ratio, m2x_interval_ratio], dim=1).unsqueeze(1)
+        mx_interval_ratio = F.interpolate(t, (s, m1x_interval_ratio.size(1)), mode='bilinear', align_corners=True)
+        mx_interval = (n2[..., 0 : 1, :] - n1[..., 0 : 1, :]).unsqueeze(3) * mx_interval_ratio
+        # cumsum to x
+        dx = torch.cumsum(mx_interval, dim=3) + n1[..., 0 : 1, :].unsqueeze(3)
+        dx = dx[..., 1 : -1, :-1]
+        # for y
+        n1y_interval_ratio = n1[:, 1, 1:] - n1[:, 1, :-1]
+        n1y_interval_ratio /= n1y_interval_ratio.sum(dim=1, keepdim=True)
+        n2y_interval_ratio = n2[:, 1, 1:] - n2[:, 1, :-1]
+        n2y_interval_ratio /= n2y_interval_ratio.sum(dim=1, keepdim=True)
+        # interpolate all x ratio
+        t = torch.stack([n1y_interval_ratio, n2y_interval_ratio], dim=2).unsqueeze(1)
+        ny_interval_ratio = F.interpolate(t, (n1y_interval_ratio.size(1), s), mode='bilinear', align_corners=True)
+        ny_interval = (m2[..., 1 : 2, :] - m1[..., 1 : 2, :]).unsqueeze(2) * ny_interval_ratio
+        # cumsum to y
+        dy = torch.cumsum(ny_interval, dim=2) + m1[..., 1 : 2, :].unsqueeze(2)
+        dy = dy[..., :-1, 1 : -1]
+        ds = torch.cat((dx, dy), dim=1)
+        gs[..., 1 : -1, 1 : -1] = ds
+        return gs
+
+    def perturb_warp(self, dd):
+        B = dd.size(0)
+        s = self.s
+        # -0.2 to 0.2
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, s), torch.linspace(-1, 1, s))
+        t = torch.stack((ix, iy), dim=0).unsqueeze(0).to('cuda').expand(B, -1, -1, -1)
+
+        tt = t.clone()
+
+        nd = random.randint(0, 4)
+        for ii in range(nd):
+            # define deformation on bd
+            pm = (torch.rand(B, 1) - 0.5) * 0.2
+            ps = (torch.rand(B, 1) - 0.5) * 1.95
+            pt = ps + pm
+            pt = pt.clamp(-0.975, 0.975)
+            # put it on one bd
+            # [1, 1] or [-1, 1] or [-1, -1] etc
+            a1 = (torch.rand(B, 2) > 0.5).float() * 2 -1
+            # select one col for every row
+            a2 = torch.rand(B, 1) > 0.5
+            a2 = torch.cat([a2, a2.bitwise_not()], dim=1)
+            a3 = a1.clone()
+            a3[a2] = ps.view(-1)
+            ps = a3.clone()
+            a3[a2] = pt.view(-1)
+            pt = a3.clone()
+            # 2 N 4
+            bds = torch.stack([
+                t[0, :, 1 : -1, 0], t[0, :, 1 : -1, -1], t[0, :, 0, 1 : -1], t[0, :, -1, 1 : -1]
+            ], dim=2)
+
+            pbd = a2.bitwise_not().float() * a1
+            # id of boundary p is on
+            pbd = torch.abs(0.5 * pbd[:, 0] + 2.5 * pbd[:, 1] + 0.5).long()
+            # ids of other boundaries
+            pbd = torch.stack([pbd + 1, pbd + 2, pbd + 3], dim=1) % 4
+            # print(pbd)
+            pbd = bds[..., pbd].permute(2, 0, 1, 3).reshape(B, 2, -1)
+
+            srcpts = torch.stack([
+                t[..., 0, 0], t[..., 0, -1], t[..., -1, 0], t[..., -1, -1],
+                ps.to('cuda')
+            ], dim=2)
+            srcpts = torch.cat([pbd, srcpts], dim=2).permute(0, 2, 1)
+            dstpts = torch.stack([
+                t[..., 0, 0], t[..., 0, -1], t[..., -1, 0], t[..., -1, -1],
+                pt.to('cuda')
+            ], dim=2)
+            dstpts = torch.cat([pbd, dstpts], dim=2).permute(0, 2, 1)
+            # print(srcpts)
+            # print(dstpts)
+            tgs = self.tpswarper(srcpts, dstpts)
+            tt = F.grid_sample(tt, tgs.permute(0, 2, 3, 1), align_corners=True)
+
+        nd = random.randint(1, 5)
+        for ii in range(nd):
+
+            pm = (torch.rand(B, 2) - 0.5) * 0.2
+            ps = (torch.rand(B, 2) - 0.5) * 1.95
+            pt = ps + pm
+            pt = pt.clamp(-0.975, 0.975)
+
+            srcpts = torch.cat([
+                t[..., -1, :], t[..., 0, :], t[..., 1 : -1, 0], t[..., 1 : -1, -1],
+                ps.unsqueeze(2).to('cuda')
+            ], dim=2).permute(0, 2, 1)
+            dstpts = torch.cat([
+                t[..., -1, :], t[..., 0, :], t[..., 1 : -1, 0], t[..., 1 : -1, -1],
+                pt.unsqueeze(2).to('cuda')
+            ], dim=2).permute(0, 2, 1)
+            tgs = self.tpswarper(srcpts, dstpts)
+            tt = F.grid_sample(tt, tgs.permute(0, 2, 3, 1), align_corners=True)
+        tgs = tt
+
+        # sample tgs to gen invtgs
+        num_sample = 512
+        # n = (H-2)*(W-2)
+        n = s * s
+        idx = torch.randperm(n)
+        idx = idx[:num_sample]
+        srcpts = tgs.reshape(-1, 2, n)[..., idx].permute(0, 2, 1)
+        dstpts = t.reshape(-1, 2, n)[..., idx].permute(0, 2, 1)
+        invtgs = self.tpswarper(srcpts, dstpts)
+        return tgs, invtgs
+
+    def equal_spacing_interpolate(self, gs, s):
+        def equal_bd(x, s):
+            # x is B 2 n
+            v0 = x[..., :-1] # B 2 n-1
+            v = x[..., 1:] - x[..., :-1]
+            vn = v.norm(dim=1, keepdim=True)
+            v = v / vn
+            c = vn.sum(dim=2, keepdim=True) #B 1 1
+            a = vn / c
+            b = torch.cumsum(a, dim=2)
+            b = torch.cat((torch.zeros(B, 1, 1, device='cuda'), b[..., :-1]), dim=2)
+
+            t = torch.linspace(1e-5, 1 - 1e-5, s).view(1, s, 1).to('cuda')
+            t = t - b # B s n-1
+            # print(t)
+
+            tt = torch.cat((t, -torch.ones(B, s, 1, device='cuda')), dim=2) # B s n
+            tt = tt[..., 1:] * tt[..., :-1] # B s n-1
+            tt = (tt < 0).float()
+            d = torch.matmul(v0, tt.permute(0, 2, 1)) + torch.matmul(v, (tt * t).permute(0, 2, 1)) # B 2 s
+            # print(d)
+            return d
+
+        gs = F.interpolate(gs, s, mode='bilinear', align_corners=True)
+        B = gs.size(0)
+        dst_cn = torch.tensor([[-1, -1], [1, -1], [1, 1], [-1, 1]], dtype=torch.float, device='cuda').expand(B, -1, -1)
+        src_cn = torch.stack([gs[..., 0, 0], gs[..., 0, -1], gs[..., -1, -1], gs[..., -1, 0]], dim=2).permute(0, 2, 1)
+        M = self.pspwarper.pspmat(src_cn, dst_cn).detach()
+        invM = self.pspwarper.pspmat(dst_cn, src_cn).detach()
+        pgs = self.pspwarper(gs.permute(0, 2, 3, 1).reshape(B, -1, 2), M).reshape(B, s, s, 2).permute(0, 3, 1, 2)
+        t = [pgs[..., 0, :], pgs[..., -1, :], pgs[..., :, 0], pgs[..., :, -1]]
+        d = []
+        for x in t:
+            d.append(equal_bd(x, s))
+        pgs[..., 0, :] = d[0]
+        pgs[..., -1, :] = d[1]
+        pgs[..., :, 0] = d[2]
+        pgs[..., :, -1] = d[3]
+        gs = self.pspwarper(pgs.permute(0, 2, 3, 1).reshape(B, -1, 2), invM).reshape(B, s, s, 2).permute(0, 3, 1, 2)
+        gs = self.global_post_warp(gs, s)
+        return gs
+
+
+
+class LocalLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def identity_loss(self, gs):
+        s = gs.size(2)
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, s), torch.linspace(-1, 1, s))
+        t = torch.stack((ix, iy), dim=0).unsqueeze(0).to('cuda').expand_as(gs)
+        loss = F.l1_loss(gs, t.detach())
+        return loss
+
+    def direct_loss(self, gs, invtgs):
+        loss = F.l1_loss(gs, invtgs.detach())
+        return loss
+
+    def warp_diff_loss(self, xd, xpd, tgs, invtgs):
+        loss_f = F.l1_loss(xd, F.grid_sample(tgs, xpd.permute(0, 2, 3, 1), align_corners=True).detach())
+        loss_b = F.l1_loss(xpd, F.grid_sample(invtgs, xd.permute(0, 2, 3, 1), align_corners=True).detach())
+        loss = loss_f + loss_b
+        return loss
+
+
+class SupervisedLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        s = 64
+        self.tpswarper = TpsWarp(s)
+
+    def fm2bm(self, fm):
+        # B 3 N N
+        # fm in [0, 1]
+        B, _, s, _ = fm.size()
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, s), torch.linspace(-1, 1, s))
+        t = torch.stack((ix, iy), dim=0).unsqueeze(0).to('cuda').expand(B, -1, -1, -1)
+        srcpts = []
+        dstpts = []
+        for ii in range(B):
+            # mask
+            m = fm[ii, 2]
+            # z s
+            z = torch.nonzero(m, as_tuple=False)
+            num_sample = 512
+            n = z.size(0)
+            # print(n)
+            idx = torch.randperm(n)
+            idx = idx[:num_sample]
+            dstpts.append(t[ii, :, z[idx, 0], z[idx, 1]])
+            srcpts.append(fm[ii, : 2, z[idx, 0], z[idx, 1]] * 2 - 1)
+        srcpts = torch.stack(srcpts, dim=0).permute(0, 2, 1)
+        dstpts = torch.stack(dstpts, dim=0).permute(0, 2, 1)
+        # z = torch.nonzero(torch.abs(srcpts - 0) < 1e-5, as_tuple=False)
+        # print(z.size(0))
+        # print(dstpts.min())
+        # print(dstpts.max())
+        bm = self.tpswarper(srcpts, dstpts)
+        # bm[bm > 1] = 1
+        # bm[bm < -1] = -1
+        return bm
+
+    def gloss(self, x, y):
+        xbd = gs_to_bd(x)
+        # y = self.fm2bm(y)
+        y = F.interpolate(y, 64, mode='bilinear', align_corners=True)
+
+        ybd = gs_to_bd(y).detach()
+        loss = F.l1_loss(xbd, ybd.detach())
+        return loss
+
+    def lloss(self, x, y, dg):
+        # sample tgs to gen invtgs
+        B, _, s, _ = dg.size()
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, s), torch.linspace(-1, 1, s))
+        t = torch.stack((ix, iy), dim=0).unsqueeze(0).to('cuda').expand(B, -1, -1, -1)
+        num_sample = 512
+        # n = (H-2)*(W-2)
+        n = s * s
+        idx = torch.randperm(n)
+        idx = idx[:num_sample]
+        # srcpts = gs_to_bd(tgs)
+        # srcpts = torch.cat([srcpts, tgs[..., 1 : -1, 1 : -1].reshape(-1, 2, n)[..., idx].permute(0, 2, 1)], dim=1)
+        srcpts = dg.reshape(-1, 2, n)[..., idx].permute(0, 2, 1)
+        # dstpts = gs_to_bd(t)
+        # dstpts = torch.cat([dstpts, t[..., 1 : -1, 1 : -1].reshape(-1, 2, n)[..., idx].permute(0, 2, 1)], dim=1)
+        dstpts = t.reshape(-1, 2, n)[..., idx].permute(0, 2, 1)
+        invdg = self.tpswarper(srcpts, dstpts)
+        # compute dl = \phi(dg^-1, y)
+        dl = F.grid_sample(invdg, y.permute(0, 2, 3, 1), align_corners=True)
+        dl = F.interpolate(dl, 64, mode='bilinear', align_corners=True)
+        loss = F.l1_loss(x, dl.detach())
+
+        # y = F.interpolate(y, 64, mode='bilinear', align_corners=True)
+        # loss = F.l1_loss(F.grid_sample(dg.detach(), x.permute(0, 2, 3, 1), align_corners=True), y)
+
+        return loss, dl.detach()

networks/paperedge_cpu.py

@@ -0,0 +1,591 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+# from torch.nn.utils import spectral_norm as SN
+# from torchvision.models.densenet import _DenseBlock
+from .tps_warp import TpsWarp, PspWarp
+from functools import partial
+# import plotly.graph_objects as go
+import random
+import numpy as np
+import cv2
+
+torch.autograd.set_detect_anomaly(True)
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError(
+                'BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError(
+                "Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.actv = nn.ReLU()
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.actv(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.actv(out)
+
+        return out
+
+
+def _make_layer(block, inplanes, planes, blocks, stride=1, dilate=False):
+    norm_layer = nn.BatchNorm2d
+    downsample = None
+
+    if stride != 1 or inplanes != planes * block.expansion:
+        downsample = nn.Sequential(
+            nn.Conv2d(inplanes, planes * block.expansion,
+                      1, stride, bias=False),
+            norm_layer(planes * block.expansion),
+        )
+
+    layers = []
+    layers.append(block(inplanes, planes, stride,
+                  downsample, norm_layer=norm_layer))
+    for _ in range(1, blocks):
+        layers.append(block(planes, planes,
+                            norm_layer=norm_layer))
+
+    return nn.Sequential(*layers)
+
+
+class Interpolate(nn.Module):
+    def __init__(self, size, mode):
+        super(Interpolate, self).__init__()
+        self.interp = nn.functional.interpolate
+        self.size = size
+        self.mode = mode
+
+    def forward(self, x):
+        x = self.interp(x, size=self.size, mode=self.mode, align_corners=True)
+        return x
+
+
+class GlobalWarper(nn.Module):
+    def __init__(self):
+        super(GlobalWarper, self).__init__()
+        modules = [
+            nn.Conv2d(5, 64, kernel_size=7, stride=2, padding=3, bias=False),
+            nn.BatchNorm2d(64),
+            nn.ReLU()
+        ]
+
+        # encoder
+        planes = [64, 128, 256, 256, 512, 512]
+        strides = [2, 2, 2, 2, 2]
+        blocks = [1, 1, 1, 1, 1]
+        for k in range(len(planes) - 1):
+            modules.append(_make_layer(
+                BasicBlock, planes[k], planes[k + 1], blocks[k], strides[k]))
+        self.encoder = nn.Sequential(*modules)
+
+        # decoder
+        modules = []
+        planes = [512, 512, 256, 128, 64]
+        strides = [2, 2, 2, 2]
+        # tsizes = [3, 5, 9, 17, 33]
+        blocks = [1, 1, 1, 1]
+        for k in range(len(planes) - 1):
+            modules += [nn.Sequential(nn.Upsample(scale_factor=strides[k],
+                                                  mode='bilinear',
+                                                  align_corners=True),
+                        _make_layer(BasicBlock, planes[k], planes[k + 1], blocks[k], 1))]
+        self.decoder = nn.Sequential(*modules)
+
+        self.to_warp = nn.Sequential(nn.Conv2d(64, 2, 1))
+        self.to_warp[0].weight.data.fill_(0.0)
+        self.to_warp[0].bias.data.fill_(0.0)
+
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, 256), torch.linspace(-1, 1, 256))
+        self.coord = torch.stack((ix, iy), dim=0).unsqueeze(0)
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, 64),
+                                torch.linspace(-1, 1, 64))
+
+        # note we mulitply a 0.9 so the network is initialized closer to GT.
+        # This is different from localwarper net
+        self.basegrid = torch.stack((ix * 0.9, iy * 0.9), dim=0).unsqueeze(0)
+
+        # # box filter
+        # ksize = 7
+        # p = int((ksize - 1) / 2)
+        # self.pad_replct = partial(F.pad, pad=(p, p, p, p), mode='replicate')
+        # bw = torch.ones(1, 1, ksize, ksize, device='cuda') / ksize / ksize
+        # self.box_filter = partial(F.conv2d, weight=bw)
+
+    def forward(self, im):
+        # print(self.to_warp[0].weight.data)
+        # coordconv
+        B = im.size(0)
+        c = self.coord.expand(B, -1, -1, -1).detach()
+        t = torch.cat((im, c), dim=1)
+
+        t = self.encoder(t)
+        t = self.decoder(t)
+        t = self.to_warp(t)
+
+        gs = t + self.basegrid
+
+        return gs
+
+
+class LocalWarper(nn.Module):
+    def __init__(self):
+        super().__init__()
+        modules = [
+            nn.Conv2d(5, 64, kernel_size=7, stride=2, padding=3, bias=False),
+            nn.BatchNorm2d(64),
+            nn.ReLU()
+        ]
+        # encoder
+        planes = [64, 128, 256, 256, 512, 512]
+        strides = [2, 2, 2, 2, 2]
+        blocks = [1, 1, 1, 1, 1]
+        for k in range(len(planes) - 1):
+            modules.append(_make_layer(
+                BasicBlock, planes[k], planes[k + 1], blocks[k], strides[k]))
+        self.encoder = nn.Sequential(*modules)
+
+        # decoder
+        modules = []
+        planes = [512, 512, 256, 128, 64]
+        strides = [2, 2, 2, 2]
+        # tsizes = [3, 5, 9, 17, 33]
+        blocks = [1, 1, 1, 1]
+        for k in range(len(planes) - 1):
+            modules += [nn.Sequential(nn.Upsample(scale_factor=strides[k], mode='bilinear', align_corners=True),
+                        _make_layer(BasicBlock, planes[k], planes[k + 1], blocks[k], 1))]
+        self.decoder = nn.Sequential(*modules)
+
+        self.to_warp = nn.Sequential(nn.Conv2d(64, 2, 1))
+        self.to_warp[0].weight.data.fill_(0.0)
+        self.to_warp[0].bias.data.fill_(0.0)
+
+        iy, ix = torch.meshgrid(
+            torch.linspace(-1, 1, 256), torch.linspace(-1, 1, 256))
+        self.coord = torch.stack((ix, iy), dim=0).unsqueeze(0)
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, 64),
+                                torch.linspace(-1, 1, 64))
+        self.basegrid = torch.stack((ix, iy), dim=0).unsqueeze(0)
+
+        # box filter
+        ksize = 5
+        p = int((ksize - 1) / 2)
+        self.pad_replct = partial(F.pad, pad=(p, p, p, p), mode='replicate')
+        bw = torch.ones(1, 1, ksize, ksize) / ksize / ksize
+        self.box_filter = partial(F.conv2d, weight=bw)
+
+    def forward(self, im):
+        c = self.coord.expand(im.size(0), -1, -1, -1).detach()
+        t = torch.cat((im, c), dim=1)
+
+        # encoder
+        t = self.encoder(t)
+        t = self.decoder(t)
+        t = self.to_warp(t)
+
+        # # filter
+        # t = self.pad_replct(t)
+        # tx = self.box_filter(t[:, 0 : 1, ...])
+        # ty = self.box_filter(t[:, 1 : 2, ...])
+        # t = torch.cat((tx, ty), dim=1)
+
+        # bd condition
+        t[..., 1, 0, :] = 0
+        t[..., 1, -1, :] = 0
+        t[..., 0, :, 0] = 0
+        t[..., 0, :, -1] = 0
+
+        gs = t + self.basegrid
+        return gs
+
+
+def gs_to_bd(gs):
+    # gs: B 2 H W
+    t = torch.cat([gs[..., 0, :], gs[..., -1, :], gs[..., 1: -1,
+                  0], gs[..., 1: -1, -1]], dim=2).permute(0, 2, 1)
+    # t: B 2(W + H - 1) 2
+    return t
+
+
+class MaskLoss(nn.Module):
+    def __init__(self, gsize):
+        super().__init__()
+        self.tpswarper = TpsWarp(gsize)
+        self.pspwarper = PspWarp()
+        # self.imsize = imsize
+        self.msk = torch.ones(1, 1, gsize, gsize)
+        self.cn = torch.tensor([[-1, -1], [1, -1], [1, 1], [-1, 1]],
+                               dtype=torch.float).unsqueeze(0)
+
+    def forward(self, gs, y, s):
+        # resize gs to s*s
+        B, _, s0, _ = gs.size()
+        tgs = F.interpolate(gs, s, mode='bilinear', align_corners=True)
+
+        # use only the boundary points
+        srcpts = gs_to_bd(tgs)
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, s),
+                                torch.linspace(-1, 1, s))
+        t = torch.stack((ix, iy), dim=0).unsqueeze(0).expand_as(tgs)
+        dstpts = gs_to_bd(t)
+
+        tgs_f = self.tpswarper(srcpts, dstpts.detach())
+        ym = self.msk.expand_as(y)
+        yh = F.grid_sample(ym, tgs_f.permute(0, 2, 3, 1), align_corners=True)
+        loss_f = F.l1_loss(yh, y)
+
+        # forward/backward consistency loss
+        tgs_b = self.tpswarper(dstpts.detach(), srcpts)
+        # tgs_b = F.interpolate(tgs, s0, mode='bilinear', align_corners=True)
+        yy = F.grid_sample(y, tgs_b.permute(0, 2, 3, 1), align_corners=True)
+        loss_b = F.l1_loss(yy, ym)
+
+        return loss_f + loss_b, tgs_f
+
+    def _dist(self, x):
+        # adjacent point distance
+        # B, 2, n
+        x = torch.cat([x[..., 0: 1].detach(), x[..., 1: -1],
+                      x[..., -1:].detach()], dim=2)
+        d = x[..., 1:] - x[..., :-1]
+        return torch.norm(d, dim=1)
+
+# class TVLoss(nn.Module):
+#     def __init__(self):
+#         super(TVLoss, self).__init__()
+
+#     def forward(self, gs):
+#         loss = self._dist(gs[..., 1:], gs[..., :-1]) + self._dist(gs[..., 1:, :], gs[..., :-1, :])
+#         return loss
+
+#     def _dist(self, x1, x0):
+#         d = torch.norm(x1 - x0, dim=1, keepdim=True)
+#         d = torch.abs(d - torch.mean(d, dim=(2, 3), keepdim=True)).mean()
+#         return d
+
+
+class WarperUtil(nn.Module):
+    def __init__(self, imsize):
+        super().__init__()
+        self.tpswarper = TpsWarp(imsize)
+        self.pspwarper = PspWarp()
+        self.s = imsize
+
+    def global_post_warp(self, gs, s):
+        # B, _, s0, _ = gs.size()
+        gs = F.interpolate(gs, s, mode='bilinear', align_corners=True)
+        # gs = F.interpolate(gs, s0, mode='bilinear', align_corners=True)
+        # extract info
+        m1 = gs[..., 0, :]
+        m2 = gs[..., -1, :]
+        n1 = gs[..., 0]
+        n2 = gs[..., -1]
+        # for x
+        m1x_interval_ratio = m1[:, 0, 1:] - m1[:, 0, :-1]
+        m1x_interval_ratio /= m1x_interval_ratio.sum(dim=1, keepdim=True)
+        m2x_interval_ratio = m2[:, 0, 1:] - m2[:, 0, :-1]
+        m2x_interval_ratio /= m2x_interval_ratio.sum(dim=1, keepdim=True)
+        # interpolate all x ratio
+        t = torch.stack(
+            [m1x_interval_ratio, m2x_interval_ratio], dim=1).unsqueeze(1)
+        mx_interval_ratio = F.interpolate(
+            t, (s, m1x_interval_ratio.size(1)), mode='bilinear', align_corners=True)
+        mx_interval = (n2[..., 0: 1, :] - n1[..., 0: 1, :]
+                       ).unsqueeze(3) * mx_interval_ratio
+        # cumsum to x
+        dx = torch.cumsum(mx_interval, dim=3) + n1[..., 0: 1, :].unsqueeze(3)
+        dx = dx[..., 1: -1, :-1]
+        # for y
+        n1y_interval_ratio = n1[:, 1, 1:] - n1[:, 1, :-1]
+        n1y_interval_ratio /= n1y_interval_ratio.sum(dim=1, keepdim=True)
+        n2y_interval_ratio = n2[:, 1, 1:] - n2[:, 1, :-1]
+        n2y_interval_ratio /= n2y_interval_ratio.sum(dim=1, keepdim=True)
+        # interpolate all x ratio
+        t = torch.stack(
+            [n1y_interval_ratio, n2y_interval_ratio], dim=2).unsqueeze(1)
+        ny_interval_ratio = F.interpolate(
+            t, (n1y_interval_ratio.size(1), s), mode='bilinear', align_corners=True)
+        ny_interval = (m2[..., 1: 2, :] - m1[..., 1: 2, :]
+                       ).unsqueeze(2) * ny_interval_ratio
+        # cumsum to y
+        dy = torch.cumsum(ny_interval, dim=2) + m1[..., 1: 2, :].unsqueeze(2)
+        dy = dy[..., :-1, 1: -1]
+        ds = torch.cat((dx, dy), dim=1)
+        gs[..., 1: -1, 1: -1] = ds
+        return gs
+
+    def perturb_warp(self, dd):
+        B = dd.size(0)
+        s = self.s
+        # -0.2 to 0.2
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, s),
+                                torch.linspace(-1, 1, s))
+        t = torch.stack((ix, iy), dim=0).unsqueeze(
+            0).expand(B, -1, -1, -1)
+
+        tt = t.clone()
+
+        nd = random.randint(0, 4)
+        for ii in range(nd):
+            # define deformation on bd
+            pm = (torch.rand(B, 1) - 0.5) * 0.2
+            ps = (torch.rand(B, 1) - 0.5) * 1.95
+            pt = ps + pm
+            pt = pt.clamp(-0.975, 0.975)
+            # put it on one bd
+            # [1, 1] or [-1, 1] or [-1, -1] etc
+            a1 = (torch.rand(B, 2) > 0.5).float() * 2 - 1
+            # select one col for every row
+            a2 = torch.rand(B, 1) > 0.5
+            a2 = torch.cat([a2, a2.bitwise_not()], dim=1)
+            a3 = a1.clone()
+            a3[a2] = ps.view(-1)
+            ps = a3.clone()
+            a3[a2] = pt.view(-1)
+            pt = a3.clone()
+            # 2 N 4
+            bds = torch.stack([
+                t[0, :, 1: -1, 0], t[0, :, 1: -1, -1], t[0,
+                                                         :, 0, 1: -1], t[0, :, -1, 1: -1]
+            ], dim=2)
+
+            pbd = a2.bitwise_not().float() * a1
+            # id of boundary p is on
+            pbd = torch.abs(0.5 * pbd[:, 0] + 2.5 * pbd[:, 1] + 0.5).long()
+            # ids of other boundaries
+            pbd = torch.stack([pbd + 1, pbd + 2, pbd + 3], dim=1) % 4
+            # print(pbd)
+            pbd = bds[..., pbd].permute(2, 0, 1, 3).reshape(B, 2, -1)
+
+            srcpts = torch.stack([
+                t[..., 0, 0], t[..., 0, -1], t[..., -1, 0], t[..., -1, -1],
+                ps
+            ], dim=2)
+            srcpts = torch.cat([pbd, srcpts], dim=2).permute(0, 2, 1)
+            dstpts = torch.stack([
+                t[..., 0, 0], t[..., 0, -1], t[..., -1, 0], t[..., -1, -1],
+                pt
+            ], dim=2)
+            dstpts = torch.cat([pbd, dstpts], dim=2).permute(0, 2, 1)
+            # print(srcpts)
+            # print(dstpts)
+            tgs = self.tpswarper(srcpts, dstpts)
+            tt = F.grid_sample(tt, tgs.permute(0, 2, 3, 1), align_corners=True)
+
+        nd = random.randint(1, 5)
+        for ii in range(nd):
+
+            pm = (torch.rand(B, 2) - 0.5) * 0.2
+            ps = (torch.rand(B, 2) - 0.5) * 1.95
+            pt = ps + pm
+            pt = pt.clamp(-0.975, 0.975)
+
+            srcpts = torch.cat([
+                t[..., -1, :], t[..., 0, :], t[..., 1: -1, 0], t[..., 1: -1, -1],
+                ps.unsqueeze(2)
+            ], dim=2).permute(0, 2, 1)
+            dstpts = torch.cat([
+                t[..., -1, :], t[..., 0, :], t[..., 1: -1, 0], t[..., 1: -1, -1],
+                pt.unsqueeze(2)
+            ], dim=2).permute(0, 2, 1)
+            tgs = self.tpswarper(srcpts, dstpts)
+            tt = F.grid_sample(tt, tgs.permute(0, 2, 3, 1), align_corners=True)
+        tgs = tt
+
+        # sample tgs to gen invtgs
+        num_sample = 512
+        # n = (H-2)*(W-2)
+        n = s * s
+        idx = torch.randperm(n)
+        idx = idx[:num_sample]
+        srcpts = tgs.reshape(-1, 2, n)[..., idx].permute(0, 2, 1)
+        dstpts = t.reshape(-1, 2, n)[..., idx].permute(0, 2, 1)
+        invtgs = self.tpswarper(srcpts, dstpts)
+        return tgs, invtgs
+
+    def equal_spacing_interpolate(self, gs, s):
+        def equal_bd(x, s):
+            # x is B 2 n
+            v0 = x[..., :-1]  # B 2 n-1
+            v = x[..., 1:] - x[..., :-1]
+            vn = v.norm(dim=1, keepdim=True)
+            v = v / vn
+            c = vn.sum(dim=2, keepdim=True)  # B 1 1
+            a = vn / c
+            b = torch.cumsum(a, dim=2)
+            b = torch.cat((torch.zeros(B, 1, 1), b[..., :-1]), dim=2)
+
+            t = torch.linspace(1e-5, 1 - 1e-5, s).view(1, s, 1)
+            t = t - b  # B s n-1
+            # print(t)
+
+            tt = torch.cat((t, -torch.ones(B, s, 1)), dim=2)  # B s n
+            tt = tt[..., 1:] * tt[..., :-1]  # B s n-1
+            tt = (tt < 0).float()
+            d = torch.matmul(v0, tt.permute(0, 2, 1)) + \
+                torch.matmul(v, (tt * t).permute(0, 2, 1))  # B 2 s
+            return d
+
+        gs = F.interpolate(gs, s, mode='bilinear', align_corners=True)
+        B = gs.size(0)
+        dst_cn = torch.tensor([[-1, -1], [1, -1], [1, 1], [-1, 1]],
+                              dtype=torch.float).expand(B, -1, -1)
+        src_cn = torch.stack([gs[..., 0, 0], gs[..., 0, -1],
+                             gs[..., -1, -1], gs[..., -1, 0]], dim=2).permute(0, 2, 1)
+        M = self.pspwarper.pspmat(src_cn, dst_cn).detach()
+        invM = self.pspwarper.pspmat(dst_cn, src_cn).detach()
+        pgs = self.pspwarper(gs.permute(0, 2, 3, 1).reshape(
+            B, -1, 2), M).reshape(B, s, s, 2).permute(0, 3, 1, 2)
+        t = [pgs[..., 0, :], pgs[..., -1, :], pgs[..., :, 0], pgs[..., :, -1]]
+        d = []
+        for x in t:
+            d.append(equal_bd(x, s))
+        pgs[..., 0, :] = d[0]
+        pgs[..., -1, :] = d[1]
+        pgs[..., :, 0] = d[2]
+        pgs[..., :, -1] = d[3]
+        gs = self.pspwarper(pgs.permute(0, 2, 3, 1).reshape(
+            B, -1, 2), invM).reshape(B, s, s, 2).permute(0, 3, 1, 2)
+        gs = self.global_post_warp(gs, s)
+        return gs
+
+
+class LocalLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def identity_loss(self, gs):
+        s = gs.size(2)
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, s),
+                                torch.linspace(-1, 1, s))
+        t = torch.stack((ix, iy), dim=0).unsqueeze(0).expand_as(gs)
+        loss = F.l1_loss(gs, t.detach())
+        return loss
+
+    def direct_loss(self, gs, invtgs):
+        loss = F.l1_loss(gs, invtgs.detach())
+        return loss
+
+    def warp_diff_loss(self, xd, xpd, tgs, invtgs):
+        loss_f = F.l1_loss(xd, F.grid_sample(
+            tgs, xpd.permute(0, 2, 3, 1), align_corners=True).detach())
+        loss_b = F.l1_loss(xpd, F.grid_sample(
+            invtgs, xd.permute(0, 2, 3, 1), align_corners=True).detach())
+        loss = loss_f + loss_b
+        return loss
+
+
+class SupervisedLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        s = 64
+        self.tpswarper = TpsWarp(s)
+
+    def fm2bm(self, fm):
+        # B 3 N N
+        # fm in [0, 1]
+        B, _, s, _ = fm.size()
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, s),
+                                torch.linspace(-1, 1, s))
+        t = torch.stack((ix, iy), dim=0).unsqueeze(
+            0).expand(B, -1, -1, -1)
+        srcpts = []
+        dstpts = []
+        for ii in range(B):
+            # mask
+            m = fm[ii, 2]
+            # z s
+            z = torch.nonzero(m, as_tuple=False)
+            num_sample = 512
+            n = z.size(0)
+            # print(n)
+            idx = torch.randperm(n)
+            idx = idx[:num_sample]
+            dstpts.append(t[ii, :, z[idx, 0], z[idx, 1]])
+            srcpts.append(fm[ii, : 2, z[idx, 0], z[idx, 1]] * 2 - 1)
+        srcpts = torch.stack(srcpts, dim=0).permute(0, 2, 1)
+        dstpts = torch.stack(dstpts, dim=0).permute(0, 2, 1)
+        # z = torch.nonzero(torch.abs(srcpts - 0) < 1e-5, as_tuple=False)
+        # print(z.size(0))
+        # print(dstpts.min())
+        # print(dstpts.max())
+        bm = self.tpswarper(srcpts, dstpts)
+        # bm[bm > 1] = 1
+        # bm[bm < -1] = -1
+        return bm
+
+    def gloss(self, x, y):
+        xbd = gs_to_bd(x)
+        # y = self.fm2bm(y)
+        y = F.interpolate(y, 64, mode='bilinear', align_corners=True)
+
+        ybd = gs_to_bd(y).detach()
+        loss = F.l1_loss(xbd, ybd.detach())
+        return loss
+
+    def lloss(self, x, y, dg):
+        # sample tgs to gen invtgs
+        B, _, s, _ = dg.size()
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, s),
+                                torch.linspace(-1, 1, s))
+        t = torch.stack((ix, iy), dim=0).unsqueeze(
+            0).expand(B, -1, -1, -1)
+        num_sample = 512
+        # n = (H-2)*(W-2)
+        n = s * s
+        idx = torch.randperm(n)
+        idx = idx[:num_sample]
+        # srcpts = gs_to_bd(tgs)
+        # srcpts = torch.cat([srcpts, tgs[..., 1 : -1, 1 : -1].reshape(-1, 2, n)[..., idx].permute(0, 2, 1)], dim=1)
+        srcpts = dg.reshape(-1, 2, n)[..., idx].permute(0, 2, 1)
+        # dstpts = gs_to_bd(t)
+        # dstpts = torch.cat([dstpts, t[..., 1 : -1, 1 : -1].reshape(-1, 2, n)[..., idx].permute(0, 2, 1)], dim=1)
+        dstpts = t.reshape(-1, 2, n)[..., idx].permute(0, 2, 1)
+        invdg = self.tpswarper(srcpts, dstpts)
+        # compute dl = \phi(dg^-1, y)
+        dl = F.grid_sample(invdg, y.permute(0, 2, 3, 1), align_corners=True)
+        dl = F.interpolate(dl, 64, mode='bilinear', align_corners=True)
+        loss = F.l1_loss(x, dl.detach())
+
+        # y = F.interpolate(y, 64, mode='bilinear', align_corners=True)
+        # loss = F.l1_loss(F.grid_sample(dg.detach(), x.permute(0, 2, 3, 1), align_corners=True), y)
+
+        return loss, dl.detach()

networks/tps_warp.py

@@ -0,0 +1,204 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class TpsWarp(nn.Module):
+    def __init__(self, s):
+        super(TpsWarp, self).__init__()
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, s),
+                                torch.linspace(-1, 1, s))
+        self.gs = torch.stack((ix, iy), dim=2).reshape((1, -1, 2))
+        self.sz = s
+
+    def forward(self, src, dst):
+        # src and dst are B.n.2
+        B, n, _ = src.size()
+        # B.n.1.2
+        delta = src.unsqueeze(2)
+        delta = delta - delta.permute(0, 2, 1, 3)
+        # B.n.n
+        K = delta.norm(dim=3)
+        # Rsq = torch.sum(delta**2, dim=3)
+        # Rsq += torch.eye(n)
+        # Rsq[Rsq == 0] = 1.
+        # K = 0.5 * Rsq * torch.log(Rsq)
+        # c = -150
+        # K = torch.exp(c * Rsq)
+        # K = torch.abs(Rsq - 0.5) - 0.5
+        # WARNING: TORCH.SQRT HAS NAN GRAD AT 0
+        # K = torch.sqrt(Rsq)
+        # print(K)
+        # K[torch.isnan(K)] = 0.
+        P = torch.cat((torch.ones((B, n, 1)), src), 2)
+        L = torch.cat((K, P), 2)
+        t = torch.cat(
+            (P.permute(0, 2, 1), torch.zeros((B, 3, 3))), 2)
+        L = torch.cat((L, t), 1)
+        # LInv = L.inverse()
+        # # wv is B.n+3.2
+        # wv = torch.matmul(LInv, torch.cat((dst, torch.zeros((B, 3, 2))), 1))
+        # the above implementation has stability problem near the boundaries
+        wv = torch.solve(
+            torch.cat((dst, torch.zeros((B, 3, 2))), 1), L)[0]
+
+        # get the grid sampler
+        s = self.gs.size(1)
+        gs = self.gs
+        delta = gs.unsqueeze(2)
+        delta = delta - src.unsqueeze(1)
+        K = delta.norm(dim=3)
+        # Rsq = torch.sum(delta**2, dim=3)
+        # K = torch.exp(c * Rsq)
+        # Rsq[Rsq == 0] = 1.
+        # K =  0.5 * Rsq * torch.log(Rsq)
+        # K = torch.abs(Rsq - 0.5) - 0.5
+        # K = torch.sqrt(Rsq)
+        # K[torch.isnan(K)] = 0.
+        gs = gs.expand(B, -1, -1)
+        P = torch.cat((torch.ones((B, s, 1)), gs), 2)
+        L = torch.cat((K, P), 2)
+        gs = torch.matmul(L, wv)
+        return gs.reshape(B, self.sz, self.sz, 2).permute(0, 3, 1, 2)
+
+
+class PspWarp(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def pspmat(self, src, dst):
+        # B, 4, 2
+        B, _, _ = src.size()
+        s = torch.cat([
+            torch.cat([src,
+                       torch.ones((B, 4, 1)),
+                       torch.zeros((B, 4, 3)),
+                       -dst[..., 0: 1] * src[..., 0: 1], -dst[..., 0: 1] * src[..., 1: 2]], dim=2),
+            torch.cat([torch.zeros((B, 4, 3)), src, torch.ones((B, 4, 1)),
+                       -dst[..., 1: 2] * src[..., 0: 1], -dst[..., 1: 2] * src[..., 1: 2]], dim=2)
+        ], dim=1)
+        t = torch.cat([dst[..., 0: 1], dst[..., 1: 2]], dim=1)
+        # M = s.inverse() @ t
+        M = torch.solve(t, s)[0]
+        # M is B 8 1
+        return M
+
+    def forward(self, xy, M):
+        # permute M to B 1 8
+        M = M.permute(0, 2, 1)
+        t = M[..., 6] * xy[..., 0] + M[..., 7] * xy[..., 1] + 1
+        u = (M[..., 0] * xy[..., 0] + M[..., 1] * xy[..., 1] + M[..., 2]) / t
+        v = (M[..., 3] * xy[..., 0] + M[..., 4] * xy[..., 1] + M[..., 5]) / t
+        return torch.stack((u, v), dim=2)
+        # for ii in range(4):
+        #     xy = src[:, ii : ii + 1, :]
+        #     uv = dst[:, ii : ii + 1, :]
+        #     t0 = [xy, torch.ones((B, 1, 1)), torch.zeros((B, 1, 3)), -uv[..., 0] * xy[..., 0], -uv[..., 0] * xy[..., 1]]
+        #     t0 = torch.cat(t0, dim=2)
+        #     t1 = [torch.zeros((B, 1, 3)), xy, torch.ones((B, 1, 1)), -uv[..., 1] * xy[..., 0], -uv[..., 1] * xy[..., 1]]
+        #     t1 = torch.cat(t1, dim=2)
+
+
+class IdwWarp(nn.Module):
+    # inverse distance weighting
+    def __init__(self, s):
+        super().__init__()
+        iy, ix = torch.meshgrid(torch.linspace(-1, 1, s),
+                                torch.linspace(-1, 1, s))
+        self.gs = torch.stack((ix, iy), dim=2).reshape((1, -1, 2)).to('cuda')
+        self.s = s
+
+    def forward(self, src, dst):
+        # B n 2
+        B, n, _ = src.size()
+        # B.n.1.2
+        delta = src.unsqueeze(2)
+        delta = delta - self.gs.unsqueeze(0)
+        # B.n.K
+        p = 1
+        Rsq = torch.sum(delta**2, dim=3)**p
+        w = 1 / Rsq
+        # turn inf to [0...1...0]
+        t = torch.isinf(w)
+        idx = t.any(dim=1).nonzero()
+        w[idx[:, 0], :, idx[:, 1]] = t[idx[:, 0], :, idx[:, 1]].float()
+        wwx = w * dst[..., 0: 1]
+        wwx = wwx.sum(dim=1) / w.sum(dim=1)
+        wwy = w * dst[..., 1: 2]
+        wwy = wwy.sum(dim=1) / w.sum(dim=1)
+        # print(wwy.size())
+        gs = torch.stack((wwx, wwy), dim=2).reshape(
+            B, self.s, self.s, 2).permute(0, 3, 1, 2)
+        return gs
+
+
+if __name__ == "__main__":
+    import cv2
+    import numpy as np
+    from hdf5storage import loadmat
+    from visdom import Visdom
+    vis = Visdom(port=10086)
+
+    # bm_path = '/nfs/bigdisk/sagnik/swat3d/bm/7/2_471_7-ec_Page_375-5LI0001.mat'
+    # img_path = '/nfs/bigdisk/sagnik/swat3d/img/7/2_471_7-ec_Page_375-5LI0001.png'
+
+    # bm = loadmat(bm_path)['bm']
+    # bm = (bm - 224) / 224.
+    # bm = cv2.resize(bm, (64, 64), cv2.INTER_LINEAR).astype(np.float32)
+
+    # im = cv2.imread(img_path) / 255.
+    # im = im[..., ::-1].copy()
+    # im = cv2.resize(im, (256, 256), cv2.INTER_AREA).astype(np.float32)
+    # im = torch.from_numpy(im.transpose(2, 0, 1)).unsqueeze(0).to('cuda')
+
+    # x = np.random.choice(np.arange(64), 50, False)
+    # y = np.random.choice(np.arange(64), 50, False)
+
+    # src = torch.tensor([[x, y]], dtype=torch.float32).permute(0, 2, 1)
+    # src = (src - 32) / 32.
+    # dst = torch.from_numpy(bm[y, x, :]).unsqueeze(0).to('cuda')
+
+    # # print(src.size())
+    # # print(dst.size())
+
+    # tpswarp = TpsWarp(64)
+    # import time
+    # t = time.time()
+    # for _ in range(100):
+    #     gs = tpswarp(src, dst)
+    # print(f'time:{time.time() - t}')
+    # gs = gs.view(-1, 64, 64, 2)
+
+    # print(gs.size())
+    # bm2x2 = F.interpolate(gs.permute(0, 3, 1, 2), size=256, mode='bilinear', align_corners=True).permute(0, 2, 3, 1)
+
+    # rim = F.grid_sample(im, bm2x2, align_corners=True)
+    # vis.images(rim, win='sk3')
+    tpswarp = TpsWarp(16)
+    import matplotlib.pyplot as plt
+    cn = torch.tensor([[-1, -1], [1, -1], [1, 1], [-1, 1], [-0.5, -1],
+                      [0, -1], [0.5, -1]], dtype=torch.float).unsqueeze(0)
+    pn = torch.tensor([[-1, -0.5], [1, -1], [1, 1], [-1, 0.5],
+                      [-0.5, -1], [0, -0.5], [0.5, -1]]).unsqueeze(0)
+    pspwarp = PspWarp()
+    # # print(cn.dtype)
+    M = pspwarp.pspmat(cn[..., 0: 4, :], pn[..., 0: 4, :])
+    invM = pspwarp.pspmat(pn[..., 0: 4, :], cn[..., 0: 4, :])
+    # iy, ix = torch.meshgrid(torch.linspace(-1, 1, 8), torch.linspace(-1, 1, 8))
+    # gs = torch.stack((ix, iy), dim=2).reshape((1, -1, 2)).to('cuda')
+    # t = pspwarp(gs, M).reshape(8, 8, 2).detach().cpu().numpy()
+    # print(M)
+
+    t = tpswarp(cn, pn)
+    from tsdeform import WarperUtil
+    wu = WarperUtil(16)
+    tgs = wu.global_post_warp(t, 16, invM, M)
+
+    t = tgs.permute(0, 2, 3, 1)[0].detach().cpu().numpy()
+
+    plt.clf()
+    plt.pcolormesh(t[..., 0], t[..., 1],
+                   np.zeros_like(t[..., 0]), edgecolors='r')
+    plt.gca().invert_yaxis()
+    plt.gca().axis('equal')
+    vis.matplot(plt, env='grid', win='mpl')

requirements.txt

@@ -0,0 +1,5 @@
+opencv_python
+scipy
+torch
+numpy>=1.20.3
+hdf5storage>=0.1.18

utils/handlers.py

@@ -0,0 +1,84 @@
+import visdom
+import numpy as np
+import csv
+import torch
+from datetime import datetime
+import os
+import cv2
+import random
+import matplotlib.pyplot as plt
+
+
+class VisPlot(object):
+    def __init__(self, port=10086, env='main'):
+        self.vis = visdom.Visdom(port=port)
+        self.env = env
+        self.vis.close('loss', env=env)
+
+    def plot_loss(self, engine, monitor_metrics, win='loss'):
+        self.vis.line(X=np.array([engine.state.iteration]),
+                      # NOTE because we use RunningAverage to log the loss, we can retrieve these numbers from state.metrics
+                      Y=np.array([[engine.state.metrics[x]
+                                 for x in monitor_metrics]]),
+                      env=self.env, win=win, update='append')
+
+    def plot_imgs(self, imgs, win='img', imhistory=False):
+        imgs = np.clip(imgs, 1e-5, 1 - 1e-5)
+        self.vis.images(imgs, env=self.env, win=win, opts={
+                        'caption': win, 'store_history': imhistory})
+
+    def plot_meshes(self, ms, win='ms'):
+        plt.close()
+        n = ms.shape[0]
+        nr = (n - 1) // 8 + 1
+        fig, axs = plt.subplots(1, 2)
+        axs = axs.ravel()
+        # fig.clf()
+
+        c = np.arange(256) / 255.0
+        c = c.reshape((16, 16))
+        for ii in range(2):
+            t = ms[ii]
+            axs[ii].pcolormesh(t[..., 0], t[..., 1], c,
+                               cmap='YlGnBu', edgecolors='black')
+            axs[ii].set_xlim(-1, 1)
+            axs[ii].set_ylim(-1, 1)
+            axs[ii].invert_yaxis()
+            # axs[ii].axis('equal', 'box')
+            axs[ii].set_aspect('equal', 'box')
+        # fig, axs = plt.subplots(1, 2)
+        # axs = axs.ravel()
+        # t = ms[0]
+        # axs[0].pcolormesh(t[..., 0], t[..., 1], np.zeros_like(t[..., 0]), edgecolors='r')
+        # axs[0].invert_yaxis()
+        # axs[0].axis('equal', 'box')
+        fig.tight_layout()
+        self.vis.matplot(fig, env=self.env, win=win)
+
+
+class CSVLogger(object):
+    def __init__(self, filename):
+        self.filename = filename
+
+    def __call__(self, engine, monitor_metrics):
+        with open(self.filename, 'a') as csvfile:
+            writer = csv.writer(csvfile, delimiter=',')
+            date_time = datetime.now().strftime('%m/%d/%Y-%H:%M:%S')
+            writer.writerow([date_time, engine.state.iteration] +
+                            [engine.state.metrics[x] for x in monitor_metrics])
+
+# class SaveRes(object):
+#     def __init__(self, resdir='./'):
+#         self.yp = []
+#         self.resdir = resdir
+
+#     def update(self, engine):
+#         self.yp.append(engine.state.output[0][1].cpu().numpy())
+
+#     def save(self, epoch_id):
+#         self.yp = np.concatenate(self.yp)
+#         savemat(os.path.join(self.resdir, 't{}.mat'.format(epoch_id)), \
+#         {'yp': self.yp})
+#         self.yp = []
+#         # self.yp = []
+#         # self.yg = []