├── public ├── assets │ ├── comparison.jpg │ └── methodology.png ├── visual comparison.zip └── paper_segmentation_dataset.zip ├── src ├── main │ ├── yolo_segmentation.py │ ├── geom_metrics.py │ ├── run.py │ └── document_recovery.py └── 1.txt └── README.md /public/assets/comparison.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HorizonParadox/DRCCBI/HEAD/public/assets/comparison.jpg -------------------------------------------------------------------------------- /public/assets/methodology.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/HorizonParadox/DRCCBI/HEAD/public/assets/methodology.png -------------------------------------------------------------------------------- /public/visual comparison.zip: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:43c3e1ed979a7b42bed57a02654765c3e4cd456e0532f92b63cf08dbee166725 3 | size 38531668 4 | -------------------------------------------------------------------------------- /public/paper_segmentation_dataset.zip: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:c42e8a8b16c3bb8070986b899203900ba3530c4158fd170930746329abcdcb79 3 | size 821889662 4 | -------------------------------------------------------------------------------- /src/main/yolo_segmentation.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from ultralytics import YOLO 3 | 4 | 5 | class YOLOSegmentation: 6 | def __init__(self, model_path): 7 | self.model = YOLO(model_path) 8 | 9 | def detect(self, img): 10 | height, width, channels = img.shape 11 | 12 | results = self.model.predict(source=img.copy(), conf=0.6, save=False, save_txt=False) 13 | result = results[0] 14 | segmentation_contours_idx = [] 15 | 16 | for seg in result.masks.xyn: 17 | seg[:, 0] *= width 18 | seg[:, 1] *= height 19 | segment = np.array(seg, dtype=np.int32) 20 | segmentation_contours_idx.append(segment) 21 | 22 | bboxes = np.array(result.boxes.xyxy.cpu(), dtype="int") 23 | scores = np.array(result.boxes.conf.cpu(), dtype="float").round(3) 24 | 25 | return bboxes, segmentation_contours_idx, scores 26 | -------------------------------------------------------------------------------- /src/main/geom_metrics.py: -------------------------------------------------------------------------------- 1 | import cv2 as cv 2 | import numpy as np 3 | from skimage.metrics import structural_similarity as ssim 4 | from skimage.metrics import normalized_root_mse as nrmse 5 | 6 | 7 | def mse(image1, image2): 8 | err = np.sum((image1.astype("float") - image2.astype("float")) ** 2) 9 | err /= float(image1.shape[0] * image2.shape[1]) 10 | return err 11 | 12 | 13 | directory_1_1 = '../../images/OCR/test_images/1_1/' 14 | directory_1_2 = '../../images/OCR/test_images/1_2/' 15 | 16 | directory_3_1 = '../../images/OCR/test_images/3_1/not_scan/' 17 | directory_3_2 = '../../images/OCR/test_images/3_2/' 18 | 19 | directory_5_1 = '../../images/OCR/test_images/5_1/' 20 | directory_5_2 = '../../images/OCR/test_images/5_2/' 21 | 22 | directory_7_1 = '../../images/OCR/test_images/7_1/' 23 | directory_7_2 = '../../images/OCR/test_images/7_2/' 24 | 25 | directory_10_1 = '../../images/OCR/test_images/10_1/' 26 | directory_10_2 = '../../images/OCR/test_images/10_2/' 27 | 28 | directory_13_1 = '../../images/OCR/test_images/13_1/' 29 | 30 | directory_28_1 = '../../images/OCR/test_images/28_1/' 31 | directory_28_2 = '../../images/OCR/test_images/28_2/' 32 | 33 | directory_46_1 = '../../images/OCR/test_images/46_1/' 34 | directory_46_2 = '../../images/OCR/test_images/46_2/' 35 | 36 | my_image = cv.imread(directory_46_2 + '46_2.jpg') 37 | scan_image = cv.imread(directory_46_2 + '46_2_scan.png') 38 | orig_image = cv.imread(directory_46_2 + '46_2_orig.jpg') 39 | any_scan_image = cv.imread(directory_46_2 + '46_2_any_scanner.jpg') 40 | cam_scanner_image = cv.imread(directory_46_2 + '46_2_cam_scanner.jpg') 41 | tap_scanner_image = cv.imread(directory_46_2 + '46_2_tap_scanner.jpg') 42 | 43 | height_image, width_image, _ = scan_image.shape 44 | scan_gray = cv.cvtColor(scan_image, cv.COLOR_BGR2GRAY) 45 | 46 | my_image = cv.resize(my_image, (width_image, height_image)) 47 | orig_image = cv.resize(orig_image, (width_image, height_image)) 48 | any_scan_image = cv.resize(any_scan_image, (width_image, height_image)) 49 | cam_scanner_image = cv.resize(cam_scanner_image, (width_image, height_image)) 50 | tap_scanner_image = cv.resize(tap_scanner_image, (width_image, height_image)) 51 | 52 | print(f'SSIM for my image: {ssim(scan_gray, cv.cvtColor(my_image, cv.COLOR_BGR2GRAY))}') 53 | print(f'SSIM for orig image: {ssim(scan_gray, cv.cvtColor(orig_image, cv.COLOR_BGR2GRAY))}') 54 | print(f'SSIM for Any scan: {ssim(scan_gray, cv.cvtColor(any_scan_image, cv.COLOR_BGR2GRAY))}') 55 | print(f'SSIM for Cam scanner: {ssim(scan_gray, cv.cvtColor(cam_scanner_image, cv.COLOR_BGR2GRAY))}') 56 | print(f'SSIM for Tap scanner: {ssim(scan_gray, cv.cvtColor(tap_scanner_image, cv.COLOR_BGR2GRAY))}') 57 | 58 | print('\n') 59 | 60 | print(f'MSE for my image: {mse(scan_image, my_image)}') 61 | print(f'MSE for orig image: {mse(scan_image, orig_image)}') 62 | print(f'MSE for Any scan: {mse(scan_image, any_scan_image)}') 63 | print(f'MSE for Cam scanner: {mse(scan_image, cam_scanner_image)}') 64 | print(f'MSE for Tap scanner: {mse(scan_image, tap_scanner_image)}') 65 | 66 | print('\n') 67 | 68 | print(f'NRMSE for my image: {nrmse(scan_image, my_image)}') 69 | print(f'NRMSE for orig image: {nrmse(scan_image, orig_image)}') 70 | print(f'NRMSE for Any scan: {nrmse(scan_image, any_scan_image)}') 71 | print(f'NRMSE for Cam scanner: {nrmse(scan_image, cam_scanner_image)}') 72 | print(f'NRMSE for Tap scanner: {nrmse(scan_image, tap_scanner_image)}') 73 | 74 | -------------------------------------------------------------------------------- /src/main/run.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | from matplotlib import pyplot as plt 4 | from src.main.document_recovery import DocumentRecovery 5 | 6 | NUM_OF_GRID_LINE = 9 7 | POLYNOMIAL_DEGREE = 3 8 | NUM_POINT_PERCENT = 0.1 9 | THRESHOLD_COEFFICIENT = 0.05 10 | INPUT_DIRECTORY = "../../images/new_med_docs" 11 | INPUT_DIRECTORY_DEPTH = "../../document_depth_maps" 12 | OUTPUT_DIRECTORY = "../../images/output" 13 | WEIGHT_DIRECTORY = '../../weights/best.pt' 14 | 15 | for image_filename in os.listdir(INPUT_DIRECTORY): 16 | image_file = os.path.join(INPUT_DIRECTORY, image_filename) 17 | if os.path.isfile(image_file): 18 | print(image_filename) 19 | 20 | model = DocumentRecovery(image_file, WEIGHT_DIRECTORY) 21 | height_image, width_image, _ = model.initial_image.shape 22 | num_points = int(width_image * NUM_POINT_PERCENT) 23 | print(f'num_points: {num_points}') 24 | 25 | num_of_mask = 0 26 | # Функция Get mask и поиск контура 27 | polynomial = model.find_polynomial(thres_coef=THRESHOLD_COEFFICIENT) 28 | for concave, conv, cont in polynomial: 29 | num_of_mask += 1 30 | 31 | # Построение контура 32 | model.plot_polynomial(concave, conv, cont) 33 | 34 | # Поиск сторон в контуре 35 | left, right, top, bottom = model.find_edges(cont, conv) 36 | 37 | # Получение приблизительной длины и ширины контура 38 | outline_length, outline_width = model.get_dimensions(left, right, top, bottom) 39 | 40 | # Отображение полученных сторон разным цветом 41 | model.plot_colorful_edges(left, right, top, bottom) 42 | 43 | # Интерполяция двух кривых 44 | new_left_x, new_left_y, new_right_x, new_right_y = model.calculate_opposite_interpolate(left, right, True) 45 | new_bottom_x, new_bottom_y, new_top_x, new_top_y = model.calculate_opposite_interpolate(bottom, top, True) 46 | 47 | # Построение линий аппроксимации (сетка) 48 | linfit_x_bt, linfit_y_bt, linfit_x_lr, linfit_y_lr = model.get_lines_of_approximation( 49 | NUM_OF_GRID_LINE, new_left_x, new_left_y, new_right_x, new_right_y, new_bottom_x, new_bottom_y, 50 | new_top_x, new_top_y) 51 | 52 | plt.imshow(model.initial_image) 53 | # Построение интерполяции аппроксимации (горизонт, вертикал) 54 | polys_bt = model.get_interpolation_approximation( 55 | width_image, NUM_OF_GRID_LINE, POLYNOMIAL_DEGREE, num_points, linfit_x_bt, linfit_y_bt) 56 | polys_lr = model.get_interpolation_approximation( 57 | height_image, NUM_OF_GRID_LINE, POLYNOMIAL_DEGREE, num_points, linfit_x_lr, linfit_y_lr) 58 | plt.show() 59 | 60 | # Поиск пересечений между линиями, заданными интерполяционными аппроксимациями 61 | points_set = model.get_intersection_interpolation_approximation( 62 | outline_length, outline_width, polys_bt, polys_lr) 63 | 64 | # Получение координат исправленной сетки 65 | XYpairs = model.get_fix_grid_coordinates(outline_length, outline_width, NUM_OF_GRID_LINE) 66 | 67 | # Восстановление искажения перспективы изображения 68 | crop_warped_image = model.remap_image(height_image, width_image, points_set, XYpairs) 69 | 70 | # Превращает документ в чёрно-белый скан 71 | # scan = DC.get_black_white_scan(crop_warped_image) 72 | 73 | image_name, image_extension = image_filename.split('.') 74 | cv2.imwrite(OUTPUT_DIRECTORY + f'/{image_name}({num_of_mask}).{image_extension}', crop_warped_image) 75 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |
2 |

