Add files via upload

Author: fh2019ustc

distorted/42_2 copy.png

@@ -0,0 +1,115 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_planes, planes, norm_fn='group', stride=1):
+        super(ResidualBlock, self).__init__()
+  
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
+        self.relu = nn.ReLU(inplace=True)
+
+        num_groups = planes // 8
+
+        if norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            if not stride == 1:
+                self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+        
+        elif norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(planes)
+            self.norm2 = nn.BatchNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.BatchNorm2d(planes)
+        
+        elif norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(planes)
+            self.norm2 = nn.InstanceNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.InstanceNorm2d(planes)
+
+        elif norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+            self.norm2 = nn.Sequential()
+            if not stride == 1:
+                self.norm3 = nn.Sequential()
+
+        if stride == 1:
+            self.downsample = None
+        
+        else:    
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)
+
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.norm1(self.conv1(y)))
+        y = self.relu(self.norm2(self.conv2(y)))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x+y)
+
+
+class BasicEncoder(nn.Module):
+    def __init__(self, input_dim=128, output_dim=128, norm_fn='batch'):
+        super(BasicEncoder, self).__init__()
+        self.norm_fn = norm_fn
+
+        if self.norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
+            
+        elif self.norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(64)
+
+        elif self.norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(64)
+
+        elif self.norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+
+        self.conv1 = nn.Conv2d(input_dim, 64, kernel_size=7, stride=2, padding=3)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.in_planes = 64
+        self.layer1 = self._make_layer(64,  stride=1)
+        self.layer2 = self._make_layer(128, stride=2)
+        self.layer3 = self._make_layer(192, stride=2)
+
+        # output convolution
+        self.conv2 = nn.Conv2d(192, output_dim, kernel_size=1)
+  
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1):
+        layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
+        layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
+        layers = (layer1, layer2)
+        
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+
+        x = self.conv2(x)
+
+        return x

distorted/63_2 copy.png

@@ -0,0 +1,128 @@
+from model import DocGeoNet
+from seg import U2NETP
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import skimage.io as io
+import numpy as np
+import cv2
+import os
+from PIL import Image
+import argparse
+import warnings
+warnings.filterwarnings('ignore')
+
+
+class Net(nn.Module):
+    def __init__(self, opt):
+        super(Net, self).__init__()
+        self.msk = U2NETP(3, 1)
+        self.DocTr = DocGeoNet()
+
+    def forward(self, x):
+        msk, _1,_2,_3,_4,_5,_6 = self.msk(x)
+        msk = (msk > 0.5).float()
+        x = msk * x
+
+        _, _, bm = self.DocTr(x)
+        bm = (2 * (bm / 255.) - 1) * 0.99
+
+        return bm
+
+
+def reload_seg_model(model, path=""):
+    if not bool(path):
+        return model
+    else:
+        model_dict = model.state_dict()
+        pretrained_dict = torch.load(path, map_location='cuda:0')
+        print(len(pretrained_dict.keys()))
+        pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict}
+        print(len(pretrained_dict.keys()))
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+
+        return model
+
+
+def reload_rec_model(model, path=""):
+    if not bool(path):
+        return model
+    else:
+        model_dict = model.state_dict()
+        pretrained_dict = torch.load(path, map_location='cuda:0')
+        print(len(pretrained_dict.keys()))
+        pretrained_dict = {k[7:]: v for k, v in pretrained_dict.items() if k[7:] in model_dict}
+        print(len(pretrained_dict.keys()))
+        model_dict.update(pretrained_dict)
+        model.load_state_dict(model_dict)
+
+        return model
+
+
+def rec(seg_model_path, rec_model_path, distorrted_path, save_path, opt):
+    print(torch.__version__)
+
+    # distorted images list
+    img_list = sorted(os.listdir(distorrted_path))
+
+    # creat save path for rectified images
+    if not os.path.exists(save_path):
+        os.makedirs(save_path)
+
+    net = Net(opt).cuda()
+    print(get_parameter_number(net))
+
+    # reload rec model
+    reload_rec_model(net.DocTr, rec_model_path)
+    reload_seg_model(net.msk, opt.seg_model_path)
+
+    net.eval()
+
+    for img_path in img_list:
+        name = img_path.split('.')[-2]  # image name
+        img_path = distorrted_path + img_path  # image path
+
+        im_ori = np.array(Image.open(img_path))[:, :, :3] / 255.  # read image 0-255 to 0-1
+        h, w, _ = im_ori.shape
+        im = cv2.resize(im_ori, (256, 256))
+        im = im.transpose(2, 0, 1)
+        im = torch.from_numpy(im).float().unsqueeze(0)
+
+        with torch.no_grad():
+            bm = net(im.cuda())
+            bm = bm.cpu()
+
+            # save rectified image
+            bm0 = cv2.resize(bm[0, 0].numpy(), (w, h))  # x flow
+            bm1 = cv2.resize(bm[0, 1].numpy(), (w, h))  # y flow
+            bm0 = cv2.blur(bm0, (3, 3))
+            bm1 = cv2.blur(bm1, (3, 3))
+            lbl = torch.from_numpy(np.stack([bm0, bm1], axis=2)).unsqueeze(0)  # h * w * 2
+            out = F.grid_sample(torch.from_numpy(im_ori).permute(2, 0, 1).unsqueeze(0).float(), lbl, align_corners=True)
+            cv2.imwrite(save_path + name + '_rec' + '.png', ((out[0] * 255).permute(1, 2, 0).numpy())[:,:,::-1].astype(np.uint8))
+
+
+def get_parameter_number(net):
+    total_num = sum(p.numel() for p in net.parameters())
+    trainable_num = sum(p.numel() for p in net.parameters() if p.requires_grad)
+    return {'Total': total_num, 'Trainable': trainable_num}
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--seg_model_path', default='./model_pretrained/preprocess.pth')
+    parser.add_argument('--rec_model_path', default='./model_pretrained/DocGeoNet.pth')
+    parser.add_argument('--distorrted_path', default='./distorted/')
+    parser.add_argument('--save_path', default='./rec/')
+    opt = parser.parse_args()
+
+    rec(seg_model_path=opt.seg_model_path,
+        rec_model_path=opt.rec_model_path,
+        distorrted_path=opt.distorrted_path,
+        save_path=opt.save_path,
+        opt=opt)
+
+if __name__ == "__main__":
+    main()