From 7c854e418db5485b294c7cbd58c2979ca734e506 Mon Sep 17 00:00:00 2001 From: TraceLumen Date: Sat, 12 Sep 2026 03:25:35 +0300 Subject: [PATCH] init --- LICENSE.md | 54 +++++ OCR_eval.py | 43 ++++ README.md | 99 +++++++++ eval.m | 64 ++++++ evalUnwarp.m | 102 +++++++++ extractor.py | 134 ++++++++++++ inference.py | 116 ++++++++++ model.py | 100 +++++++++ ocr_img.txt | 62 ++++++ requirements.txt | 6 + seg.py | 552 +++++++++++++++++++++++++++++++++++++++++++++++ update.py | 106 +++++++++ 12 files changed, 1438 insertions(+) create mode 100644 LICENSE.md create mode 100644 OCR_eval.py create mode 100644 README.md create mode 100644 eval.m create mode 100644 evalUnwarp.m create mode 100644 extractor.py create mode 100644 inference.py create mode 100644 model.py create mode 100644 ocr_img.txt create mode 100644 requirements.txt create mode 100644 seg.py create mode 100644 update.py diff --git a/LICENSE.md b/LICENSE.md new file mode 100644 index 0000000..7ef5a1c --- /dev/null +++ b/LICENSE.md @@ -0,0 +1,54 @@ +# License + +Copyright Β© Hao Feng 2024. All Rights Reserved. + +## 1. Definitions + +1.1 "Algorithm" refers to the deep learning algorithm contained in this repository, including all associated code, documentation, and data. + +1.2 "Author" refers to Hao Feng, the creator and copyright holder of the Algorithm. + +1.3 "Non-Commercial Use" means use for academic research, personal study, or non-profit projects, without any direct or indirect commercial advantage. + +1.4 "Commercial Use" means any use intended for or directed toward commercial advantage or monetary compensation. + +## 2. Grant of Rights + +2.1 Non-Commercial Use: The Author hereby grants you a worldwide, royalty-free, non-exclusive license to use, copy, modify, and distribute the Algorithm for Non-Commercial Use, subject to the conditions in Section 3. + +2.2 Commercial Use: Any Commercial Use of the Algorithm is strictly prohibited without explicit prior written permission from the Author. + +## 3. Conditions + +3.1 For Non-Commercial Use: + a) Attribution: You must give appropriate credit to the Author, provide a link to this license, and indicate if changes were made. + b) Share-Alike: If you modify, transform, or build upon the Algorithm, you must distribute your contributions under the same license as this one. + c) No additional restrictions: You may not apply legal terms or technological measures that legally restrict others from doing anything this license permits. + +3.2 For Commercial Use: + a) Prior Contact: Before any Commercial Use, you must contact the Author at haof@mail.ustc.edu.cn and obtain explicit written permission. + b) Separate Agreement: Commercial Use terms will be stipulated in a separate commercial license agreement. + +## 4. Disclaimer of Warranty + +The Algorithm 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 non-infringement. In no event shall the Author be liable for any claim, damages, or other liability arising from, out of, or in connection with the Algorithm or the use or other dealings in the Algorithm. + +## 5. Limitation of Liability + +In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, shall the Author be liable to you for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this license or out of the use or inability to use the Algorithm. + +## 6. Termination + +6.1 This license and the rights granted hereunder will terminate automatically upon any breach by you of the terms of this license. + +6.2 All sections which by their nature should survive the termination of this license shall survive such termination. + +## 7. Miscellaneous + +7.1 If any provision of this license is held to be unenforceable, such provision shall be reformed only to the extent necessary to make it enforceable. + +7.2 This license represents the complete agreement concerning the subject matter hereof. + +By using the Algorithm, you acknowledge that you have read this license, understand it, and agree to be bound by its terms and conditions. If you do not agree to the terms and conditions of this license, do not use, modify, or distribute the Algorithm. + +For permissions beyond the scope of this license, please contact the Author at haof@mail.ustc.edu.cn. diff --git a/OCR_eval.py b/OCR_eval.py new file mode 100644 index 0000000..e77efc7 --- /dev/null +++ b/OCR_eval.py @@ -0,0 +1,43 @@ +def Levenshtein_Distance(str1, str2): + matrix = [[ i + j for j in range(len(str2) + 1)] for i in range(len(str1) + 1)] + for i in range(1, len(str1)+1): + for j in range(1, len(str2)+1): + if(str1[i-1] == str2[j-1]): + d = 0 + else: + d = 1 + matrix[i][j] = min(matrix[i-1][j]+1, matrix[i][j-1]+1, matrix[i-1][j-1]+d) + + return matrix[len(str1)][len(str2)] + +def cal_cer_ed(path_ours, tail='_rec'): + path_gt='./GT/' + N=66 + cer1=[] + cer2=[] + ed1=[] + ed2=[] + check=[0 for _ in range(N+1)] + lis=[1,2,3,4,5,6,7,9,10,21,22,23,24,27,30,31,32,36,38,40,41,44,45,46,47,48,50,51,52,53] # DocTr (Setting 1) + # lis=[1,9,10,12,19,20,21,22,23,24,30,31,32,34,35,36,37,38,39,40,44,45,46,47,49] # DewarpNet (Setting 2) + for i in range(1,N): + if i not in lis: + continue + gt=Image.open(path_gt+str(i)+'.png') + img1=Image.open(path_ours+str(i)+'_1' + tail) + img2=Image.open(path_ours+str(i)+'_2' + tail) + content_gt=pytesseract.image_to_string(gt) + content1=pytesseract.image_to_string(img1) + content2=pytesseract.image_to_string(img2) + l1=Levenshtein_Distance(content_gt,content1) + l2=Levenshtein_Distance(content_gt,content2) + ed1.append(l1) + ed2.append(l2) + cer1.append(l1/len(content_gt)) + cer2.append(l2/len(content_gt)) + check[i]=cer1[-1] + print('CER: ', (np.mean(cer1)+np.mean(cer2)) / 2.) + print('ED: ', (np.mean(ed1)+np.mean(ed2)) / 2.) + +def evalu(path_ours, tail): + cal_cer_ed(path_ours, tail) diff --git a/README.md b/README.md new file mode 100644 index 0000000..60fc6fa --- /dev/null +++ b/README.md @@ -0,0 +1,99 @@ +πŸ”₯ ***2025.3.24:*** **Good news! Our work has been accepted by International Journal of Computer Vision (IJCV).** + +πŸ”₯ ***2024.4.28:*** **Good news! The code and pre-trained model of DocScanner are now released!** + +πŸš€ **Good news! The [online demo](https://docai.doctrp.top:20443/) for DocScanner is now live, allowing for easy image upload and correction.** + +πŸ”₯ **Good news! Our new work [DocTr++: Deep Unrestricted Document Image Rectification](https://github.com/fh2019ustc/DocTr-Plus) comes out, capable of rectifying various distorted document images in the wild.** + +πŸ”₯ **Good news! A comprehensive list of [Awesome Document Image Rectification](https://github.com/fh2019ustc/Awesome-Document-Image-Rectification) methods is available.** + +# DocScanner + +