Geometry Restoration and Dewarping of Camera-Captured Document Images

3 | 4 | [**Valery Istomin**](https://www.linkedin.com/in/valery-istomin-90a473247/)1, [**Oleg Perezyabov**](https://www.linkedin.com/in/oleg-pereziabov-a6b287254/)2 and [**Ilya Afanasyev**](https://www.linkedin.com/in/ilya-afanasyev-8783291a/)3,4 5 | 6 | 1Shuya Branch of the Ivanovo State University, Shuya, Russia
7 | 2BIA Technologies, St. Petersburg, Russia
8 | 3Saint Petersburg Electrotechnical University "LETI", St. Petersburg, Russia
9 | 4Innopolis University, Innopolis, Russia 10 | 11 | Paper PDF 12 | Benchmark 13 | Dataset 14 | 15 |
16 | 17 | This research focuses on developing a method for restoring the topology of digital images of paper documents captured by a camera, using algorithms for detection, segmentation, geometry restoration, and dewarping. Our methodology employs deep learning (DL) for document outline detection, followed by computer vision (CV) to create a topological 2D grid using cubic polynomial interpolation and correct nonlinear distortions by remapping the image. Using classical CV methods makes the document topology restoration process more efficient and faster, as it requires significantly fewer computational resources and memory. We developed a new pipeline for automatic document dewarping and reconstruction, along with a framework and annotated dataset to demonstrate its efficiency. Our experiments confirm the promise of our methodology and its superiority over existing benchmarks (including mobile apps and popular DL solutions, such as RectiNet, DocGeoNet, and DocTr++) both visually and in terms of document readability via Optical Character Recognition (OCR) and geometry restoration metrics. This paves the way for creating high-quality digital copies of paper documents and enhancing the efficiency of OCR systems. 18 | 19 | Keywords: Document Image Dewarping, Image Distortions, Geometry Restoration 20 | 21 | ![teaser](public/assets/methodology.png) 22 | 23 | The flowchart of the document geometry restoration and dewarping algorithm: 24 | 1) Identifying the document mask using the YOLOv8 model; 25 | 2) Detecting the contour edges of the document, approximating the corners, and segmenting the contour into fragments corresponding to each side of the document; 26 | 3) Creating a 2D grid of the document by interpolating its opposite sides with evenly spaced curved lines, approximating each line with a cubic polynomial; 27 | 4) Detecting the intersection points of the curved lines, constructing the resulting grid for image transformation, and creating a transformation map based on the 2D points, followed by remapping the original image. 28 | 29 | ![teaser](public/assets/comparison.jpg) 30 | 31 | The comparison of documents reconstructed by popular desktop DL models - DocTr++, DocGeoNet, RectiNet and our algorithm. 32 | 33 | ## Citation 34 | 35 | If you find this project useful, please consider citing: 36 | 37 | ```bibtex 38 | @article{istomin2025geometry, 39 | title={Geometry Restoration and Dewarping of Camera-Captured Document Images}, 40 | author={Istomin, Valery and Pereziabov, Oleg and Afanasyev, Ilya}, 41 | journal={arXiv preprint arXiv:2501.03145}, 42 | year={2025} 43 | } 44 | -------------------------------------------------------------------------------- /src/1.txt: -------------------------------------------------------------------------------- 1 | 10ENGLISHPersonalizingyourbrushingexperiencePhilipsSonicareautomaticallystartsinthedefaultCleanmode.Topersonalizeyourbrushing:1PriortoturningonthePhilipsSonicare,pressthepersonalizedbrushingbuttontocyclethroughthemodesandroutines.ThegreenLEDindicatestheselectedmodeorroutine.BrushingmodesCleanmodeStandardmodeforsuperiorteethcleaning.Whitemode2minutesofCleanmode,withanadditional30secondsofWhitemodetofocusonyourvisiblefrontteeth.Whitemodebrushinginstructions1Brushthefirst2minutesasexplainedinsection'Brushinginstructions'.2Afterthe2minutesofCleanmode,theWhitemodestartswithachangeinbrushingsoundandmotion.Thisisyoursignaltostartbrushingtheupperfrontteethfor15seconds.3Atthenextbeepandpause,movetothebottomfrontteethforthefinal15secondsofbrushing.MassagemodeModeforgentlegumstimulation. 2 | | 10 ENGLISH i Personalizing your brushing experience Philips Sonicare automatically starts in the default Clean mode. To personalize your brushing: Prior to turning on the Philips Sonicare, press the personalized brushing button to cycle through the modes and routines. __-D The green LED indicates the selected mode or routine. | J Brushing modes | Clean mode Standard mode for superior teeth cleaning. White mode 2 minutes of Clean mode, with an additional : 30 seconds of White mode to focus on your visible front teeth. White mode brushing instructions Brush the first 2 minutes as explained in section ‘Brushing instructions’. After the 2 minutes of Clean mode, the White mode starts with a change in brushing sound | and motion. This is your signal to start brushing | the upper front teeth for 15 seconds. — | | At the next beep and pause, move to the bottom front teeth for the final 15 seconds of brushing. Massage mode Mode for gentle gum stimulation. | | 3 | 10 ENGLISH i i e es Personalizing your brushing experience” = Sen ; Philit icare automatically Pans lin Wire ault Clean mode.T> personalize your br ushin Eq Prior to turning on the Philips Sonicare, ! Press the personalized brushing button to cycle through the modes and routines. > The green LED indicates the selected mode | Or routine, ’ Brushing modes ee USng Modes Clean mode Standard mode for superior teeth cleaning. White mode 2 minutes of Clean mode, with an additional 30 seconds of White mode to focus on your visible front teeth White mode brushing instructions Brush the first 2 minutes as explained in section ‘Brushing instructions’. After the 2 minutes of Clean mode, the ae mode starts with a change in brushing soun | and motion. This is your signal to start brushing the upper front teeth for 15 seconds. } IEW At the next beep and pause, move to the bottom front teeth for the final 15 seconds | of brushing. : Massage mode Mode for gentle gum stimulation 4 | a oa \eae 4 i ~ gee | oe z \nsek* \e S| ) £ y ij ee ane i Yome : soit y/ i y Bs Ho . Hiss voce , | an EA2 ‘ ee Fe 54 coal + ‘ os a F - Od a — \ 4 bg Uf Y ) 2 7 Zt a aS ¥9 ) = \ ‘\ ' y Se 8 C 0 S ‘ - 4 el * i Ss st ’ 7q ‘ ie) t z\ ee . er = Ne o \ gt = \ ~ , : a Le) \ \ \ — Ne ” \ . \ Pelee 3 TF sm, \J Na as ie ips Ae ee 2 Seep mas \) SS Vue y \ s sas a ; y a 3 \ \ \ : , - _ ~~ © 9 \ > - rs Ne / sGa. i , ee 4 A. REN Vena t — 2 ; G . on\ J pave eee Z wm : é we \% eh . a Ps x\ A, i e ‘ 4 in \ Se e\¥ Oe * z= 3 \ i ™ \ P ka =a _ \ \ yom & , a poe aa ee | \ e ,\ & oe | S a cE \ ee sa ; Pee . , ‘ SS - a] yo i SS a < =< si, ess { the Philips ton tO sav? a ee a | TEM Prior to turning on d prushing isi i Fa te * ~ is a | press the personaliz© and routines: Ae | : _ a | cycle through the Vienne selected Mo —— es = | dD The green LED indicates \ : — or routine. { | Brushing modes Clean mode . leaning: | \ Standard mode for superior es \ White mode dditional Se u\\ | 2 minutes of Clean mode, win es on your | .. | . 30 seconds of White mode Sa \eNS | visible front teeth. tions | Ree: eee | ing instruct! | Bie | White mode brushing ins' Fedo | Battie \ cas eae = | i a Pee \ion eee ] Brush the first 2 minutes is © P Re Gi sue ing instructions - : | : Ko oe ee 3 ; section ‘Brushing instr’ sie the White ) . : - | After the 2 minutes of oe brushing sound ; mode starts with a change 1 rare brushing \ ; and motion. This is your Se on ( the upper front teeth for 1 ff z ve to At the next beep and Pe final 15 seconds Marte ; bottom front teeth for t ss 2 t of brushing. i) ais re Massage mode a | a Mode for gentle gum stimulation. " : : . X ¥ ‘: i = 2 fae y —— 4 seis e 4 , . VS : ; er i Baer GIR Ps Lae Persea BE are a AL Seg as 5 | | " cy. (; X Sawer ) aw | Pe | i> r \ Go. ox og Bn | oO , oa i % my id & ; a We Zz eas , ; Fi . \ , . : es Ss op. 3 $ 4 . " _ \\ a | | * A — fo ' | 10 Rhy a : Sy N ae Z i = - S ~ _ -_,* 4H " Le camel ~ eae | Per =e | 5 ~Cnalizing your brushing experience vee er Cie SOnicare automatically starts in the default b= 4 iy 'MOde To personalize your brushing: \ Prior : f . ' nel to turning on the Philips Sonicare, j Press the personalized brushing button to ] | ) Tee through the modes and routines. | € green LED indicates the selected mode = j Or routine, 4 ] Brushing modes : - 4 | Clean mode j Standard mode for superior teeth cleaning. 2 minutes of Clean mode, with an additional : : 30 seconds of White mode to focus on you" visible front teeth, | ) White mode brushing instructions | Brush the first 2 minutes as explained section ‘Brushing instructions . ite After the 2 minutes of Clean mode, the ae! | mode starts with a change in br ushing can and motion. This is your signal to start brushing — the upper front teeth for 15 seconds. if At the next beep and pause, move to the bottom front teeth for the final 15 seconds of brushing. : Massage Mode | Ode for gentle gum stimulation. | | | | = —<—— ae a rd - < , — - Ia 6 | = Ant) ) pees + a co \ \y * ( — 7 ~ 0 \ . 4 i fj y ) r i = H/ ha B \ Ij { \ \ \ i xX oe y) iy J e:, \ 2 : , Ss » \\ \\ =e | a : \ roa Wa os \Y Fr , yet lis. -~ ; rene MS ies a ; v4 \ > pee fst ‘ = onis\ A 4 \ _ n ” ; LY oy 5 . val ~) \ A ney wis Bi 3H Bs 7], > \ ¥ ee =. - 7 { eo / an ~ Pop \ + j F NA ; PP ee = \ N Hae owl i a uu i » eB \ a “ prise ’ “eA ee - \ \ a & « wane : \ 4 eel rr, SN — ' C { ee a % ENGLISH rae a pl a ee ae P ner shing experienc ro ee Personalizing your Brust TT erat ‘ = ep | Philips Sonicare automatically et. ishing: ‘ ae ae j Clean mode.To per sonalize yOu! brt g; | ee | she | Fi % A icares , “Se a | Prior’'to cucming On the Philips Sonicar" a = "a | : rning CO" d brushing button ; a Press the personalize - 7 - | BS des and routines. " -- | cycle through the mo 4 mode a | > The green LED indicates the selecte | f ica | Or routine. | | | Se Brushing modes. § "of Clean mode (ore . | Standard mode for superior teeth cleaning: White mode | 2 minutes of Clean mode, with an additional | . 30 seconds of White mode to focus on your . Visible front teeth. . White mode brushing instructions | ; Brush the first 2 minutes as explained in - section ‘Brushing instructions’. After the 2 minutes of Clean mode, the White \ mode starts with a change in brushing sound \ and motion. This is your signal to start brushing the upper front teeth for 15 seconds: ( BEB At the next beep and pause, move t© the bottom front teeth for the final 15 seconds of brushing. Massage mode | Mode for gentle gum stimulation. \ | \ | \ ; y - Pn, 7a rags : : A a ra — £ wecpasn ra —— cA 7 | 8 ALGER AE BA eee x bess + Horn End \ Insert \ ete eee = So 2 a et \ pees SS = oy \e ied 7aN m il 4 ~ % Sat 1. Q a Sy y i & \ ) ce | (aa ~ \, i Ks : f q 7 8 Q 0 \ a, I= { 1 7 5 a, : \ { i. ou ae ] \ 4 + i \ \ my Ni a \ e. " ? Mf \ \ ae “a y) a € " ph ; \er YF G re 3.1. % Van en .% ee Vibe Sh i ‘ 0 ba “i = 5 Z _ | — 4 ‘ ; SF ge aoe on \ ce \ < > |) ee =: \ OF a \ \ i a \ oa 4 \ ; ae .B \ . ” eo evi’ a A P \ \ 1s ams @k Pad" | , a ay PO ae . ; RG 3 AY W ty wi . ra j fr Th UN, PriSc Curl bes 4 4 soca ils ' ™ — aaa : cranes \ ’ a a - \ NS wes Sos | . ae 7 s ‘ _ > [ee Oe * = a ere — | en er eS ; y - SLisy pepe Le wt ee A. : g experione= ee — : mi * — | IE | Personatizin pr brush ee sefault — Sea ; iy ts in the . = igen ae ; ’ bi | i e ~ tically t40 ching: ; + i a F iilips Sonicare jutomatt our prushing } Ser eS ks - - a Clean mode.To persona conicares SSS ee ; = as; | Philips to j See si * " El Prior to turning On he rushing button a ee Se, eco press the personaliz® des and routines: | ae 2 or os \ hs cycle through the mo! ine selected mo | 7. Pe a oa | » The green LED indicates ’ - ; ea | Brushing modes Bak | uss es | Clean mode -geeth cleaning Rare arg Standard mode for superior FS” ior Re, Whit 4+ion | Se a = 2 sie orcad mode, with an addition! AAS Pues : es of Clean mo aig Son yO oN : AG ecemeecr vite mode Cue 1: 77 | # Еслиа на изображении несколько объектов... 78 | single_image_flag = False 79 | 80 | # Определяет два объекта с наибольшими оценками 81 | for bbox, seg, score in zip(bboxes, segments, scores): 82 | yolo_result.append([bbox, seg, score]) 83 | max1 = max(yolo_result, key=lambda x: x[2]) 84 | yolo_result.remove(max1) 85 | max2 = max(yolo_result, key=lambda x: x[2]) 86 | 87 | # Проверяет пересечение между сегментами двух объектов 88 | segments_intersection_flag = self.find_mask_intersection(max1[1], max2[1]) 89 | 90 | if segments_intersection_flag: 91 | # Если есть пересечение, добавляет два объекта в yolo_output 92 | yolo_output = [max1, max2] 93 | else: 94 | # Если пересечения нет, устанавливает флаг одного изображения и добавляет объект с наибольшей оценкой 95 | single_image_flag = True 96 | yolo_output = [max([max1, max2], key=lambda x: x[2])] 97 | 98 | print(f'SHAPE: {self.initial_image.shape}') 99 | 100 | # Обрабатывает каждый объект в yolo_output 101 | for bbox, seg, score in yolo_output: 102 | if single_image_flag or segments_intersection_flag: 103 | print(f'Score: {score}') 104 | 105 | # Создает маску объекта 106 | mask = np.zeros(self.initial_image.shape[:2], np.uint8) 107 | cv2.drawContours(mask, [seg], -1, (255, 255, 255), -1, cv2.LINE_AA) 108 | self.plot_image(mask) 109 | 110 | # Определяет параметры фильтрации маски 111 | radius = int(min(self.initial_image.shape[0], self.initial_image.shape[1]) * 0.01) 112 | if radius < 2: 113 | radius = 3 114 | epsilon = radius // 2 115 | print(f'Radius: {radius}') 116 | 117 | # Применяет фильтрацию маски 118 | fil_mask = cv2.ximgproc.guidedFilter(self.initial_image, mask, radius, epsilon) 119 | _, bw = cv2.threshold(fil_mask, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) 120 | masks.append(bw) 121 | self.plot_image(bw) 122 | return masks 123 | 124 | def find_polynomial(self, thres_coef): 125 | """ 126 | Находит полиномиальные кривые на масках документа. 127 | 128 | Аргументы: 129 | - thres_coef: Коэффициент порога для аппроксимации полиномиальной кривой. 130 | 131 | Возвращает: 132 | - results: Список с результатами, содержащими аппроксимированные полиномиальные кривые, выпуклые оболочки 133 | и контуры для каждой маски. 134 | """ 135 | 136 | masks = self.find_document_masks_yolo() 137 | results = [] 138 | for mask in masks: 139 | # Находит контуры в маске 140 | contours, hierarchy = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) 141 | # Находит наибольший контур по площади 142 | big_contour = max(contours, key=cv2.contourArea) 143 | result = np.zeros_like(self.initial_image) 144 | cv2.drawContours(result, [big_contour], -1, (0, 255, 0), cv2.FILLED) 145 | epsilon = thres_coef * cv2.arcLength(big_contour, True) 146 | # Аппроксимирует контур полиномиальной кривой 147 | approx = cv2.approxPolyDP(big_contour, epsilon, True) 148 | # Вычисляет выпуклую оболочку контура 149 | hull = cv2.convexHull(approx) 150 | results.append((approx, hull, big_contour)) 151 | return results 152 | 153 | def plot_polynomial(self, concave, conv, cont): 154 | """ 155 | Визуализирует полиномиальные кривые на изображении документа. 156 | 157 | Аргументы: 158 | - concave: Список контуров аппроксимированных полиномиальных кривых. 159 | - conv: Список контуров выпуклых оболочек полиномиальных кривых. 160 | - cont: Список контуров исходных контуров. 161 | 162 | Вывод: 163 | - Отображает изображение документа с нарисованными полиномиальными кривыми. 164 | """ 165 | 166 | img = self.initial_image.copy() 167 | cv2.drawContours(img, conv, -1, (255, 0, 0), 8) # Выпуклые оболочки - синий цвет 168 | cv2.drawContours(img, cont, -1, (0, 255, 0), 8) # Исходные контуры - зеленый цвет 169 | cv2.drawContours(img, concave, -1, (0, 0, 255), 8) # Аппроксимированные полиномиальные кривые - красный цвет 170 | 171 | self.plot_image(img) 172 | 173 | @staticmethod 174 | def sort_points(points): 175 | """ 176 | Сортирует точки в порядке их угла относительно центра фигуры. 177 | 178 | Аргументы: 179 | - points: Массив точек для сортировки. 180 | 181 | Возвращает: 182 | - Отсортированный массив точек. 183 | """ 184 | 185 | figure_center = np.mean(points, axis=0) 186 | points_list = list(points) 187 | points_list.sort(key=lambda point: np.arctan2(point[0] - figure_center[0], point[1] - figure_center[1])) 188 | return np.array(points_list) 189 | 190 | @staticmethod 191 | def line_point_cross_product(line, point) -> np.ndarray: 192 | """ 193 | Вычисляет векторное произведение между линией и точкой. 194 | 195 | Аргументы: 196 | - line: Массив из двух точек, представляющих линию. 197 | - point: Точка. 198 | 199 | Возвращает: 200 | - Векторное произведение в виде массива NumPy. 201 | """ 202 | 203 | v1 = np.subtract(line[1], line[0]) 204 | v2 = np.subtract(line[1], point) 205 | xp = np.cross(v1, v2) 206 | return xp 207 | 208 | def find_edges(self, contour, convex): 209 | """ 210 | Определяет края многоугольника на изображении. 211 | 212 | Аргументы: 213 | - contour: Контур многоугольника. 214 | - convex: Выпуклая оболочка многоугольника. 215 | 216 | Возвращает: 217 | - left: Край, лежащий слева от диагональных линий. 218 | - right: Край, лежащий справа от диагональных линий. 219 | - top: Край, лежащий сверху от диагональных линий. 220 | - bottom: Край, лежащий снизу от диагональных линий. 221 | """ 222 | 223 | img_test = self.initial_image.copy() 224 | 225 | # Извлечение граней и углов из массивов 226 | edges = np.squeeze(np.array(contour), axis=1) 227 | corners = np.reshape(np.array(convex), (4, 2)) 228 | 229 | # Сортировка углов по порядку 230 | corners = self.sort_points(corners) 231 | 232 | # Формирование диагональных линий 233 | diag_lines = np.array([[corners[0], corners[2]], [corners[1], corners[3]]]) 234 | 235 | # Отрисовка диагональных линий на изображении 236 | for num, line in enumerate(diag_lines): 237 | cv2.line(img_test, line[0], line[1], (0, 255 - 100 * num, 0), 2) 238 | 239 | self.plot_image(img_test) 240 | 241 | # Выделение краёв, лежащих слева от диагональных линий 242 | left = [edge for edge in np.array(edges).tolist() if 243 | (self.line_point_cross_product(diag_lines[0], edge) < 0 < self.line_point_cross_product( 244 | diag_lines[1], edge))] 245 | 246 | # Выделение краёв, лежащих снизу от диагональных линий 247 | bottom = [edge for edge in np.array(edges).tolist() if 248 | (self.line_point_cross_product(diag_lines[0], edge) < 0 and self.line_point_cross_product( 249 | diag_lines[1], edge) < 0)] 250 | 251 | # Выделение краёв, лежащих справа от диагональных линий 252 | right = [edge for edge in np.array(edges).tolist() if 253 | (self.line_point_cross_product(diag_lines[0], edge) > 0 > self.line_point_cross_product( 254 | diag_lines[1], edge))] 255 | 256 | # Выделение краёв, лежащих сверху от диагональных линий 257 | top = [edge for edge in np.array(edges).tolist() if 258 | (self.line_point_cross_product(diag_lines[0], edge) > 0 and self.line_point_cross_product( 259 | diag_lines[1], edge) > 0)] 260 | 261 | top.sort() 262 | 263 | right = right[::-1] # Инвертирование порядка краёв справа 264 | 265 | return left, right, top, bottom 266 | 267 | @staticmethod 268 | def get_dimensions(left, right, top, bottom): 269 | """ 270 | Вычисляет размеры многоугольника на изображении. 271 | 272 | Аргументы: 273 | - left: Край, лежащий слева от диагональных линий. 274 | - right: Край, лежащий справа от диагональных линий. 275 | - top: Край, лежащий сверху от диагональных линий. 276 | - bottom: Край, лежащий снизу от диагональных линий. 277 | 278 | Возвращает: 279 | - length: Длина многоугольника. 280 | - width: Ширина многоугольника. 281 | """ 282 | 283 | # Вычисление длины краёв 284 | left_length = round(spatial.distance.euclidean(left[0], left[-1])) 285 | right_length = round(spatial.distance.euclidean(right[0], right[-1])) 286 | top_length = round(spatial.distance.euclidean(top[0], top[-1])) 287 | bottom_length = round(spatial.distance.euclidean(bottom[0], bottom[-1])) 288 | 289 | # Вычисление средних значений длины и ширины 290 | length = int(np.mean([left_length, right_length])) 291 | width = int(np.mean([top_length, bottom_length])) 292 | 293 | return length, width 294 | 295 | def plot_colorful_edges(self, left, right, top, bottom): 296 | """ 297 | Визуализирует грани многоугольника на изображении разными цветами. 298 | 299 | Аргументы: 300 | - left: Грани, лежащие слева от диагональных линий. 301 | - right: Грани, лежащие справа от диагональных линий. 302 | - top: Грани, расположенные сверху от диагональных линий. 303 | - bottom: Грани, расположенные снизу от диагональных линий. 304 | """ 305 | 306 | img = self.initial_image.copy() 307 | for point in left: 308 | cv2.circle(img, point, 2, [0, 0, 255], -1) 309 | for point in bottom: 310 | cv2.circle(img, point, 2, [0, 255, 0], -1) 311 | for point in right: 312 | cv2.circle(img, point, 2, [255, 0, 0], -1) 313 | for point in top: 314 | cv2.circle(img, point, 2, [0, 255, 255], -1) 315 | self.plot_image(img) 316 | 317 | @staticmethod 318 | def interpolate(x, y, n_points): 319 | """ 320 | Интерполирует значения координат x и y для генерации новых точек. 321 | 322 | Аргументы: 323 | - x: Список значений координаты x. 324 | - y: Список значений координаты y. 325 | - n_points: Желаемое количество новых точек для интерполяции. 326 | 327 | Возвращает: 328 | - x1: Интерполированные значения координаты x. 329 | - y1: Интерполированные значения координаты y. 330 | """ 331 | 332 | interp_x = interp1d(np.arange(len(x)), x) 333 | interp_y = interp1d(np.arange(len(y)), y) 334 | new_indices = np.linspace(0, len(x) - 1, n_points) 335 | x1 = interp_x(new_indices) 336 | y1 = interp_y(new_indices) 337 | return x1, y1 338 | 339 | def calculate_opposite_interpolate(self, side_1, side_2, smooth=False): 340 | """ 341 | Вычисляет интерполированные координаты противоположных сторон. 342 | 343 | Аргументы: 344 | - side_1: Список координат первой стороны. 345 | - side_2: Список координат второй стороны. 346 | - smooth: Флаг, указывающий на необходимость сглаживания данных. 347 | 348 | Возвращает: 349 | - new_side_1_x: Интерполированные значения координаты x для первой стороны. 350 | - new_side_1_y: Интерполированные значения координаты y для первой стороны. 351 | - new_side_2_x: Интерполированные значения координаты x для второй стороны. 352 | - new_side_2_y: Интерполированные значения координаты y для второй стороны. 353 | """ 354 | 355 | side_1_x, side_1_y = zip(*side_1) 356 | side_2_x, side_2_y = zip(*side_2) 357 | pts_num_lr = np.max((len(side_1_x), len(side_2_x))) 358 | 359 | new_side_1_x, new_side_1_y = self.interpolate(side_1_x, side_1_y, pts_num_lr) 360 | new_side_2_x, new_side_2_y = self.interpolate(side_2_x, side_2_y, pts_num_lr) 361 | 362 | if smooth: 363 | window = 51 364 | coef = 5 365 | new_side_1_x = savgol_filter(new_side_1_x, window, coef) 366 | new_side_1_y = savgol_filter(new_side_1_y, window, coef) 367 | new_side_2_x = savgol_filter(new_side_2_x, window, coef) 368 | new_side_2_y = savgol_filter(new_side_2_y, window, coef) 369 | return new_side_1_x, new_side_1_y, new_side_2_x, new_side_2_y 370 | 371 | def get_lines_of_approximation(self, num, l_x, l_y, r_x, r_y, b_x, b_y, t_x, t_y): 372 | """ 373 | Получает линии аппроксимации для указанного числа точек. 374 | 375 | Аргументы: 376 | - num: Число точек для аппроксимации. 377 | - l_x: Интерполированные значения координаты x для левой стороны. 378 | - l_y: Интерполированные значения координаты y для левой стороны. 379 | - r_x: Интерполированные значения координаты x для правой стороны. 380 | - r_y: Интерполированные значения координаты y для правой стороны. 381 | - b_x: Интерполированные значения координаты x для нижней стороны. 382 | - b_y: Интерполированные значения координаты y для нижней стороны. 383 | - t_x: Интерполированные значения координаты x для верхней стороны. 384 | - t_y: Интерполированные значения координаты y для верхней стороны. 385 | 386 | Возвращает: 387 | - linfit_x_bt: Функция аппроксимации для координаты x по горизонтали. 388 | - linfit_y_bt: Функция аппроксимации для координаты y по горизонтали. 389 | - linfit_x_lr: Функция аппроксимации для координаты x по вертикали. 390 | - linfit_y_lr: Функция аппроксимации для координаты y по вертикали. 391 | """ 392 | 393 | linfit_x_lr = interp1d([1, num], np.vstack((l_x, r_x)), axis=0) 394 | linfit_y_lr = interp1d([1, num], np.vstack((l_y, r_y)), axis=0) 395 | 396 | linfit_x_bt = interp1d([1, num], np.vstack((b_x, t_x)), axis=0) 397 | linfit_y_bt = interp1d([1, num], np.vstack((b_y, t_y)), axis=0) 398 | 399 | for i in range(1, num + 1): 400 | plt.imshow(self.initial_image) 401 | plt.plot(linfit_x_bt(i), linfit_y_bt(i)) 402 | plt.plot(linfit_x_lr(i), linfit_y_lr(i)) 403 | plt.show() 404 | return linfit_x_bt, linfit_y_bt, linfit_x_lr, linfit_y_lr 405 | 406 | @staticmethod 407 | def calculate_approximation(p, x, y, shape_flag): 408 | """ 409 | Вычисляет аппроксимацию для заданных коэффициентов и координат. 410 | 411 | Аргументы: 412 | - p: Коэффициенты аппроксимации. 413 | - x: Координата x. 414 | - y: Координата y. 415 | - shape_flag: Флаг формы (0 для горизонтальной аппроксимации, 1 для вертикальной). 416 | 417 | Возвращает: 418 | - Разность между аппроксимацией и исходными координатами. 419 | """ 420 | 421 | if shape_flag == 0: 422 | return p[0] * y ** 3 + p[1] * y ** 2 + p[2] * y + p[3] - x 423 | else: 424 | return p[0] * x ** 3 + p[1] * x ** 2 + p[2] * x + p[3] - y 425 | 426 | def get_interpolation_approximation(self, shape, num_ln, degree, num_p, linfit_x, linfit_y): 427 | """ 428 | Вычисляет аппроксимацию интерполяции. 429 | 430 | Аргументы: 431 | - shape: Размерность изображения (высота или ширина, в зависимости от флага формы). 432 | - num_ln: Количество линий интерполяции. 433 | - degree: Степень аппроксимации. 434 | - num_p: Количество точек аппроксимации. 435 | - linfit_x: Функция интерполяции для координаты x. 436 | - linfit_y: Функция интерполяции для координаты y. 437 | 438 | Возвращает: 439 | - Список полигонов с аппроксимированными координатами. 440 | """ 441 | 442 | polys_bt = [] 443 | shape_flag = 0 if shape == self.initial_image.shape[0] else 1 444 | 445 | for i in range(1, num_ln + 1): 446 | x = linfit_x(i) 447 | y = linfit_y(i) 448 | 449 | p0 = np.zeros(degree + 1) 450 | result = least_squares(self.calculate_approximation, p0, args=(x, y, shape_flag)) 451 | if shape_flag == 0: 452 | y_new = np.linspace(0, shape, num_p) 453 | x_new = result.x[0] * y_new ** 3 + result.x[1] * y_new ** 2 + result.x[2] * y_new + result.x[3] 454 | else: 455 | x_new = np.linspace(0, shape, num_p) 456 | y_new = result.x[0] * x_new ** 3 + result.x[1] * x_new ** 2 + result.x[2] * x_new + result.x[3] 457 | 458 | poly_coor = np.hstack((np.reshape(x_new, (x_new.shape[0], 1)), np.reshape(y_new, (y_new.shape[0], 1)))) 459 | polys_bt.append(poly_coor) 460 | plt.plot(x_new, y_new) 461 | 462 | return polys_bt 463 | 464 | @staticmethod 465 | def flatten(t): 466 | """ 467 | Выполняет сглаживание вложенных списков. 468 | 469 | Аргументы: 470 | - t: вложенный список 471 | 472 | Возвращает: 473 | - Одномерный список, полученный путем сглаживания вложенных списков 474 | """ 475 | return list(itertools.chain.from_iterable(t)) 476 | 477 | def find_intersection(self, curve_1, curve_2, tolerance): 478 | """ 479 | Находит точку пересечения двух кривых. 480 | 481 | Аргументы: 482 | - curve_1: список координат точек первой кривой 483 | - curve_2: список координат точек второй кривой 484 | - tolerance: пороговое значение для определения близости точек 485 | 486 | Возвращает: 487 | - Координаты точки пересечения двух кривых 488 | """ 489 | 490 | # Строим деревья для каждой кривой 491 | tree_big = spatial.cKDTree(curve_1) 492 | tree_small = spatial.cKDTree(curve_2) 493 | # Выполняем запрос на поиск ближайших соседей между деревьями с использованием порогового значения tolerance 494 | inds = tree_small.query_ball_tree(tree_big, r=tolerance) 495 | # Объединяем индексы найденных соседей в один список 496 | small_inds = self.flatten([i for i in inds if i]) 497 | # Получаем уникальные индексы 498 | unique_inds = np.unique(small_inds).tolist() 499 | # Получаем координаты точек из curve_1, соответствующие уникальным индексам 500 | unique_coords = [curve_1[i] for i in unique_inds] 501 | return np.mean(np.asarray(unique_coords), axis=0) 502 | 503 | def get_intersection_interpolation_approximation(self, length, width, polys_bt, polys_lr): 504 | """ 505 | Вычисляет пересечения между линиями интерполяции аппроксимации. 506 | 507 | Аргументы: 508 | - length: длина изображения 509 | - width: ширина изображения 510 | - polys_bt: список полиномиальных координат для линий, идущих в направлении bottom-top 511 | - polys_lr: список полиномиальных координат для линий, идущих в направлении left-right 512 | 513 | Возвращает: 514 | - Список координат пересечений линий интерполяции аппроксимации 515 | """ 516 | 517 | plt.imshow(self.initial_image) 518 | points_set = [] 519 | # Вычисление значения tolerance как округленного значения 1% от длины или ширины изображения 520 | # (в зависимости от того, какая величина больше) 521 | tolerance = ceil(length * 0.01) if length > width else ceil(width * 0.01) 522 | if tolerance < 20: 523 | tolerance = 20.0 524 | print(f'TOLERANCE {tolerance}') 525 | for line1, line2 in itertools.product(polys_bt, polys_lr): 526 | # Поиск пересечения между текущими линиями line1 и line2, используя заданный tolerance 527 | a = self.find_intersection(line1, line2, tolerance) 528 | points_set.append([a[0], a[1]]) 529 | plt.plot(round(a[0]), round(a[1]), 'ro') 530 | plt.show() 531 | return points_set 532 | 533 | def get_fix_grid_coordinates(self, length, width, num_of_line): 534 | """ 535 | Вычисляет координаты фиксированной сетки на изображении. 536 | 537 | Аргументы: 538 | - length: длина изображения 539 | - width: ширина изображения 540 | - num_of_line: количество линий в сетке 541 | 542 | Возвращает: 543 | - Массив координат фиксированной сетки 544 | 545 | """ 546 | 547 | x = np.linspace(0, width, num_of_line) 548 | y = np.linspace(0, length, num_of_line) 549 | xv, yv = np.meshgrid(x, y) 550 | XYpairs = np.dstack([xv, yv]).reshape(-1, 2) 551 | plt.imshow(self.initial_image) 552 | plt.plot(XYpairs[:, 0], XYpairs[:, 1], marker='.', color='k', linestyle='none') 553 | plt.show() 554 | return XYpairs 555 | 556 | @staticmethod 557 | def crop_image(img): 558 | """ 559 | Обрезает изображение в соответствии с контурами на нем. 560 | 561 | Параметры: 562 | - img: исходное изображение (массив пикселей). 563 | 564 | Возвращает: 565 | - cropped_img: обрезанное изображение. 566 | """ 567 | 568 | gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) 569 | # Применение пороговой обработки для получения бинарного изображения 570 | thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_OTSU + cv2.THRESH_BINARY)[1] 571 | # Поиск контуров на бинарном изображении 572 | contour = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) 573 | contour = contour[0] if len(contour) == 2 else contour[1] 574 | # Сортировка контуров по площади в порядке убывания 575 | contour = sorted(contour, key=cv2.contourArea, reverse=True) 576 | if len(contour) > 0: 577 | # Получение ограничивающего прямоугольника для самого большого контура 578 | x, y, w, h = cv2.boundingRect(contour[0]) 579 | # Обрезка изображения согласно ограничивающему прямоугольнику 580 | img = img[y:y + h, x:x + w] 581 | return img 582 | 583 | def remap_image(self, height, width, points_set, xy_pairs): 584 | """ 585 | Переотображает исходное изображение с использованием заданных точек. 586 | 587 | Параметры: 588 | - height: высота переотображаемого изображения. 589 | - width: ширина переотображаемого изображения. 590 | - points_set: список точек, определяющих переотображение. 591 | - xy_pairs: координаты сетки для переотображения. 592 | 593 | Возвращает: 594 | - crop_warped_image_BGR: переотображенное и обрезанное изображение в формате BGR. 595 | """ 596 | 597 | img = self.initial_image.copy() 598 | # Создание комплексных чисел для формирования сетки 599 | comp_w = complex(f'{width}j') 600 | comp_h = complex(f'{height}j') 601 | # Создание сетки 602 | grid_x, grid_y = np.mgrid[0:(height - 1):comp_h, 0:(width - 1):comp_w] 603 | # Переотображение точек на сетку с использованием кубической интерполяции 604 | grid_z = griddata(np.fliplr(xy_pairs), np.fliplr(points_set), (grid_x, grid_y), method='cubic') 605 | # Извлечение координат x и y из переотображенных точек 606 | map_x = np.append([], [ar[:, 1] for ar in grid_z]).reshape(height, width) 607 | map_y = np.append([], [ar[:, 0] for ar in grid_z]).reshape(height, width) 608 | # Преобразование координат в тип float32 609 | map_x_32 = map_x.astype('float32') 610 | map_y_32 = map_y.astype('float32') 611 | # Переотображение исходного изображения с использованием полученных координат 612 | warped_image = cv2.remap(img, map_x_32, map_y_32, cv2.INTER_CUBIC) 613 | warped_image = np.fliplr(np.rot90(warped_image, 2)) 614 | # Обрезка переотображенного изображения 615 | # crop_warped_image = self.crop_image(warped_image) 616 | crop_warped_image_BGR = cv2.cvtColor(warped_image, cv2.COLOR_RGB2BGR) 617 | plt.subplot(121, title='before') 618 | plt.imshow(img) 619 | plt.axis('off') 620 | plt.subplot(122, title='after') 621 | plt.imshow(warped_image) 622 | plt.axis('off') 623 | plt.tight_layout() 624 | plt.show() 625 | return crop_warped_image_BGR 626 | --------------------------------------------------------------------------------