diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bee8a64 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__ diff --git a/OCR_eval.py b/OCR_eval.py deleted file mode 100644 index e77efc7..0000000 --- a/OCR_eval.py +++ /dev/null @@ -1,43 +0,0 @@ -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/eval.m b/eval.m deleted file mode 100644 index 0189322..0000000 --- a/eval.m +++ /dev/null @@ -1,64 +0,0 @@ -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 deleted file mode 100644 index 6f84634..0000000 --- a/evalUnwarp.m +++ /dev/null @@ -1,102 +0,0 @@ -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/inference.py b/inference.py index c04a144..e698164 100644 --- a/inference.py +++ b/inference.py @@ -1,115 +1,85 @@ +import os, sys, argparse, warnings +import torch, torch.nn as nn, torch.nn.functional as F +import numpy as np, cv2 +from PIL import Image 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') +# Запоминаем CWD пользователя ДО смены директории +USER_CWD = os.getcwd() +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +os.chdir(SCRIPT_DIR) + class Net(nn.Module): def __init__(self): - super(Net, self).__init__() + super().__init__() self.msk = U2NETP(3, 1) - self.bm = DocScanner() # 矫正 + 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 + msk, *_ = self.msk(x) + bm = self.bm((msk > 0.5).float() * x, iters=12, test_mode=True) + return (2 * (bm / 286.8) - 1) * 0.99 -def reload_seg_model(model, path=""): - if not bool(path): - return model +def load_model(model, path, strip_prefix=False): + state_dict = model.state_dict() + pretrained = torch.load(path, map_location='cuda:0') + if strip_prefix: + pretrained = {k[6:]: v for k, v in pretrained.items() if k[6:] in state_dict} 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 + pretrained = {k: v for k, v in pretrained.items() if k in state_dict} + state_dict.update(pretrained) + model.load_state_dict(state_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 resolve_user_path(path): + """Если путь относительный — считаем его относительно CWD пользователя, не скрипта""" + if os.path.isabs(path): + return path + return os.path.join(USER_CWD, path) 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/') + parser.add_argument('-i', '--input', required=True) + parser.add_argument('-o', '--output', required=True) 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) + input_path = resolve_user_path(opt.input) + output_path = resolve_user_path(opt.output) + + net = Net().cuda().eval() + load_model(net.msk, f'{SCRIPT_DIR}/model_pretrained/seg.pth', strip_prefix=True) + load_model(net.bm, f'{SCRIPT_DIR}/model_pretrained/DocScanner-L.pth') + + image = np.array(Image.open(input_path))[:, :, :3] / 255.0 + height, width = image.shape[:2] + tensor = torch.from_numpy(cv2.resize(image, (288, 288)).transpose(2, 0, 1)).float().unsqueeze(0) + + with torch.no_grad(): + bm = net(tensor.cuda()).cpu() + + flow = torch.from_numpy(np.stack([ + cv2.blur(cv2.resize(bm[0, 0].numpy(), (width, height)), (3, 3)), + cv2.blur(cv2.resize(bm[0, 1].numpy(), (width, height)), (3, 3)) + ], axis=2)).unsqueeze(0) + + out = F.grid_sample( + torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float(), + flow, align_corners=True + ) + result = ((out[0] * 255).permute(1, 2, 0).numpy()[:, :, ::-1]).astype(np.uint8) + + out_dir = os.path.dirname(output_path) + if out_dir: + os.makedirs(out_dir, exist_ok=True) + cv2.imwrite(output_path, result) + print(f"[OK] {output_path}") if __name__ == "__main__": diff --git a/model_pretrained/DocScanner-L.pth b/model_pretrained/DocScanner-L.pth new file mode 100644 index 0000000..901891d Binary files /dev/null and b/model_pretrained/DocScanner-L.pth differ diff --git a/model_pretrained/seg.pth b/model_pretrained/seg.pth new file mode 100644 index 0000000..8b9c6d7 Binary files /dev/null and b/model_pretrained/seg.pth differ diff --git a/ocr_img.txt b/ocr_img.txt deleted file mode 100644 index fe2211d..0000000 --- a/ocr_img.txt +++ /dev/null @@ -1,62 +0,0 @@ -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