init
This commit is contained in:
54
LICENSE.md
Normal file
54
LICENSE.md
Normal file
@ -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.
|
||||
43
OCR_eval.py
Normal file
43
OCR_eval.py
Normal file
@ -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)
|
||||
99
README.md
Normal file
99
README.md
Normal file
@ -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
|
||||
|
||||
<p>
|
||||
<a href='https://drive.google.com/file/d/1mmCUj90rHyuO1SmpLt361youh-07Y0sD/view?usp=share_link' target="_blank"><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a>
|
||||
<a href='https://docai.doctrp.top:20443/' target="_blank"><img src='https://img.shields.io/badge/Online-Demo-green'></a>
|
||||
</p>
|
||||
|
||||
|
||||
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).
|
||||
|
||||
|
||||
|
||||
|
||||
## 🚀 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.
|
||||
|
||||
<img width="1534" alt="image" src="https://github.com/fh2019ustc/DocScanner/assets/50725551/9eca3f7d-1570-4246-a3db-0a1cf1eece2d">
|
||||
|
||||
### Examples
|
||||
|
||||
|
||||
|
||||
|
||||
## 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)).
|
||||
|
||||
64
eval.m
Normal file
64
eval.m
Normal file
@ -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
|
||||
102
evalUnwarp.m
Normal file
102
evalUnwarp.m
Normal file
@ -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
|
||||
134
extractor.py
Normal file
134
extractor.py
Normal file
@ -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
|
||||
116
inference.py
Normal file
116
inference.py
Normal file
@ -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()
|
||||
100
model.py
Normal file
100
model.py
Normal file
@ -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
|
||||
62
ocr_img.txt
Normal file
62
ocr_img.txt
Normal file
@ -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
|
||||
6
requirements.txt
Normal file
6
requirements.txt
Normal file
@ -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
|
||||
552
seg.py
Normal file
552
seg.py
Normal file
@ -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)
|
||||
106
update.py
Normal file
106
update.py
Normal file
@ -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
|
||||
Reference in New Issue
Block a user