+ + +

+ + +This is a PyTorch/GPU re-implementation of the paper [DocScanner: Robust Document Image Rectification with Progressive Learning](https://drive.google.com/file/d/1mmCUj90rHyuO1SmpLt361youh-07Y0sD/view?usp=share_link). + +![image](https://user-images.githubusercontent.com/50725551/209266364-aee68a88-090d-4f21-919a-092f19570d86.png) + + +## πŸš€ Demo [(Link)](https://docai.doctrp.top:20443/) +***Note***：The model version used in the demo corresponds to ***"DocScanner-L"*** as described in the paper. +1. Upload the distorted document image to be rectified in the left box. +2. Click the "Submit" button. +3. The rectified image will be displayed in the right box. + +image + +### Examples +![image](https://user-images.githubusercontent.com/50725551/223947040-eac8389c-bed8-433d-b23b-679c926fba8f.png) +![image](https://user-images.githubusercontent.com/50725551/223946953-3a46d6a3-4361-41ef-bb5c-f235392e1f88.png) + + +## Training +- We train the **Document Localization Module** using the [Doc3D](https://github.com/fh2019ustc/doc3D-dataset) dataset. Besides, [DTD](https://www.robots.ox.ac.uk/~vgg/data/dtd/) dataset is exploited for background data enhancement. +- We train the **Progressive Rectification Module** using the [Doc3D](https://github.com/fh2019ustc/doc3D-dataset) dataset. Here we use the background-excluded document images for training. + +## Inference +1. Put the [pre-trained DocScanner-L](https://drive.google.com/drive/folders/1W1_DJU8dfEh6FqDYqFQ7ypR38Z8c5r4D?usp=sharing) to `$ROOT/model_pretrained/`. +2. Put the distorted images in `$ROOT/distorted/`. +3. Run the script and the rectified images are saved in `$ROOT/rectified/` by default. + ``` + python inference.py + ``` + +## Evaluation +- ***Important.*** In the [DocUNet Benchmark](https://www3.cs.stonybrook.edu/~cvl/docunet.html), the '64_1.png' and '64_2.png' distorted images are rotated by 180 degrees, which do not match the GT documents. It is ignored by most of the existing works. Before the evaluation, please make a check. Note that the performances in most of the existing work are computed with these two ***mistaken*** samples. +- For reproducing the following quantitative performance on the ***corrected*** [DocUNet Benchmark](https://www3.cs.stonybrook.edu/~cvl/docunet.html), please use the geometric rectified images available from [Google Drive](https://drive.google.com/drive/folders/1QBe26xJwIl38sWqK2ZE9ke5nu0Mpr4dW?usp=sharing). For the ***corrected*** performance of [other methods](https://github.com/fh2019ustc/Awesome-Document-Image-Rectification), please refer to the paper [DocScanner](https://arxiv.org/pdf/2110.14968v2.pdf). +- ***Image Metrics:*** We use the same evaluation code for MS-SSIM and LD as [DocUNet Benchmark](https://www3.cs.stonybrook.edu/~cvl/docunet.html) dataset based on Matlab 2019a. Please compare the scores according to your Matlab version. We provide our Matlab interface file at ```$ROOT/ssim_ld_eval.m```. +- ***OCR Metrics:*** The index of 30 documents (60 images) of [DocUNet Benchmark](https://www3.cs.stonybrook.edu/~cvl/docunet.html) used for our OCR evaluation is ```$ROOT/ocr_img.txt``` (*Setting 1*). Please refer to [DewarpNet](https://github.com/cvlab-stonybrook/DewarpNet) for the index of 25 documents (50 images) of [DocUNet Benchmark](https://www3.cs.stonybrook.edu/~cvl/docunet.html) used for their OCR evaluation (*Setting 2*). We provide the OCR evaluation code at ```$ROOT/OCR_eval.py```. The version of pytesseract is 0.3.8, and the version of [Tesseract](https://digi.bib.uni-mannheim.de/tesseract/) in Windows is recent 5.0.1.20220118. Note that in different operating systems, the calculated performance has slight differences. +- ***W_v and W_h Index:*** The layout results of [DocUNet Benchmark](https://www3.cs.stonybrook.edu/~cvl/docunet.html) is available at [Google Drive](https://drive.google.com/drive/folders/1PcfWIowjM0AVKhZrRwGChM-2VAcUwWrF?usp=sharing). + +| Method | MS-SSIM | LD | Li-D | ED (*Setting 1*) | CER | ED (*Setting 2*) | CER | Para. (M) | +|:-----------------------:|:------------:|:-----------:| :-------:|:----------------:|:--------------:|:---------------------:|:--------------:|:--------------:| +| *DocScanner-T* | 0.5123 | 7.92 | 2.04 | 501.82 | 0.1823 | 809.46 | 0.2068 | 2.6 | +| *DocScanner-B* | 0.5134 | 7.62 | 1.88 | 434.11 | 0.1652 | 671.48 | 0.1789 | 5.2 | +| *DocScanner-L* | 0.5178 | 7.45 | 1.86 | 390.43 | 0.1486 | 632.34 | 0.1648 | 8.5 | + +## Citation +Please cite the related works in your publications if it helps your research: + +``` +@inproceedings{feng2021doctr, + title={DocTr: Document Image Transformer for Geometric Unwarping and Illumination Correction}, + author={Feng, Hao and Wang, Yuechen and Zhou, Wengang and Deng, Jiajun and Li, Houqiang}, + booktitle={Proceedings of the 29th ACM International Conference on Multimedia}, + pages={273--281}, + year={2021} +} +``` + +``` +@inproceedings{feng2022docgeonet, + title={Geometric Representation Learning for Document Image Rectification}, + author={Feng, Hao and Zhou, Wengang and Deng, Jiajun and Wang, Yuechen and Li, Houqiang}, + booktitle={Proceedings of the European Conference on Computer Vision}, + year={2022} +} +``` + +``` +@article{feng2025docscanner, + title={DocScanner: Robust document image rectification with progressive learning}, + author={Feng, Hao and Zhou, Wengang and Deng, Jiajun and Tian, Qi and Li, Houqiang}, + journal={International Journal of Computer Vision}, + pages={1--20}, + year={2025} +} +``` + +## Acknowledgement +The codes are largely based on [DocUNet](https://www3.cs.stonybrook.edu/~cvl/docunet.html) and [DewarpNet](https://github.com/cvlab-stonybrook/DewarpNet). Thanks for their wonderful works. + +## Contact +For commercial usage, please contact Hao Feng ([haof@mail.ustc.edu.cn](haof@mail.ustc.edu.cn)). + diff --git a/eval.m b/eval.m new file mode 100644 index 0000000..0189322 --- /dev/null +++ b/eval.m @@ -0,0 +1,64 @@ +path_rec = "xxx"; % rectified image path +path_scan = './scan/'; % scan image path +label_path = './layout/'; % layout result path + +tarea = 598400; +ms1 = 0; +ld1 = 0; +lid1 = 0; +ms2 = 0; +ld2 = 0; +lid2 = 0; +wv = 0; +wh = 0; + +sprintf(path_rec) +for i=1:65 + path_rec_1 = sprintf("%s%d%s", path_rec, i, '_1 copy_rec.png'); % rectified image path + path_rec_2 = sprintf("%s%d%s", path_rec, i, '_2 copy_rec.png'); % rectified image path + path_scan_new = sprintf("%s%d%s", path_scan, i, '.png'); % corresponding scan image path + bbox_i_path = sprintf("%s%d%s", label_path, i, '.txt'); % corresponding layout txt path + + % imread and rgb2gray + A1 = imread(path_rec_1); + A2 = imread(path_rec_2); + +% if i == 64 +% A1 = rot90(A1,-2); +% A2 = rot90(A2,-2); +% end + + ref = imread(path_scan_new); + A1 = rgb2gray(A1); + A2 = rgb2gray(A2); + ref = rgb2gray(ref); + bbox_i = read_txt(bbox_i_path); + bbox_i = bbox_i + 1; % python index starts from 0 + + % resize + b = sqrt(tarea/size(ref,1)/size(ref,2)); + ref = imresize(ref,b); + A1 = imresize(A1,[size(ref,1),size(ref,2)]); + A2 = imresize(A2,[size(ref,1),size(ref,2)]); + scaled_bbox_i = bbox_i * b * 0.5; + scaled_bbox_i = round(scaled_bbox_i); + scaled_bbox_i = max(scaled_bbox_i, 1); + + % calculate + [ms_1, ld_1, lid_1, W_v_1, W_h_1] = evalUnwarp(A1, ref, scaled_bbox_i); + [ms_2, ld_2, lid_2, W_v_2, W_h_2] = evalUnwarp(A2, ref, scaled_bbox_i); + ms1 = ms1 + ms_1; + ms2 = ms2 + ms_2; + ld1 = ld1 + ld_1; + ld2 = ld2 + ld_2; + lid1 = lid1 + lid_1; + lid2 = lid2 + lid_2; + wv = wv + W_v_1 + W_v_2; + wh = wh + W_h_1 + W_h_2; +end + +ms = (ms1 + ms2) / 130 % MS-SSIM +ld = (ld1 + ld2) / 130 % local distortion +li_d = (lid1 + lid2) / 130 % line distortion +wv = wv / 130 % wv index +wh = wh / 130 % wh index diff --git a/evalUnwarp.m b/evalUnwarp.m new file mode 100644 index 0000000..6f84634 --- /dev/null +++ b/evalUnwarp.m @@ -0,0 +1,102 @@ +function [ms, ld, li_d, wv, wh] = evalUnwarp(A, ref, data) +%EVALUNWARP compute MSSSIM and LD between the unwarped image and the scan +% A: unwarped image +% ref: reference image, the scan image +% ms: returned MS-SSIM value +% ld: returned local distortion value +% Matlab image processing toolbox is necessary to compute ssim. The weights +% for multi-scale ssim is directly adopted from: +% +% Wang, Zhou, Eero P. Simoncelli, and Alan C. Bovik. "Multiscale structural +% similarity for image quality assessment." In Signals, Systems and Computers, +% 2004. Conference Record of the Thirty-Seventh Asilomar Conference on, 2003. +% +% Local distortion relies on the paper: +% Liu, Ce, Jenny Yuen, and Antonio Torralba. "Sift flow: Dense correspondence +% across scenes and its applications." In PAMI, 2010. +% +% and its implementation: +% https://people.csail.mit.edu/celiu/SIFTflow/ + +x = A; +y = ref; + +im1=imresize(imfilter(y,fspecial('gaussian',7,1.),'same','replicate'),0.5,'bicubic'); +im2=imresize(imfilter(x,fspecial('gaussian',7,1.),'same','replicate'),0.5,'bicubic'); + +im1=im2double(im1); +im2=im2double(im2); + +cellsize=3; +gridspacing=1; + +sift1 = mexDenseSIFT(im1,cellsize,gridspacing); +sift2 = mexDenseSIFT(im2,cellsize,gridspacing); + +SIFTflowpara.alpha=2*255; +SIFTflowpara.d=40*255; +SIFTflowpara.gamma=0.005*255; +SIFTflowpara.nlevels=4; +SIFTflowpara.wsize=2; +SIFTflowpara.topwsize=10; +SIFTflowpara.nTopIterations = 60; +SIFTflowpara.nIterations= 30; + + +[vx,vy,~]=SIFTflowc2f(sift1,sift2,SIFTflowpara); + +rows1p = size(im1,1); +cols1p = size(im1,2); + +% Li-D +rowstd_sum = 0; +for i = 1:rows1p + rowstd = std(vy(i, :),1); + rowstd_sum = rowstd_sum + rowstd; +end +rowstd_mean = rowstd_sum / rows1p; + +colstd_sum = 0; +for i = 1:cols1p + colstd = std(vx(:, i),1); + colstd_sum = colstd_sum + colstd; +end +colstd_mean = colstd_sum / cols1p; + +li_d = (rowstd_mean + colstd_mean) / 2; + + +% LD +d = sqrt(vx.^2 + vy.^2); +ld = mean(d(:)); + + +% MS-SSIM +wt = [0.0448 0.2856 0.3001 0.2363 0.1333]; +ss = zeros(5, 1); +for s = 1 : 5 + ss(s) = ssim(x, y); + x = impyramid(x, 'reduce'); + y = impyramid(y, 'reduce'); +end +ms = wt * ss; + + +% wv and wh +rowstd_sum = 0; +for i = 1:size(data, 1) + rowstd_top = std(vy(data(i,2), data(i,1):data(i,3)),1) / (data(i,3)-data(i,1)); + rowstd_bot = std(vy(data(i,4), data(i,1):data(i,3)),1) / (data(i,3)-data(i,1)); + rowstd_sum = rowstd_sum + rowstd_top + rowstd_bot; +end +wv = rowstd_sum / (2 * size(data, 1)); + +colstd_sum = 0; +for i = 1:size(data, 1) + colstd_left = std(vx(data(i,2):data(i,4), data(i,1)),1) / (data(i,4)- data(i,2)); + colstd_right = std(vx(data(i,2):data(i,4), data(i,3)),1) / (data(i,4)- data(i,2)); + colstd_sum = colstd_sum + colstd_left + colstd_right; +end +wh = colstd_sum / (2 * size(data, 1)); + +end diff --git a/extractor.py b/extractor.py new file mode 100644 index 0000000..aaa8006 --- /dev/null +++ b/extractor.py @@ -0,0 +1,134 @@ +import torch.nn as nn + + +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 == '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) + + 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 BottleneckBlock(nn.Module): + def __init__(self, in_planes, planes, norm_fn='group', stride=1): + super(BottleneckBlock, self).__init__() + + self.conv1 = nn.Conv2d(in_planes, planes // 4, kernel_size=1, padding=0) + self.conv2 = nn.Conv2d(planes // 4, planes // 4, kernel_size=3, padding=1, stride=stride) + self.conv3 = nn.Conv2d(planes // 4, planes, kernel_size=1, padding=0) + self.relu = nn.ReLU(inplace=True) + + if norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(planes // 4) + self.norm2 = nn.BatchNorm2d(planes // 4) + self.norm3 = nn.BatchNorm2d(planes) + if not stride == 1: + self.norm4 = nn.BatchNorm2d(planes) + + elif norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(planes // 4) + self.norm2 = nn.InstanceNorm2d(planes // 4) + self.norm3 = nn.InstanceNorm2d(planes) + if not stride == 1: + self.norm4 = nn.InstanceNorm2d(planes) + + if stride == 1: + self.downsample = None + else: + self.downsample = nn.Sequential( + nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4) + + def forward(self, x): + y = x + y = self.relu(self.norm1(self.conv1(y))) + y = self.relu(self.norm2(self.conv2(y))) + y = self.relu(self.norm3(self.conv3(y))) + + if self.downsample is not None: + x = self.downsample(x) + + return self.relu(x + y) + + +class BasicEncoder(nn.Module): + def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0): + super(BasicEncoder, self).__init__() + self.norm_fn = norm_fn + + if self.norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(64) + + elif self.norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(64) + + self.conv1 = nn.Conv2d(3, 80, kernel_size=7, stride=2, padding=3) + self.relu1 = nn.ReLU(inplace=True) + + self.in_planes = 80 + self.layer1 = self._make_layer(80, stride=1) + self.layer2 = self._make_layer(160, stride=2) + self.layer3 = self._make_layer(240, stride=2) + + # output convolution + self.conv2 = nn.Conv2d(240, 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 diff --git a/inference.py b/inference.py new file mode 100644 index 0000000..c04a144 --- /dev/null +++ b/inference.py @@ -0,0 +1,116 @@ +from model import DocScanner +from seg import U2NETP + +import torch +import torch.nn as nn +import torch.nn.functional as F +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): + super(Net, self).__init__() + self.msk = U2NETP(3, 1) + self.bm = DocScanner() # 矫正 + + def forward(self, x): + msk, _1,_2,_3,_4,_5,_6 = self.msk(x) + msk = (msk > 0.5).float() + x = msk * x + + bm = self.bm(x, iters=12, test_mode=True) + bm = (2 * (bm / 286.8) - 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') + pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict} + 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') + pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} + 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): + # distorted images list + img_list = os.listdir(distorrted_path) + + # creat save path for rectified images + if not os.path.exists(save_path): + os.makedirs(save_path) + + # net init + net = Net().cuda() + # reload seg model + reload_seg_model(net.msk, seg_model_path) + # reload rec model + reload_rec_model(net.bm, rec_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. + h, w, _ = im_ori.shape + im = cv2.resize(im_ori, (288, 288)) + 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 main(): + parser = argparse.ArgumentParser() + parser.add_argument('--seg_model_path', default='./model_pretrained/seg.pth') + parser.add_argument('--rec_model_path', default='./model_pretrained/DocScanner-L.pth') + parser.add_argument('--distorrted_path', default='./distorted/') + parser.add_argument('--rectified_path', default='./rectified/') + 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.rectified_path) + + +if __name__ == "__main__": + main() diff --git a/model.py b/model.py new file mode 100644 index 0000000..1cf5800 --- /dev/null +++ b/model.py @@ -0,0 +1,100 @@ +from update import BasicUpdateBlock +from extractor import BasicEncoder + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def bilinear_sampler(img, coords, mode='bilinear', mask=False): + """ Wrapper for grid_sample, uses pixel coordinates """ + H, W = img.shape[-2:] + xgrid, ygrid = coords.split([1, 1], dim=-1) + xgrid = 2 * xgrid / (W - 1) - 1 + ygrid = 2 * ygrid / (H - 1) - 1 + + grid = torch.cat([xgrid, ygrid], dim=-1) + img = F.grid_sample(img, grid, align_corners=True) + if mask: + mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) + return img, mask.float() + + return img + + +def coords_grid(batch, ht, wd): + coords = torch.meshgrid(torch.arange(ht), torch.arange(wd)) + coords = torch.stack(coords[::-1], dim=0).float() + return coords[None].repeat(batch, 1, 1, 1) + + +class DocScanner(nn.Module): + def __init__(self): + super(DocScanner, self).__init__() + + self.hidden_dim = hdim = 160 + self.context_dim = 160 + + self.fnet = BasicEncoder(output_dim=320, norm_fn='instance') + self.update_block = BasicUpdateBlock(hidden_dim=hdim) + + def freeze_bn(self): + for m in self.modules(): + if isinstance(m, nn.BatchNorm2d): + m.eval() + + def initialize_flow(self, img): + N, C, H, W = img.shape + coodslar = coords_grid(N, H, W).to(img.device) + coords0 = coords_grid(N, H // 8, W // 8).to(img.device) + coords1 = coords_grid(N, H // 8, W // 8).to(img.device) + + return coodslar, coords0, coords1 + + def upsample_flow(self, flow, mask): + N, _, H, W = flow.shape + mask = mask.view(N, 1, 9, 8, 8, H, W) + mask = torch.softmax(mask, dim=2) + + up_flow = F.unfold(8 * flow, [3, 3], padding=1) + up_flow = up_flow.view(N, 2, 9, 1, 1, H, W) + + up_flow = torch.sum(mask * up_flow, dim=2) + up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) + + return up_flow.reshape(N, 2, 8 * H, 8 * W) + + def forward(self, image1, iters=12, flow_init=None, test_mode=False): + image1 = image1.contiguous() + + fmap1 = self.fnet(image1) + + warpfea = fmap1 + + net, inp = torch.split(fmap1, [160, 160], dim=1) + net = torch.tanh(net) + inp = torch.relu(inp) + + coodslar, coords0, coords1 = self.initialize_flow(image1) + + if flow_init is not None: + coords1 = coords1 + flow_init + + flow_predictions = [] + for itr in range(iters): + coords1 = coords1.detach() + flow = coords1 - coords0 + + net, up_mask, delta_flow = self.update_block(net, inp, warpfea, flow) + + coords1 = coords1 + delta_flow + flow_up = self.upsample_flow(coords1 - coords0, up_mask) + bm_up = coodslar + flow_up + + warpfea = bilinear_sampler(fmap1, coords1.permute(0, 2, 3, 1)) + flow_predictions.append(bm_up) + + if test_mode: + return bm_up + + return flow_predictions \ No newline at end of file diff --git a/ocr_img.txt b/ocr_img.txt new file mode 100644 index 0000000..fe2211d --- /dev/null +++ b/ocr_img.txt @@ -0,0 +1,62 @@ +The images for OCR evaluation of DocUNet Benchmark. +# Setting 1 (Setting from DocTr) +# Total 30 * 2 = 60 images. +./scan/1.png +./scan/2.png +./scan/3.png +./scan/4.png +./scan/5.png +./scan/6.png +./scan/7.png +./scan/9.png +./scan/10.png +./scan/21.png +./scan/22.png +./scan/23.png +./scan/24.png +./scan/27.png +./scan/30.png +./scan/31.png +./scan/32.png +./scan/36.png +./scan/38.png +./scan/40.png +./scan/41.png +./scan/44.png +./scan/45.png +./scan/46.png +./scan/47.png +./scan/48.png +./scan/50.png +./scan/51.png +./scan/52.png +./scan/53.png + +# Setting 2 (Setting from DewarpNet) +# Link: https://github.com/cvlab-stonybrook/DewarpNet/blob/master/eval/ocr_eval/ocr_files.txt +# Total 25 * 2 = 50 images. +./scan/1.png +./scan/9.png +./scan/10.png +./scan/12.png +./scan/19.png +./scan/20.png +./scan/21.png +./scan/22.png +./scan/23.png +./scan/24.png +./scan/30.png +./scan/31.png +./scan/32.png +./scan/34.png +./scan/35.png +./scan/36.png +./scan/37.png +./scan/38.png +./scan/39.png +./scan/40.png +./scan/44.png +./scan/45.png +./scan/46.png +./scan/47.png +./scan/49.png diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..e4fc19a --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +numpy==1.19.0 +opencv_python==4.2.0.34 +Pillow==9.4.0 +scikit_image==0.17.2 +skimage==0.0 +torch==1.5.1+cu101 diff --git a/seg.py b/seg.py new file mode 100644 index 0000000..1e4e861 --- /dev/null +++ b/seg.py @@ -0,0 +1,552 @@ +import torch +import torch.nn as nn +from torchvision import models +import torch.nn.functional as F +import numpy as np + + +class sobel_net(nn.Module): + def __init__(self): + super().__init__() + self.conv_opx = nn.Conv2d(1, 1, 3, bias=False) + self.conv_opy = nn.Conv2d(1, 1, 3, bias=False) + sobel_kernelx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype='float32').reshape((1, 1, 3, 3)) + sobel_kernely = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype='float32').reshape((1, 1, 3, 3)) + self.conv_opx.weight.data = torch.from_numpy(sobel_kernelx) + self.conv_opy.weight.data = torch.from_numpy(sobel_kernely) + + for p in self.parameters(): + p.requires_grad = False + + def forward(self, im): # input rgb + x = (0.299 * im[:, 0, :, :] + 0.587 * im[:, 1, :, :] + 0.114 * im[:, 2, :, :]).unsqueeze(1) # rgb2gray + gradx = self.conv_opx(x) + grady = self.conv_opy(x) + + x = (gradx ** 2 + grady ** 2) ** 0.5 + x = (x - x.min()) / (x.max() - x.min()) + x = F.pad(x, (1, 1, 1, 1)) + + x = torch.cat([im, x], dim=1) + return x + + +class REBNCONV(nn.Module): + def __init__(self, in_ch=3, out_ch=3, dirate=1): + super(REBNCONV, self).__init__() + + self.conv_s1 = nn.Conv2d(in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate) + self.bn_s1 = nn.BatchNorm2d(out_ch) + self.relu_s1 = nn.ReLU(inplace=True) + + def forward(self, x): + hx = x + xout = self.relu_s1(self.bn_s1(self.conv_s1(hx))) + + return xout + + +## upsample tensor 'src' to have the same spatial size with tensor 'tar' +def _upsample_like(src, tar): + src = F.interpolate(src, size=tar.shape[2:], mode='bilinear', align_corners=False) + + return src + + +### RSU-7 ### +class RSU7(nn.Module): # UNet07DRES(nn.Module): + + def __init__(self, in_ch=3, mid_ch=12, out_ch=3): + super(RSU7, self).__init__() + + self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1) + + self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1) + self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1) + self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1) + self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1) + self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1) + self.pool5 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=1) + + self.rebnconv7 = REBNCONV(mid_ch, mid_ch, dirate=2) + + self.rebnconv6d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1) + + def forward(self, x): + hx = x + hxin = self.rebnconvin(hx) + + hx1 = self.rebnconv1(hxin) + hx = self.pool1(hx1) + + hx2 = self.rebnconv2(hx) + hx = self.pool2(hx2) + + hx3 = self.rebnconv3(hx) + hx = self.pool3(hx3) + + hx4 = self.rebnconv4(hx) + hx = self.pool4(hx4) + + hx5 = self.rebnconv5(hx) + hx = self.pool5(hx5) + + hx6 = self.rebnconv6(hx) + + hx7 = self.rebnconv7(hx6) + + hx6d = self.rebnconv6d(torch.cat((hx7, hx6), 1)) + hx6dup = _upsample_like(hx6d, hx5) + + hx5d = self.rebnconv5d(torch.cat((hx6dup, hx5), 1)) + hx5dup = _upsample_like(hx5d, hx4) + + hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1)) + hx4dup = _upsample_like(hx4d, hx3) + + hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1)) + hx3dup = _upsample_like(hx3d, hx2) + + hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1)) + hx2dup = _upsample_like(hx2d, hx1) + + hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1)) + + return hx1d + hxin + + +### RSU-6 ### +class RSU6(nn.Module): # UNet06DRES(nn.Module): + + def __init__(self, in_ch=3, mid_ch=12, out_ch=3): + super(RSU6, self).__init__() + + self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1) + + self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1) + self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1) + self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1) + self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1) + self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1) + + self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=2) + + self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1) + + def forward(self, x): + hx = x + + hxin = self.rebnconvin(hx) + + hx1 = self.rebnconv1(hxin) + hx = self.pool1(hx1) + + hx2 = self.rebnconv2(hx) + hx = self.pool2(hx2) + + hx3 = self.rebnconv3(hx) + hx = self.pool3(hx3) + + hx4 = self.rebnconv4(hx) + hx = self.pool4(hx4) + + hx5 = self.rebnconv5(hx) + + hx6 = self.rebnconv6(hx5) + + hx5d = self.rebnconv5d(torch.cat((hx6, hx5), 1)) + hx5dup = _upsample_like(hx5d, hx4) + + hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1)) + hx4dup = _upsample_like(hx4d, hx3) + + hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1)) + hx3dup = _upsample_like(hx3d, hx2) + + hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1)) + hx2dup = _upsample_like(hx2d, hx1) + + hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1)) + + return hx1d + hxin + + +### RSU-5 ### +class RSU5(nn.Module): # UNet05DRES(nn.Module): + + def __init__(self, in_ch=3, mid_ch=12, out_ch=3): + super(RSU5, self).__init__() + + self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1) + + self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1) + self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1) + self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1) + self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1) + + self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=2) + + self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1) + + def forward(self, x): + hx = x + + hxin = self.rebnconvin(hx) + + hx1 = self.rebnconv1(hxin) + hx = self.pool1(hx1) + + hx2 = self.rebnconv2(hx) + hx = self.pool2(hx2) + + hx3 = self.rebnconv3(hx) + hx = self.pool3(hx3) + + hx4 = self.rebnconv4(hx) + + hx5 = self.rebnconv5(hx4) + + hx4d = self.rebnconv4d(torch.cat((hx5, hx4), 1)) + hx4dup = _upsample_like(hx4d, hx3) + + hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1)) + hx3dup = _upsample_like(hx3d, hx2) + + hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1)) + hx2dup = _upsample_like(hx2d, hx1) + + hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1)) + + return hx1d + hxin + + +### RSU-4 ### +class RSU4(nn.Module): # UNet04DRES(nn.Module): + + def __init__(self, in_ch=3, mid_ch=12, out_ch=3): + super(RSU4, self).__init__() + + self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1) + + self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1) + self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1) + self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1) + + self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=2) + + self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1) + self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1) + + def forward(self, x): + hx = x + + hxin = self.rebnconvin(hx) + + hx1 = self.rebnconv1(hxin) + hx = self.pool1(hx1) + + hx2 = self.rebnconv2(hx) + hx = self.pool2(hx2) + + hx3 = self.rebnconv3(hx) + + hx4 = self.rebnconv4(hx3) + + hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1)) + hx3dup = _upsample_like(hx3d, hx2) + + hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1)) + hx2dup = _upsample_like(hx2d, hx1) + + hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1)) + + return hx1d + hxin + + +### RSU-4F ### +class RSU4F(nn.Module): # UNet04FRES(nn.Module): + + def __init__(self, in_ch=3, mid_ch=12, out_ch=3): + super(RSU4F, self).__init__() + + self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1) + + self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1) + self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=2) + self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=4) + + self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=8) + + self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=4) + self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=2) + self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1) + + def forward(self, x): + hx = x + + hxin = self.rebnconvin(hx) + + hx1 = self.rebnconv1(hxin) + hx2 = self.rebnconv2(hx1) + hx3 = self.rebnconv3(hx2) + + hx4 = self.rebnconv4(hx3) + + hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1)) + hx2d = self.rebnconv2d(torch.cat((hx3d, hx2), 1)) + hx1d = self.rebnconv1d(torch.cat((hx2d, hx1), 1)) + + return hx1d + hxin + + +##### U^2-Net #### +class U2NET(nn.Module): + + def __init__(self, in_ch=3, out_ch=1): + super(U2NET, self).__init__() + self.edge = sobel_net() + + self.stage1 = RSU7(in_ch, 32, 64) + self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage2 = RSU6(64, 32, 128) + self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage3 = RSU5(128, 64, 256) + self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage4 = RSU4(256, 128, 512) + self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage5 = RSU4F(512, 256, 512) + self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage6 = RSU4F(512, 256, 512) + + # decoder + self.stage5d = RSU4F(1024, 256, 512) + self.stage4d = RSU4(1024, 128, 256) + self.stage3d = RSU5(512, 64, 128) + self.stage2d = RSU6(256, 32, 64) + self.stage1d = RSU7(128, 16, 64) + + self.side1 = nn.Conv2d(64, out_ch, 3, padding=1) + self.side2 = nn.Conv2d(64, out_ch, 3, padding=1) + self.side3 = nn.Conv2d(128, out_ch, 3, padding=1) + self.side4 = nn.Conv2d(256, out_ch, 3, padding=1) + self.side5 = nn.Conv2d(512, out_ch, 3, padding=1) + self.side6 = nn.Conv2d(512, out_ch, 3, padding=1) + + self.outconv = nn.Conv2d(6, out_ch, 1) + + def forward(self, x): + x = self.edge(x) + hx = x + + # stage 1 + hx1 = self.stage1(hx) + hx = self.pool12(hx1) + + # stage 2 + hx2 = self.stage2(hx) + hx = self.pool23(hx2) + + # stage 3 + hx3 = self.stage3(hx) + hx = self.pool34(hx3) + + # stage 4 + hx4 = self.stage4(hx) + hx = self.pool45(hx4) + + # stage 5 + hx5 = self.stage5(hx) + hx = self.pool56(hx5) + + # stage 6 + hx6 = self.stage6(hx) + hx6up = _upsample_like(hx6, hx5) + + # -------------------- decoder -------------------- + hx5d = self.stage5d(torch.cat((hx6up, hx5), 1)) + hx5dup = _upsample_like(hx5d, hx4) + + hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1)) + hx4dup = _upsample_like(hx4d, hx3) + + hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1)) + hx3dup = _upsample_like(hx3d, hx2) + + hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1)) + hx2dup = _upsample_like(hx2d, hx1) + + hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1)) + + # side output + d1 = self.side1(hx1d) + + d2 = self.side2(hx2d) + d2 = _upsample_like(d2, d1) + + d3 = self.side3(hx3d) + d3 = _upsample_like(d3, d1) + + d4 = self.side4(hx4d) + d4 = _upsample_like(d4, d1) + + d5 = self.side5(hx5d) + d5 = _upsample_like(d5, d1) + + d6 = self.side6(hx6) + d6 = _upsample_like(d6, d1) + + d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1)) + + return torch.sigmoid(d0), torch.sigmoid(d1), torch.sigmoid(d2), torch.sigmoid(d3), torch.sigmoid( + d4), torch.sigmoid(d5), torch.sigmoid(d6) + + +### U^2-Net small ### +class U2NETP(nn.Module): + + def __init__(self, in_ch=3, out_ch=1): + super(U2NETP, self).__init__() + + self.stage1 = RSU7(in_ch, 16, 64) + self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage2 = RSU6(64, 16, 64) + self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage3 = RSU5(64, 16, 64) + self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage4 = RSU4(64, 16, 64) + self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage5 = RSU4F(64, 16, 64) + self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True) + + self.stage6 = RSU4F(64, 16, 64) + + # decoder + self.stage5d = RSU4F(128, 16, 64) + self.stage4d = RSU4(128, 16, 64) + self.stage3d = RSU5(128, 16, 64) + self.stage2d = RSU6(128, 16, 64) + self.stage1d = RSU7(128, 16, 64) + + self.side1 = nn.Conv2d(64, out_ch, 3, padding=1) + self.side2 = nn.Conv2d(64, out_ch, 3, padding=1) + self.side3 = nn.Conv2d(64, out_ch, 3, padding=1) + self.side4 = nn.Conv2d(64, out_ch, 3, padding=1) + self.side5 = nn.Conv2d(64, out_ch, 3, padding=1) + self.side6 = nn.Conv2d(64, out_ch, 3, padding=1) + + self.outconv = nn.Conv2d(6, out_ch, 1) + + def forward(self, x): + hx = x + + # stage 1 + hx1 = self.stage1(hx) + hx = self.pool12(hx1) + + # stage 2 + hx2 = self.stage2(hx) + hx = self.pool23(hx2) + + # stage 3 + hx3 = self.stage3(hx) + hx = self.pool34(hx3) + + # stage 4 + hx4 = self.stage4(hx) + hx = self.pool45(hx4) + + # stage 5 + hx5 = self.stage5(hx) + hx = self.pool56(hx5) + + # stage 6 + hx6 = self.stage6(hx) + hx6up = _upsample_like(hx6, hx5) + + # decoder + hx5d = self.stage5d(torch.cat((hx6up, hx5), 1)) + hx5dup = _upsample_like(hx5d, hx4) + + hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1)) + hx4dup = _upsample_like(hx4d, hx3) + + hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1)) + hx3dup = _upsample_like(hx3d, hx2) + + hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1)) + hx2dup = _upsample_like(hx2d, hx1) + + hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1)) + + # side output + d1 = self.side1(hx1d) + + d2 = self.side2(hx2d) + d2 = _upsample_like(d2, d1) + + d3 = self.side3(hx3d) + d3 = _upsample_like(d3, d1) + + d4 = self.side4(hx4d) + d4 = _upsample_like(d4, d1) + + d5 = self.side5(hx5d) + d5 = _upsample_like(d5, d1) + + d6 = self.side6(hx6) + d6 = _upsample_like(d6, d1) + + d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1)) + + return torch.sigmoid(d0), torch.sigmoid(d1), torch.sigmoid(d2), torch.sigmoid(d3), torch.sigmoid( + d4), torch.sigmoid(d5), torch.sigmoid(d6) diff --git a/update.py b/update.py new file mode 100644 index 0000000..51b7e36 --- /dev/null +++ b/update.py @@ -0,0 +1,106 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class FlowHead(nn.Module): + def __init__(self, input_dim=128, hidden_dim=256): + super(FlowHead, self).__init__() + self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1) + self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + return self.conv2(self.relu(self.conv1(x))) + + +class ConvGRU(nn.Module): + def __init__(self, hidden_dim=128, input_dim=192+128): + super(ConvGRU, self).__init__() + self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1) + self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1) + self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1) + + def forward(self, h, x): + hx = torch.cat([h, x], dim=1) + + z = torch.sigmoid(self.convz(hx)) + r = torch.sigmoid(self.convr(hx)) + q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1))) + + h = (1-z) * h + z * q + return h + + +class SepConvGRU(nn.Module): + def __init__(self, hidden_dim=128, input_dim=192+128): + super(SepConvGRU, self).__init__() + self.convz1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2)) + self.convr1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2)) + self.convq1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2)) + + self.convz2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0)) + self.convr2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0)) + self.convq2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0)) + + def forward(self, h, x): + # horizontal + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz1(hx)) + r = torch.sigmoid(self.convr1(hx)) + q = torch.tanh(self.convq1(torch.cat([r*h, x], dim=1))) + h = (1-z) * h + z * q + + # vertical + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz2(hx)) + r = torch.sigmoid(self.convr2(hx)) + q = torch.tanh(self.convq2(torch.cat([r*h, x], dim=1))) + h = (1-z) * h + z * q + + return h + + +class BasicMotionEncoder(nn.Module): + def __init__(self): + super(BasicMotionEncoder, self).__init__() + self.convc1 = nn.Conv2d(320, 240, 1, padding=0) + self.convc2 = nn.Conv2d(240, 160, 3, padding=1) + self.convf1 = nn.Conv2d(2, 160, 7, padding=3) + self.convf2 = nn.Conv2d(160, 80, 3, padding=1) + self.conv = nn.Conv2d(160+80, 160-2, 3, padding=1) + + def forward(self, flow, corr): + cor = F.relu(self.convc1(corr)) + cor = F.relu(self.convc2(cor)) + flo = F.relu(self.convf1(flow)) + flo = F.relu(self.convf2(flo)) + + cor_flo = torch.cat([cor, flo], dim=1) + out = F.relu(self.conv(cor_flo)) + return torch.cat([out, flow], dim=1) + + +class BasicUpdateBlock(nn.Module): + def __init__(self, hidden_dim=128): + super(BasicUpdateBlock, self).__init__() + self.encoder = BasicMotionEncoder() + self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=160+160) + self.flow_head = FlowHead(hidden_dim, hidden_dim=320) + + self.mask = nn.Sequential( + nn.Conv2d(hidden_dim, 288, 3, padding=1), + nn.ReLU(inplace=True), + nn.Conv2d(288, 64*9, 1, padding=0)) + + def forward(self, net, inp, corr, flow): + motion_features = self.encoder(flow, corr) + inp = torch.cat([inp, motion_features], dim=1) + + net = self.gru(net, inp) + + delta_flow = self.flow_head(net) + + mask = .25 * self.mask(net) + + return net, mask, delta_flow