table_line.py 144 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340334133423343334433453346334733483349335033513352335333543355335633573358335933603361336233633364336533663367336833693370337133723373337433753376337733783379338033813382338333843385338633873388338933903391339233933394339533963397339833993400340134023403340434053406340734083409341034113412341334143415341634173418341934203421342234233424342534263427342834293430343134323433343434353436343734383439344034413442344334443445344634473448344934503451345234533454345534563457345834593460346134623463346434653466346734683469347034713472347334743475347634773478347934803481348234833484348534863487348834893490349134923493349434953496349734983499350035013502350335043505350635073508350935103511351235133514351535163517351835193520352135223523352435253526352735283529353035313532353335343535353635373538353935403541354235433544354535463547354835493550355135523553355435553556355735583559356035613562356335643565356635673568356935703571357235733574357535763577357835793580358135823583358435853586358735883589359035913592359335943595359635973598359936003601360236033604360536063607360836093610361136123613361436153616361736183619362036213622362336243625362636273628362936303631363236333634363536363637363836393640364136423643364436453646364736483649365036513652365336543655365636573658365936603661366236633664366536663667366836693670367136723673367436753676367736783679368036813682368336843685368636873688368936903691369236933694369536963697369836993700370137023703370437053706370737083709
  1. #!/usr/bin/env python3
  2. # -*- coding: utf-8 -*-
  3. """
  4. Created on Thu Sep 9 23:11:51 2020
  5. table line detect
  6. @author: chineseocr
  7. """
  8. import copy
  9. import io
  10. import logging
  11. import pickle
  12. import sys
  13. import traceback
  14. import tensorflow as tf
  15. import tensorflow.keras.backend as K
  16. from tensorflow.keras.models import Model
  17. from tensorflow.keras.layers import Input, concatenate, Conv2D, MaxPooling2D, BatchNormalization, UpSampling2D
  18. from tensorflow.keras.layers import LeakyReLU
  19. from otr.utils import letterbox_image, get_table_line, adjust_lines, line_to_line, draw_boxes
  20. import numpy as np
  21. import cv2
  22. import time
  23. from format_convert import _global
  24. from format_convert.utils import log
  25. def dice_coef(y_true, y_pred, smooth=1e-5):
  26. y_true_f = K.flatten(y_true)
  27. y_pred_f = K.flatten(y_pred)
  28. intersection = K.sum(y_true_f * y_pred_f)
  29. return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)
  30. def dice_coef_loss():
  31. def dice_coef_loss_fixed(y_true, y_pred):
  32. return -dice_coef(y_true, y_pred)
  33. return dice_coef_loss_fixed
  34. def focal_loss(gamma=3., alpha=.5):
  35. # 3 0.85 2000e acc-0.6 p-0.99 r-0.99 val_acc-0.56 val_p-0.86 val_r-0.95
  36. # 2 0.85 double_gpu acc-
  37. # 3 0.25 gpu 50e acc-0.5 p-0.99 r-0.99 val_acc-0.45 val_p-0.96 val_r-0.88
  38. # 2 0.25 gpu acc-
  39. # 3 0.5 double_gpu acc-0.6 p-0.99 r-0.99 val_acc-0.60 val_p-0.93 val_r-0.93
  40. def focal_loss_fixed(y_true, y_pred):
  41. pt_1 = tf.where(tf.equal(y_true, 1), y_pred, tf.ones_like(y_pred))
  42. pt_0 = tf.where(tf.equal(y_true, 0), y_pred, tf.zeros_like(y_pred))
  43. return -K.sum(alpha * K.pow(1. - pt_1, gamma) * K.log(K.epsilon()+pt_1))-K.sum((1-alpha) * K.pow( pt_0, gamma) * K.log(1. - pt_0 + K.epsilon()))
  44. return focal_loss_fixed
  45. def table_net_large(input_shape=(1152, 896, 3), num_classes=1):
  46. inputs = Input(shape=input_shape)
  47. # 512
  48. use_bias = False
  49. down0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(inputs)
  50. down0a = BatchNormalization()(down0a)
  51. down0a = LeakyReLU(alpha=0.1)(down0a)
  52. down0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(down0a)
  53. down0a = BatchNormalization()(down0a)
  54. down0a = LeakyReLU(alpha=0.1)(down0a)
  55. down0a_pool = MaxPooling2D((2, 2), strides=(2, 2))(down0a)
  56. # 256
  57. down0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(down0a_pool)
  58. down0 = BatchNormalization()(down0)
  59. down0 = LeakyReLU(alpha=0.1)(down0)
  60. down0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(down0)
  61. down0 = BatchNormalization()(down0)
  62. down0 = LeakyReLU(alpha=0.1)(down0)
  63. down0_pool = MaxPooling2D((2, 2), strides=(2, 2))(down0)
  64. # 128
  65. down1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(down0_pool)
  66. down1 = BatchNormalization()(down1)
  67. down1 = LeakyReLU(alpha=0.1)(down1)
  68. down1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(down1)
  69. down1 = BatchNormalization()(down1)
  70. down1 = LeakyReLU(alpha=0.1)(down1)
  71. down1_pool = MaxPooling2D((2, 2), strides=(2, 2))(down1)
  72. # 64
  73. down2 = Conv2D(128, (3, 3), padding='same', use_bias=use_bias)(down1_pool)
  74. down2 = BatchNormalization()(down2)
  75. down2 = LeakyReLU(alpha=0.1)(down2)
  76. down2 = Conv2D(128, (3, 3), padding='same', use_bias=use_bias)(down2)
  77. down2 = BatchNormalization()(down2)
  78. down2 = LeakyReLU(alpha=0.1)(down2)
  79. down2_pool = MaxPooling2D((2, 2), strides=(2, 2))(down2)
  80. # 32
  81. down3 = Conv2D(256, (3, 3), padding='same', use_bias=use_bias)(down2_pool)
  82. down3 = BatchNormalization()(down3)
  83. down3 = LeakyReLU(alpha=0.1)(down3)
  84. down3 = Conv2D(256, (3, 3), padding='same', use_bias=use_bias)(down3)
  85. down3 = BatchNormalization()(down3)
  86. down3 = LeakyReLU(alpha=0.1)(down3)
  87. down3_pool = MaxPooling2D((2, 2), strides=(2, 2))(down3)
  88. # 16
  89. down4 = Conv2D(512, (3, 3), padding='same', use_bias=use_bias)(down3_pool)
  90. down4 = BatchNormalization()(down4)
  91. down4 = LeakyReLU(alpha=0.1)(down4)
  92. down4 = Conv2D(512, (3, 3), padding='same', use_bias=use_bias)(down4)
  93. down4 = BatchNormalization()(down4)
  94. down4 = LeakyReLU(alpha=0.1)(down4)
  95. down4_pool = MaxPooling2D((2, 2), strides=(2, 2))(down4)
  96. # 8
  97. center = Conv2D(1024, (3, 3), padding='same', use_bias=use_bias)(down4_pool)
  98. center = BatchNormalization()(center)
  99. center = LeakyReLU(alpha=0.1)(center)
  100. center = Conv2D(1024, (3, 3), padding='same', use_bias=use_bias)(center)
  101. center = BatchNormalization()(center)
  102. center = LeakyReLU(alpha=0.1)(center)
  103. # center
  104. up4 = UpSampling2D((2, 2))(center)
  105. up4 = concatenate([down4, up4], axis=3)
  106. up4 = Conv2D(512, (3, 3), padding='same', use_bias=use_bias)(up4)
  107. up4 = BatchNormalization()(up4)
  108. up4 = LeakyReLU(alpha=0.1)(up4)
  109. up4 = Conv2D(512, (3, 3), padding='same', use_bias=use_bias)(up4)
  110. up4 = BatchNormalization()(up4)
  111. up4 = LeakyReLU(alpha=0.1)(up4)
  112. up4 = Conv2D(512, (3, 3), padding='same', use_bias=use_bias)(up4)
  113. up4 = BatchNormalization()(up4)
  114. up4 = LeakyReLU(alpha=0.1)(up4)
  115. # 16
  116. up3 = UpSampling2D((2, 2))(up4)
  117. up3 = concatenate([down3, up3], axis=3)
  118. up3 = Conv2D(256, (3, 3), padding='same', use_bias=use_bias)(up3)
  119. up3 = BatchNormalization()(up3)
  120. up3 = LeakyReLU(alpha=0.1)(up3)
  121. up3 = Conv2D(256, (3, 3), padding='same', use_bias=use_bias)(up3)
  122. up3 = BatchNormalization()(up3)
  123. up3 = LeakyReLU(alpha=0.1)(up3)
  124. up3 = Conv2D(256, (3, 3), padding='same', use_bias=use_bias)(up3)
  125. up3 = BatchNormalization()(up3)
  126. up3 = LeakyReLU(alpha=0.1)(up3)
  127. # 32
  128. up2 = UpSampling2D((2, 2))(up3)
  129. up2 = concatenate([down2, up2], axis=3)
  130. up2 = Conv2D(128, (3, 3), padding='same', use_bias=use_bias)(up2)
  131. up2 = BatchNormalization()(up2)
  132. up2 = LeakyReLU(alpha=0.1)(up2)
  133. up2 = Conv2D(128, (3, 3), padding='same', use_bias=use_bias)(up2)
  134. up2 = BatchNormalization()(up2)
  135. up2 = LeakyReLU(alpha=0.1)(up2)
  136. up2 = Conv2D(128, (3, 3), padding='same', use_bias=use_bias)(up2)
  137. up2 = BatchNormalization()(up2)
  138. up2 = LeakyReLU(alpha=0.1)(up2)
  139. # 64
  140. up1 = UpSampling2D((2, 2))(up2)
  141. up1 = concatenate([down1, up1], axis=3)
  142. up1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(up1)
  143. up1 = BatchNormalization()(up1)
  144. up1 = LeakyReLU(alpha=0.1)(up1)
  145. up1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(up1)
  146. up1 = BatchNormalization()(up1)
  147. up1 = LeakyReLU(alpha=0.1)(up1)
  148. up1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(up1)
  149. up1 = BatchNormalization()(up1)
  150. up1 = LeakyReLU(alpha=0.1)(up1)
  151. # 128
  152. up0 = UpSampling2D((2, 2))(up1)
  153. up0 = concatenate([down0, up0], axis=3)
  154. up0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(up0)
  155. up0 = BatchNormalization()(up0)
  156. up0 = LeakyReLU(alpha=0.1)(up0)
  157. up0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(up0)
  158. up0 = BatchNormalization()(up0)
  159. up0 = LeakyReLU(alpha=0.1)(up0)
  160. up0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(up0)
  161. up0 = BatchNormalization()(up0)
  162. up0 = LeakyReLU(alpha=0.1)(up0)
  163. # 256
  164. up0a = UpSampling2D((2, 2))(up0)
  165. up0a = concatenate([down0a, up0a], axis=3)
  166. up0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(up0a)
  167. up0a = BatchNormalization()(up0a)
  168. up0a = LeakyReLU(alpha=0.1)(up0a)
  169. up0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(up0a)
  170. up0a = BatchNormalization()(up0a)
  171. up0a = LeakyReLU(alpha=0.1)(up0a)
  172. up0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(up0a)
  173. up0a = BatchNormalization()(up0a)
  174. up0a = LeakyReLU(alpha=0.1)(up0a)
  175. # 512
  176. classify = Conv2D(num_classes, (1, 1), activation='sigmoid')(up0a)
  177. model = Model(inputs=inputs, outputs=classify)
  178. return model
  179. def table_net(input_shape=(1152, 896, 3), num_classes=1):
  180. inputs = Input(shape=input_shape)
  181. # 512
  182. use_bias = False
  183. down0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(inputs)
  184. down0a = BatchNormalization()(down0a)
  185. down0a = LeakyReLU(alpha=0.1)(down0a)
  186. down0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(down0a)
  187. down0a = BatchNormalization()(down0a)
  188. down0a = LeakyReLU(alpha=0.1)(down0a)
  189. down0a_pool = MaxPooling2D((2, 2), strides=(2, 2))(down0a)
  190. # 256
  191. down0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(down0a_pool)
  192. down0 = BatchNormalization()(down0)
  193. down0 = LeakyReLU(alpha=0.1)(down0)
  194. down0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(down0)
  195. down0 = BatchNormalization()(down0)
  196. down0 = LeakyReLU(alpha=0.1)(down0)
  197. down0_pool = MaxPooling2D((2, 2), strides=(2, 2))(down0)
  198. # 128
  199. down1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(down0_pool)
  200. down1 = BatchNormalization()(down1)
  201. down1 = LeakyReLU(alpha=0.1)(down1)
  202. down1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(down1)
  203. down1 = BatchNormalization()(down1)
  204. down1 = LeakyReLU(alpha=0.1)(down1)
  205. down1_pool = MaxPooling2D((2, 2), strides=(2, 2))(down1)
  206. # 64
  207. down2 = Conv2D(128, (3, 3), padding='same', use_bias=use_bias)(down1_pool)
  208. down2 = BatchNormalization()(down2)
  209. down2 = LeakyReLU(alpha=0.1)(down2)
  210. down2 = Conv2D(128, (3, 3), padding='same', use_bias=use_bias)(down2)
  211. down2 = BatchNormalization()(down2)
  212. down2 = LeakyReLU(alpha=0.1)(down2)
  213. down2_pool = MaxPooling2D((2, 2), strides=(2, 2))(down2)
  214. # 32
  215. up1 = UpSampling2D((2, 2))(down2)
  216. up1 = concatenate([down1, up1], axis=3)
  217. up1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(up1)
  218. up1 = BatchNormalization()(up1)
  219. up1 = LeakyReLU(alpha=0.1)(up1)
  220. up1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(up1)
  221. up1 = BatchNormalization()(up1)
  222. up1 = LeakyReLU(alpha=0.1)(up1)
  223. up1 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(up1)
  224. up1 = BatchNormalization()(up1)
  225. up1 = LeakyReLU(alpha=0.1)(up1)
  226. # 128
  227. up0 = UpSampling2D((2, 2))(up1)
  228. up0 = concatenate([down0, up0], axis=3)
  229. up0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(up0)
  230. up0 = BatchNormalization()(up0)
  231. up0 = LeakyReLU(alpha=0.1)(up0)
  232. up0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(up0)
  233. up0 = BatchNormalization()(up0)
  234. up0 = LeakyReLU(alpha=0.1)(up0)
  235. up0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(up0)
  236. up0 = BatchNormalization()(up0)
  237. up0 = LeakyReLU(alpha=0.1)(up0)
  238. # 256
  239. up0a = UpSampling2D((2, 2))(up0)
  240. up0a = concatenate([down0a, up0a], axis=3)
  241. up0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(up0a)
  242. up0a = BatchNormalization()(up0a)
  243. up0a = LeakyReLU(alpha=0.1)(up0a)
  244. up0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(up0a)
  245. up0a = BatchNormalization()(up0a)
  246. up0a = LeakyReLU(alpha=0.1)(up0a)
  247. up0a = Conv2D(16, (3, 3), padding='same', use_bias=use_bias)(up0a)
  248. up0a = BatchNormalization()(up0a)
  249. up0a = LeakyReLU(alpha=0.1)(up0a)
  250. # 512
  251. classify = Conv2D(num_classes, (1, 1), activation='sigmoid')(up0a)
  252. model = Model(inputs=inputs, outputs=classify)
  253. return model
  254. model = table_net((None, None, 3), 2)
  255. def draw_pixel(pred, prob=0.2, is_test=1):
  256. if not is_test:
  257. return
  258. else:
  259. import matplotlib.pyplot as plt
  260. _array = []
  261. for _h in range(len(pred)):
  262. _line = []
  263. for _w in range(len(pred[_h])):
  264. _prob = pred[_h][_w]
  265. if _prob[0] > prob:
  266. _line.append((0, 0, 255))
  267. elif _prob[1] > prob:
  268. _line.append((255, 0, 0))
  269. else:
  270. _line.append((255, 255, 255))
  271. _array.append(_line)
  272. plt.axis('off')
  273. plt.imshow(np.array(_array))
  274. plt.show()
  275. return
  276. def expansionAndShrinkage(pred,width=3):
  277. pred_array = np.array(pred)
  278. print("pred_array=====",pred_array.shape)
  279. _array = pred_array[...,0]
  280. _l = [_array]
  281. for _i in range(width):
  282. tmp_array = np.pad(_array[:-(_i+1),...],((_i+1,0),(0,0)))
  283. _l.append(tmp_array)
  284. for _i in range(width):
  285. tmp_array = np.pad(_array[_i+1:,...],((0,_i+1),(0,0)))
  286. _l.append(tmp_array)
  287. for _i in range(width):
  288. tmp_array = np.pad(_array[...,:-(_i+1)],((0,0),(_i+1,0)))
  289. _l.append(tmp_array)
  290. for _i in range(width):
  291. tmp_array = np.pad(_array[...,_i+1:],((0,0),(0,_i+1)))
  292. _l.append(tmp_array)
  293. for _a in _l:
  294. print(_a.shape)
  295. h_array = np.stack(_l,axis=0)
  296. h_array = np.max(h_array,axis=0,keepdims=False)
  297. _array = pred_array[...,1]
  298. _l = [_array]
  299. for _i in range(width):
  300. tmp_array = np.pad(_array[:-(_i+1),...],((_i+1,0),(0,0)))
  301. _l.append(tmp_array)
  302. for _i in range(width):
  303. tmp_array = np.pad(_array[_i+1:,...],((0,_i+1),(0,0)))
  304. _l.append(tmp_array)
  305. for _i in range(width):
  306. tmp_array = np.pad(_array[...,:-(_i+1)],((0,0),(_i+1,0)))
  307. _l.append(tmp_array)
  308. for _i in range(width):
  309. tmp_array = np.pad(_array[...,_i+1:],((0,0),(0,_i+1)))
  310. _l.append(tmp_array)
  311. v_array = np.stack(_l,axis=0)
  312. print("v_array=====",v_array.shape)
  313. v_array = np.max(v_array,axis=0,keepdims=False)
  314. print("h_array=====",h_array.shape)
  315. print("v_array=====",v_array.shape)
  316. last_array = np.stack([h_array,v_array],axis=-1)
  317. print("pred_array=====",last_array.shape)
  318. return last_array
  319. def getIOU(bbox0, bbox1):
  320. width = abs(max(bbox0[2],bbox1[2])-min(bbox0[0],bbox1[0]))-(abs(bbox0[2]-bbox0[0])+abs(bbox1[2]-bbox1[0]))
  321. height = abs(max(bbox0[3],bbox1[3])-min(bbox0[1],bbox1[1]))-(abs(bbox0[3]-bbox0[1])+abs(bbox1[3]-bbox1[1]))
  322. if width < 0 and height < 0:
  323. iou = abs(width*height/min(abs((bbox0[2]-bbox0[0])*(bbox0[3]-bbox0[1])),
  324. abs((bbox1[2]-bbox1[0])*(bbox1[3]-bbox1[1]))))
  325. # print("getIOU", iou)
  326. return iou
  327. return 0
  328. def lines_cluster(list_lines,line_width):
  329. after_len = 0
  330. log("len lines %d"%len(list_lines))
  331. append_width = line_width//2
  332. while 1:
  333. c_lines = []
  334. first_len = after_len
  335. for _line in list_lines:
  336. bbox = _line["bbox"]
  337. _find = False
  338. for c_l_i in range(len(c_lines)):
  339. c_l = c_lines[len(c_lines)-c_l_i-1]
  340. bbox1 = c_l["bbox"]
  341. bboxa = [max(0,bbox[0]-append_width),max(0,bbox[1]-append_width),bbox[2]+append_width,bbox[3]+append_width]
  342. bboxb = [max(0,bbox1[0]-append_width),max(0,bbox1[1]-append_width),bbox1[2]+append_width,bbox1[3]+append_width]
  343. _iou = getIOU(bboxa,bboxb)
  344. if _iou>0:
  345. new_bbox = [min(bbox[0],bbox[2],bbox1[0],bbox1[2]),min(bbox[1],bbox[3],bbox1[1],bbox1[3]),max(bbox[0],bbox[2],bbox1[0],bbox1[2]),max(bbox[1],bbox[3],bbox1[1],bbox1[3])]
  346. _find = True
  347. c_l["bbox"] = new_bbox
  348. break
  349. if not _find:
  350. c_lines.append(_line)
  351. after_len = len(c_lines)
  352. if first_len==after_len:
  353. break
  354. list_lines = c_lines
  355. return c_lines
  356. def points2lines(pred,sourceP_LB=True, prob=0.2, line_width=8, padding=3, min_len=10,
  357. cell_width=13):
  358. _time = time.time()
  359. log("starting points2lines")
  360. height = len(pred)
  361. width = len(pred[0])
  362. _sum = list(np.sum(np.array((pred[...,0]>prob)).astype(int),axis=1))
  363. h_index = -1
  364. h_lines = []
  365. v_lines = []
  366. _step = line_width
  367. while 1:
  368. h_index += 1
  369. if h_index>=height:
  370. break
  371. w_index = -1
  372. if sourceP_LB:
  373. h_i = height-1-h_index
  374. else:
  375. h_i = h_index
  376. _start = None
  377. if _sum[h_index]<min_len:
  378. continue
  379. while 1:
  380. w_index += _step
  381. if w_index>=width:
  382. break
  383. _h,_v = pred[h_i][w_index]
  384. if _h>prob:
  385. if _start is None:
  386. _start = w_index
  387. else:
  388. if _start is not None:
  389. _end = w_index-1
  390. _bbox = [_start,h_i,_end,h_i]
  391. _dict = {"bbox":_bbox}
  392. h_lines.append(_dict)
  393. _start = None
  394. w_index -= _step//2
  395. log("starting points2lines 1")
  396. w_index = -1
  397. _sum = list(np.sum(np.array((pred[...,1]>prob)).astype(int),axis=0))
  398. _step = line_width
  399. while 1:
  400. w_index += 1
  401. if w_index>=width:
  402. break
  403. if _sum[w_index]<min_len:
  404. continue
  405. h_index = -1
  406. _start = None
  407. while 1:
  408. h_index += _step
  409. if h_index>=height:
  410. break
  411. if sourceP_LB:
  412. h_i = height-1-h_index
  413. else:
  414. h_i = h_index
  415. _h,_v = pred[h_index][w_index]
  416. if _v>prob:
  417. if _start is None:
  418. _start = h_i
  419. else:
  420. if _start is not None:
  421. _end = last_h
  422. _bbox = [w_index,_start,w_index,_end]
  423. _dict = {"bbox":_bbox}
  424. v_lines.append(_dict)
  425. _start = None
  426. h_index -= _step//2
  427. last_h = h_i
  428. log("starting points2lines 2")
  429. for _line in h_lines:
  430. _bbox = _line["bbox"]
  431. _bbox = [max(_bbox[0]-2,0),(_bbox[1]+_bbox[3])/2,_bbox[2]+2,(_bbox[1]+_bbox[3])/2]
  432. _line["bbox"] = _bbox
  433. for _line in v_lines:
  434. _bbox = _line["bbox"]
  435. _bbox = [(_bbox[0]+_bbox[2])/2,max(_bbox[1]-2,0),(_bbox[0]+_bbox[2])/2,_bbox[3]+2]
  436. _line["bbox"] = _bbox
  437. h_lines = lines_cluster(h_lines,line_width=line_width)
  438. v_lines = lines_cluster(v_lines,line_width=line_width)
  439. list_line = []
  440. for _line in h_lines:
  441. _bbox = _line["bbox"]
  442. _bbox = [max(_bbox[0]-1,0),(_bbox[1]+_bbox[3])/2,_bbox[2]+1,(_bbox[1]+_bbox[3])/2]
  443. list_line.append(_bbox)
  444. for _line in v_lines:
  445. _bbox = _line["bbox"]
  446. _bbox = [(_bbox[0]+_bbox[2])/2,max(_bbox[1]-1,0),(_bbox[0]+_bbox[2])/2,_bbox[3]+1]
  447. list_line.append(_bbox)
  448. log("points2lines cost %.2fs"%(time.time()-_time))
  449. # import matplotlib.pyplot as plt
  450. # plt.figure()
  451. # for _line in list_line:
  452. # x0,y0,x1,y1 = _line
  453. # plt.plot([x0,x1],[y0,y1])
  454. # for _line in list_line:
  455. # x0,y0,x1,y1 = _line.bbox
  456. # plt.plot([x0,x1],[y0,y1])
  457. # for point in list_crosspoints:
  458. # plt.scatter(point.get("point")[0],point.get("point")[1])
  459. # plt.show()
  460. return list_line
  461. def points2lines_bak(pred, sourceP_LB=True, prob=0.2, line_width=7, padding=3, min_len=10,
  462. cell_width=13):
  463. def inBbox(bbox,point,line_width):
  464. x,y = point
  465. if x>=bbox[0]-line_width and x<=bbox[2]+line_width and y>=bbox[1]-line_width and y<=bbox[3]+line_width:
  466. return True,[min(x,bbox[0]),min(y,bbox[1]),max(x,bbox[2]),max(y,bbox[3])]
  467. return False,None
  468. _time = time.time()
  469. height = len(pred)
  470. width = len(pred[0])
  471. clust_horizontal = []
  472. clust_vertical = []
  473. h_index = -1
  474. _step = line_width
  475. _sum = list(np.sum(np.array((pred[...,1]>prob)).astype(int),axis=0))
  476. _last = False
  477. _current = False
  478. while 1:
  479. h_index += 2
  480. if h_index>=height:
  481. break
  482. w_index = -1
  483. if sourceP_LB:
  484. h_i = height-1-h_index
  485. else:
  486. h_i = h_index
  487. while 1:
  488. w_index += 2
  489. if w_index>=width:
  490. break
  491. if _sum[w_index]<min_len:
  492. continue
  493. _h,_v = pred[h_index][w_index]
  494. if _v>prob:
  495. _find = False
  496. _point = (w_index,h_i)
  497. for l_h_i in range(len(clust_vertical)):
  498. l_h = clust_vertical[len(clust_vertical)-l_h_i-1]
  499. bbox = l_h.get("bbox")
  500. b_in,_bbox = inBbox(bbox,_point,line_width)
  501. if b_in:
  502. _find = True
  503. l_h.get("points").append(_point)
  504. l_h["bbox"] = _bbox
  505. break
  506. if not _find:
  507. clust_vertical.append({"points":[_point],"bbox":[w_index,h_i,w_index,h_i]})
  508. w_index = -1
  509. _sum = list(np.sum(np.array((pred[...,0]>prob)).astype(int),axis=1))
  510. while 1:
  511. w_index += 2
  512. if w_index>=width:
  513. break
  514. h_index = -1
  515. while 1:
  516. h_index += 2
  517. if h_index>=height:
  518. break
  519. if _sum[h_index]<min_len:
  520. continue
  521. if sourceP_LB:
  522. h_i = height-1-h_index
  523. else:
  524. h_i = h_index
  525. _h,_v = pred[h_index][w_index]
  526. if _h>prob:
  527. _find = False
  528. _point = (w_index,h_i)
  529. for l_h_i in range(len(clust_horizontal)):
  530. l_h = clust_horizontal[len(clust_horizontal)-l_h_i-1]
  531. bbox = l_h.get("bbox")
  532. b_in,_bbox = inBbox(bbox,_point,line_width)
  533. if b_in:
  534. _find = True
  535. l_h.get("points").append(_point)
  536. l_h["bbox"] = _bbox
  537. break
  538. if not _find:
  539. clust_horizontal.append({"points":[_point],"bbox":[w_index,h_i,w_index,h_i]})
  540. tmp_vertical = []
  541. for _dict in clust_vertical:
  542. _bbox = _dict.get("bbox")
  543. if _bbox[2]-_bbox[0]>=min_len or _bbox[3]-_bbox[1]>=min_len:
  544. tmp_vertical.append([(_bbox[0]+_bbox[2])/2,_bbox[1]-padding,(_bbox[0]+_bbox[2])/2,_bbox[3]+padding])
  545. tmp_horizontal = []
  546. for _dict in clust_horizontal:
  547. _bbox = _dict.get("bbox")
  548. if _bbox[2]-_bbox[0]>=min_len or _bbox[3]-_bbox[1]>=min_len:
  549. tmp_horizontal.append([_bbox[0]-padding,(_bbox[1]+_bbox[3])/2,_bbox[2]+padding,(_bbox[1]+_bbox[3])/2])
  550. #merge lines
  551. tmp_vertical.sort(key=lambda x:x[3],reverse=True)
  552. tmp_horizontal.sort(key=lambda x:x[0])
  553. pop_index = []
  554. final_vertical = []
  555. for _line in tmp_vertical:
  556. _find = False
  557. x0,y0,x1,y1 = _line
  558. for _line2 in final_vertical:
  559. x2,y2,x3,y3 = _line2
  560. if abs(x0-x2)<line_width and abs(y0-y3)<cell_width or abs(y1-y2)<cell_width:
  561. _find = True
  562. final_vertical.append([x0,min(y0,y2),x1,max(y1,y3)])
  563. break
  564. if not _find:
  565. final_vertical.append(_line)
  566. final_horizontal = []
  567. for _line in tmp_horizontal:
  568. _find = False
  569. x0,y0,x1,y1 = _line
  570. for _line2 in final_horizontal:
  571. x2,y2,x3,y3 = _line2
  572. if abs(y0-y2)<line_width and abs(x0-x3)<cell_width or abs(x1-x2)<cell_width:
  573. _find = True
  574. final_horizontal.append([min(x0,x2),y0,max(x1,x3),y1])
  575. break
  576. if not _find:
  577. final_horizontal.append(_line)
  578. list_line = []
  579. for _line in final_vertical:
  580. list_line.append(_line)
  581. for _line in final_horizontal:
  582. list_line.append(_line)
  583. log("points2lines cost %.2fs"%(time.time()-_time))
  584. # import matplotlib.pyplot as plt
  585. # plt.figure()
  586. # for _line in list_line:
  587. # x0,y0,x1,y1 = _line
  588. # plt.plot([x0,x1],[y0,y1])
  589. # for _line in list_line:
  590. # x0,y0,x1,y1 = _line.bbox
  591. # plt.plot([x0,x1],[y0,y1])
  592. # for point in list_crosspoints:
  593. # plt.scatter(point.get("point")[0],point.get("point")[1])
  594. # plt.show()
  595. return list_line
  596. def get_line_from_binary_image(image_np, point_value=1, axis=0):
  597. """
  598. 根据像素点的变化,将像素点为特定值的转化为line,即找出端点坐标。
  599. 需要二值化的图。
  600. 仅支持竖线横线。
  601. :param image_np: numpy格式 image
  602. :param point_value: 像素点的特定值
  603. :param axis: 是否是行,否则为列
  604. :return: line list
  605. """
  606. def get_axis_points(_list, axis=0):
  607. _list.sort(key=lambda x: (x[1-axis], x[axis]))
  608. standard_axis = points[axis][1-axis]
  609. axis_points = []
  610. sub_points = []
  611. for p in _list:
  612. if p[1-axis] == standard_axis:
  613. sub_points.append(p)
  614. else:
  615. standard_axis = p[1-axis]
  616. if sub_points:
  617. axis_points.append(sub_points)
  618. sub_points = []
  619. # 最后一行/列
  620. if sub_points:
  621. axis_points.append(sub_points)
  622. return axis_points
  623. def get_axis_lines(_list, axis=0):
  624. # 逐行/列判断,一行/列可能多条横线/竖线
  625. points_lines = []
  626. for axis_list in _list:
  627. sub_line = [axis_list[0]]
  628. for p in axis_list:
  629. # 设置基准点
  630. standard_p = sub_line[-1]
  631. # 判断连续
  632. if p[axis] - standard_p[axis] == 1:
  633. sub_line.append(p)
  634. else:
  635. points_lines.append(sub_line)
  636. sub_line = [p]
  637. # 最后一行/列
  638. if sub_line:
  639. points_lines.append(sub_line)
  640. # 许多点组成的line转为两点line
  641. lines = []
  642. for line in points_lines:
  643. line.sort(key=lambda x: (x[axis], x[1-axis]))
  644. lines.append([line[0][0], line[0][1], line[-1][0], line[-1][1]])
  645. return lines
  646. # 取值大于point_value的点的坐标
  647. ys, xs = np.where(image_np >= point_value)
  648. points = [[xs[i], ys[i]] for i in range(len(xs))]
  649. # 提出所有相同x或相同y的点
  650. # 提取行/列
  651. axis_points = get_axis_points(points, axis)
  652. # 提取每行/列的横线/竖线
  653. axis_lines = get_axis_lines(axis_points, axis)
  654. # print("axis_lines", axis_lines)
  655. return axis_lines
  656. def table_preprocess(img_data, prob=0.2):
  657. try:
  658. log("into table_preprocess, prob is " + str(prob))
  659. start_time = time.time()
  660. # 二进制数据流转np.ndarray [np.uint8: 8位像素]
  661. img = cv2.imdecode(np.frombuffer(img_data, np.uint8), cv2.IMREAD_COLOR)
  662. # 将bgr转为rbg
  663. image_np = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
  664. # 模型输入
  665. inputs = np.array([image_np])
  666. # # 压缩numpy
  667. # compressed_array = io.BytesIO()
  668. # np.savez_compressed(compressed_array, inputs)
  669. # compressed_array.seek(0)
  670. # inputs_compressed = compressed_array.read()
  671. log("otr preprocess time: " + str(round(float(time.time()-start_time), 4)) + "s")
  672. return image_np, inputs
  673. except Exception as e:
  674. log("table pre process failed!")
  675. traceback.print_exc()
  676. return [-1], [-1]
  677. def table_postprocess(img_new, pred, prob=0.2, is_test=0):
  678. try:
  679. # 横线预测结果
  680. # row_pred = pred[..., 0] > hprob
  681. # row_pred = row_pred.astype(np.uint8)
  682. # # 竖线预测结果
  683. # col_pred = pred[..., 1] > vprob
  684. # col_pred = col_pred.astype(np.uint8)
  685. # # 打印模型输出
  686. # cv2.imshow("predict", (col_pred+row_pred)*255)
  687. # cv2.waitKey(0)
  688. start_time = time.time()
  689. list_line = points2lines(pred, False, prob=prob)
  690. mat_plot(list_line, "points2lines", is_test)
  691. log("points2lines " + str(time.time()-start_time))
  692. # 清除短线
  693. # print(img_new.shape)
  694. start_time = time.time()
  695. list_line = delete_short_lines(list_line, img_new.shape)
  696. mat_plot(list_line, "delete_short_lines", is_test)
  697. log("delete_short_lines " + str(time.time()-start_time))
  698. # # 清除无交点线--无需清除,会影响后面的结果
  699. # start_time = time.time()
  700. # list_line = delete_no_cross_lines(list_line)
  701. # mat_plot(list_line, "delete_no_cross_lines", is_test)
  702. # log("delete_no_cross_lines " + str(time.time()-start_time))
  703. # 分成横竖线
  704. start_time = time.time()
  705. list_rows = []
  706. list_cols = []
  707. for line in list_line:
  708. if line[0] == line[2]:
  709. list_cols.append(line)
  710. elif line[1] == line[3]:
  711. list_rows.append(line)
  712. log("divide rows and cols " + str(time.time()-start_time))
  713. # 合并错开线
  714. start_time = time.time()
  715. list_rows = merge_line(list_rows, axis=0)
  716. list_cols = merge_line(list_cols, axis=1)
  717. mat_plot(list_rows+list_cols, "merge_line", is_test)
  718. log("merge_line " + str(time.time()-start_time))
  719. # 计算交点、分割线
  720. start_time = time.time()
  721. cross_points = get_points(list_rows, list_cols, (img_new.shape[0], img_new.shape[1]))
  722. if not cross_points:
  723. return []
  724. log("get_points " + str(time.time()-start_time))
  725. # 清掉外围的没用的线
  726. # list_rows, list_cols = delete_outline(list_rows, list_cols, cross_points)
  727. # mat_plot(list_rows+list_cols, "delete_outline", is_test)
  728. # 多个表格分割线
  729. start_time = time.time()
  730. list_rows, list_cols = fix_in_split_lines(list_rows, list_cols, img_new)
  731. split_lines, split_y = get_split_line(cross_points, list_cols, img_new)
  732. log("get_split_line " + str(time.time()-start_time))
  733. # 修复边框
  734. start_time = time.time()
  735. new_rows, new_cols, long_rows, long_cols = fix_outline(img_new, list_rows, list_cols, cross_points,
  736. split_y)
  737. # 如有补线
  738. if new_rows or new_cols:
  739. # 连接至补线的延长线
  740. if long_rows:
  741. list_rows = long_rows
  742. if long_cols:
  743. list_cols = long_cols
  744. # 新的补线
  745. if new_rows:
  746. list_rows += new_rows
  747. if new_cols:
  748. list_cols += new_cols
  749. list_rows, list_cols = fix_in_split_lines(list_rows, list_cols, img_new)
  750. # 修复边框后重新计算交点、分割线
  751. cross_points = get_points(list_rows, list_cols, (img_new.shape[0], img_new.shape[1]))
  752. cv_plot(cross_points, img_new.shape, 0, is_test)
  753. split_lines, split_y = get_split_line(cross_points, list_cols, img_new)
  754. print("fix new split_y", split_y)
  755. print("fix new split_lines", split_lines)
  756. # 修复内部缺线
  757. # cross_points = fix_inner(list_rows, list_cols, cross_points, split_y)
  758. # if not cross_points:
  759. # return []
  760. mat_plot(list_rows+list_cols, "fix_outline", is_test)
  761. split_lines_show = []
  762. for _l in split_lines:
  763. split_lines_show.append([_l[0][0], _l[0][1], _l[1][0], _l[1][1]])
  764. mat_plot(split_lines_show+list_cols,
  765. "split_lines", is_test)
  766. log("fix_outline " + str(time.time()-start_time))
  767. # 修复表格4个角
  768. start_time = time.time()
  769. list_rows, list_cols = fix_corner(list_rows, list_cols, split_y, threshold=0)
  770. mat_plot(list_rows+list_cols, "fix_corner", is_test)
  771. log("fix_corner " + str(time.time()-start_time))
  772. # 修复内部缺线
  773. start_time = time.time()
  774. list_rows, list_cols = fix_inner(list_rows, list_cols, cross_points, split_y)
  775. mat_plot(list_rows+list_cols, "fix_inner", is_test)
  776. log("fix_inner " + str(time.time()-start_time))
  777. # 合并错开线
  778. start_time = time.time()
  779. list_rows = merge_line(list_rows, axis=0)
  780. list_cols = merge_line(list_cols, axis=1)
  781. mat_plot(list_rows+list_cols, "merge_line", is_test)
  782. log("merge_line " + str(time.time()-start_time))
  783. list_line = list_rows + list_cols
  784. # 打印处理后线
  785. mat_plot(list_line, "all", is_test)
  786. log("otr postprocess table_line " + str(time.time()-start_time))
  787. return list_line
  788. except Exception as e:
  789. log("table post process failed!")
  790. traceback.print_exc()
  791. return [-1]
  792. def table_line(img, model, size=(512, 1024), prob=0.2, is_test=0):
  793. log("into table_line, prob is " + str(prob))
  794. sizew, sizeh = size
  795. img_new = cv2.resize(img, (sizew, sizeh), interpolation=cv2.INTER_AREA)
  796. start_time = time.time()
  797. pred = model.predict(np.array([img_new]))
  798. log("otr model predict time " + str(time.time()-start_time))
  799. pred = pred[0]
  800. draw_pixel(pred, prob, is_test)
  801. # 横线预测结果
  802. # row_pred = pred[..., 0] > hprob
  803. # row_pred = row_pred.astype(np.uint8)
  804. # # 竖线预测结果
  805. # col_pred = pred[..., 1] > vprob
  806. # col_pred = col_pred.astype(np.uint8)
  807. # # 打印模型输出
  808. # cv2.imshow("predict", (col_pred+row_pred)*255)
  809. # cv2.waitKey(0)
  810. start_time = time.time()
  811. list_line = points2lines(pred, False, prob=prob)
  812. mat_plot(list_line, "points2lines", is_test)
  813. log("points2lines " + str(time.time()-start_time))
  814. # 清除短线
  815. # print(img_new.shape)
  816. start_time = time.time()
  817. list_line = delete_short_lines(list_line, img_new.shape)
  818. mat_plot(list_line, "delete_short_lines", is_test)
  819. log("delete_short_lines " + str(time.time()-start_time))
  820. # # 清除无交点线--无需清除,会影响后面的结果
  821. # start_time = time.time()
  822. # list_line = delete_no_cross_lines(list_line)
  823. # mat_plot(list_line, "delete_no_cross_lines", is_test)
  824. # log("delete_no_cross_lines " + str(time.time()-start_time))
  825. # 分成横竖线
  826. start_time = time.time()
  827. list_rows = []
  828. list_cols = []
  829. for line in list_line:
  830. if line[0] == line[2]:
  831. list_cols.append(line)
  832. elif line[1] == line[3]:
  833. list_rows.append(line)
  834. log("divide rows and cols " + str(time.time()-start_time))
  835. # 合并错开线
  836. start_time = time.time()
  837. list_rows = merge_line(list_rows, axis=0)
  838. list_cols = merge_line(list_cols, axis=1)
  839. mat_plot(list_rows+list_cols, "merge_line", is_test)
  840. log("merge_line " + str(time.time()-start_time))
  841. # 计算交点、分割线
  842. start_time = time.time()
  843. cross_points = get_points(list_rows, list_cols, (img_new.shape[0], img_new.shape[1]))
  844. if not cross_points:
  845. return []
  846. log("get_points " + str(time.time()-start_time))
  847. # 清掉外围的没用的线
  848. # list_rows, list_cols = delete_outline(list_rows, list_cols, cross_points)
  849. # mat_plot(list_rows+list_cols, "delete_outline", is_test)
  850. # 多个表格分割线
  851. start_time = time.time()
  852. list_rows, list_cols = fix_in_split_lines(list_rows, list_cols, img_new)
  853. split_lines, split_y = get_split_line(cross_points, list_cols, img_new)
  854. log("get_split_line " + str(time.time()-start_time))
  855. # 修复边框
  856. start_time = time.time()
  857. new_rows, new_cols, long_rows, long_cols = fix_outline(img_new, list_rows, list_cols, cross_points,
  858. split_y)
  859. # 如有补线
  860. if new_rows or new_cols:
  861. # 连接至补线的延长线
  862. if long_rows:
  863. list_rows = long_rows
  864. if long_cols:
  865. list_cols = long_cols
  866. # 新的补线
  867. if new_rows:
  868. list_rows += new_rows
  869. if new_cols:
  870. list_cols += new_cols
  871. list_rows, list_cols = fix_in_split_lines(list_rows, list_cols, img_new)
  872. # 修复边框后重新计算交点、分割线
  873. cross_points = get_points(list_rows, list_cols, (img_new.shape[0], img_new.shape[1]))
  874. cv_plot(cross_points, img_new.shape, 0, is_test)
  875. split_lines, split_y = get_split_line(cross_points, list_cols, img_new)
  876. print("fix new split_y", split_y)
  877. print("fix new split_lines", split_lines)
  878. # 修复内部缺线
  879. # cross_points = fix_inner(list_rows, list_cols, cross_points, split_y)
  880. # if not cross_points:
  881. # return []
  882. mat_plot(list_rows+list_cols, "fix_outline", is_test)
  883. split_lines_show = []
  884. for _l in split_lines:
  885. split_lines_show.append([_l[0][0], _l[0][1], _l[1][0], _l[1][1]])
  886. mat_plot(split_lines_show+list_cols,
  887. "split_lines", is_test)
  888. log("fix_outline " + str(time.time()-start_time))
  889. # 修复表格4个角
  890. start_time = time.time()
  891. list_rows, list_cols = fix_corner(list_rows, list_cols, split_y, threshold=0)
  892. mat_plot(list_rows+list_cols, "fix_corner", is_test)
  893. log("fix_corner " + str(time.time()-start_time))
  894. # 修复内部缺线
  895. start_time = time.time()
  896. list_rows, list_cols = fix_inner(list_rows, list_cols, cross_points, split_y)
  897. mat_plot(list_rows+list_cols, "fix_inner", is_test)
  898. log("fix_inner " + str(time.time()-start_time))
  899. # 合并错开线
  900. start_time = time.time()
  901. list_rows = merge_line(list_rows, axis=0)
  902. list_cols = merge_line(list_cols, axis=1)
  903. mat_plot(list_rows+list_cols, "merge_line", is_test)
  904. log("merge_line " + str(time.time()-start_time))
  905. list_line = list_rows + list_cols
  906. # 打印处理后线
  907. mat_plot(list_line, "all", is_test)
  908. log("otr postprocess table_line " + str(time.time()-start_time))
  909. return list_line
  910. def table_line2(img, model, size=(512, 1024), hprob=0.5, vprob=0.5, row=50, col=30, alph=15):
  911. sizew, sizeh = size
  912. # [..., ::-1] 最后一维内部反向输出
  913. # inputBlob, fx, fy = letterbox_image(img[..., ::-1], (sizew, sizeh))
  914. # pred = model.predict(np.array([np.array(inputBlob)]))
  915. # pred = model.predict(np.array([np.array(inputBlob)/255.0]))
  916. img_new = cv2.resize(img, (sizew, sizeh), interpolation=cv2.INTER_AREA)
  917. # log("into table_line 1")
  918. pred = model.predict(np.array([img_new]))
  919. # log("into table_line 2")
  920. pred = pred[0]
  921. draw_pixel(pred)
  922. _time = time.time()
  923. points2lines(pred)
  924. log("points2lines takes %ds"%(time.time()-_time))
  925. vpred = pred[..., 1] > vprob # 横线
  926. hpred = pred[..., 0] > hprob # 竖线
  927. vpred = vpred.astype(int)
  928. hpred = hpred.astype(int)
  929. # print("vpred shape", vpred)
  930. # print("hpred shape", hpred)
  931. colboxes = get_table_line(vpred, axis=1, lineW=col)
  932. rowboxes = get_table_line(hpred, axis=0, lineW=row)
  933. # log("into table_line 3")
  934. # if len(rowboxes) > 0:
  935. # rowboxes = np.array(rowboxes)
  936. # rowboxes[:, [0, 2]] = rowboxes[:, [0, 2]]/fx
  937. # rowboxes[:, [1, 3]] = rowboxes[:, [1, 3]]/fy
  938. # rowboxes = rowboxes.tolist()
  939. # if len(colboxes) > 0:
  940. # colboxes = np.array(colboxes)
  941. # colboxes[:, [0, 2]] = colboxes[:, [0, 2]]/fx
  942. # colboxes[:, [1, 3]] = colboxes[:, [1, 3]]/fy
  943. # colboxes = colboxes.tolist()
  944. nrow = len(rowboxes)
  945. ncol = len(colboxes)
  946. for i in range(nrow):
  947. for j in range(ncol):
  948. rowboxes[i] = line_to_line(rowboxes[i], colboxes[j], 10)
  949. colboxes[j] = line_to_line(colboxes[j], rowboxes[i], 10)
  950. # log("into table_line 4")
  951. # 删掉贴着边框的line
  952. temp_list = []
  953. threshold = 5
  954. for line in rowboxes:
  955. if line[1]-0 <= threshold or size[1]-line[1] <= threshold:
  956. continue
  957. # 内部排序
  958. if line[0] > line[2]:
  959. line = [line[2], line[3], line[0], line[1]]
  960. temp_list.append(line)
  961. rowboxes = temp_list
  962. temp_list = []
  963. for line in colboxes:
  964. if line[0]-0 <= threshold or size[0]-line[0] <= threshold:
  965. continue
  966. # 内部排序
  967. if line[1] > line[3]:
  968. line = [line[2], line[3], line[0], line[1]]
  969. temp_list.append(line)
  970. colboxes = temp_list
  971. return rowboxes, colboxes, img_new
  972. def fix_in_split_lines(_rows, _cols, _img):
  973. # 补线贴着边缘无法得到split_y,导致无法分区
  974. for _row in _rows:
  975. if _row[1] >= _img.shape[0] - 5:
  976. _row[1] = _img.shape[0] - 6
  977. _row[3] = _img.shape[0] - 6
  978. print("_row", _row)
  979. if _row[1] <= 0 + 5:
  980. _row[1] = 6
  981. _row[3] = 6
  982. for _col in _cols:
  983. if _col[3] >= _img.shape[0] - 5:
  984. _col[3] = _img.shape[0] - 6
  985. if _col[1] <= 0 + 5:
  986. _col[1] = 6
  987. return _rows, _cols
  988. def mat_plot(list_line, name="", is_test=1):
  989. if not is_test:
  990. return
  991. from matplotlib import pyplot as plt
  992. plt.figure()
  993. plt.title(name)
  994. for _line in list_line:
  995. x0, y0, x1, y1 = _line
  996. plt.plot([x0, x1], [y0, y1])
  997. plt.show()
  998. def cv_plot(_list, img_shape, line_or_point=1, is_test=1):
  999. if is_test == 0:
  1000. return
  1001. img_print = np.zeros(img_shape, np.uint8)
  1002. img_print.fill(255)
  1003. if line_or_point:
  1004. for line in _list:
  1005. cv2.line(img_print, (int(line[0]), int(line[1])), (int(line[2]), int(line[3])),
  1006. (255, 0, 0))
  1007. cv2.imshow("cv_plot", img_print)
  1008. cv2.waitKey(0)
  1009. else:
  1010. for point in _list:
  1011. cv2.circle(img_print, (int(point[0]), int(point[1])), 1, (255, 0, 0), 2)
  1012. cv2.imshow("cv_plot", img_print)
  1013. cv2.waitKey(0)
  1014. def delete_no_cross_lines(list_lines):
  1015. def get_cross_point(l1, l2):
  1016. # https://www.zhihu.com/question/381406535/answer/1095948349
  1017. flag = 0
  1018. # 相交一定是一条横线一条竖线
  1019. if (l1[0] == l1[2] and l2[1] == l2[3]) or (l1[1] == l1[3] and l2[0] == l2[2]):
  1020. if l1[0] <= l2[0] <= l1[2] and l2[1] <= l1[1] <= l2[3]:
  1021. flag = 1
  1022. elif l2[0] <= l1[0] <= l2[2] and l1[1] <= l2[1] <= l1[3]:
  1023. flag = 1
  1024. return flag
  1025. new_list_lines = []
  1026. for i in range(len(list_lines)):
  1027. line1 = list_lines[i]
  1028. find_flag = 0
  1029. for j in range(i+1, len(list_lines)):
  1030. line2 = list_lines[j]
  1031. if get_cross_point(line1, line2):
  1032. # print("delete_no_cross_lines", line1, line2)
  1033. find_flag = 1
  1034. if line2 not in new_list_lines:
  1035. new_list_lines.append(line2)
  1036. if find_flag and line1 not in new_list_lines:
  1037. new_list_lines.append(line1)
  1038. return new_list_lines
  1039. def delete_short_lines(list_lines, image_shape, scale=100):
  1040. x_min_len = max(5, int(image_shape[0] / scale))
  1041. y_min_len = max(5, int(image_shape[1] / scale))
  1042. new_list_lines = []
  1043. for line in list_lines:
  1044. if line[0] == line[2]:
  1045. if abs(line[3] - line[1]) >= y_min_len:
  1046. # print("y_min_len", abs(line[3] - line[1]), y_min_len)
  1047. new_list_lines.append(line)
  1048. else:
  1049. if abs(line[2] - line[0]) >= x_min_len:
  1050. # print("x_min_len", abs(line[2] - line[0]), x_min_len)
  1051. new_list_lines.append(line)
  1052. return new_list_lines
  1053. def get_outline(points, image_np):
  1054. # 取出x, y的最大值最小值
  1055. x_min = points[0][0]
  1056. x_max = points[-1][0]
  1057. points.sort(key=lambda x: (x[1], x[0]))
  1058. y_min = points[0][1]
  1059. y_max = points[-1][1]
  1060. # 创建空图
  1061. # outline_img = np.zeros(image_size, np.uint8)
  1062. outline_img = np.copy(image_np)
  1063. cv2.rectangle(outline_img, (x_min-5, y_min-5), (x_max+5, y_max+5), (0, 0, 0), 2)
  1064. # cv2.imshow("outline_img", outline_img)
  1065. # cv2.waitKey(0)
  1066. return outline_img
  1067. def get_split_line(points, col_lines, image_np, threshold=5):
  1068. # print("get_split_line", image_np.shape)
  1069. points.sort(key=lambda x: (x[1], x[0]))
  1070. # 遍历y坐标,并判断y坐标与上一个y坐标是否存在连接线
  1071. i = 0
  1072. split_line_y = []
  1073. for point in points:
  1074. # 从已分开的线下面开始判断
  1075. if split_line_y:
  1076. if point[1] <= split_line_y[-1] + threshold:
  1077. last_y = point[1]
  1078. continue
  1079. if last_y <= split_line_y[-1] + threshold:
  1080. last_y = point[1]
  1081. continue
  1082. if i == 0:
  1083. last_y = point[1]
  1084. i += 1
  1085. continue
  1086. current_line = (last_y, point[1])
  1087. split_flag = 1
  1088. for col in col_lines:
  1089. # 只要找到一条col包含就不是分割线
  1090. if current_line[0] >= col[1]-3 and current_line[1] <= col[3]+3:
  1091. split_flag = 0
  1092. # print("img", img.shape)
  1093. # print("col", col)
  1094. # print("current_line", current_line)
  1095. break
  1096. if split_flag:
  1097. split_line_y.append(current_line[0]+5)
  1098. split_line_y.append(current_line[1]-5)
  1099. last_y = point[1]
  1100. # 加上收尾分割线
  1101. points.sort(key=lambda x: (x[1], x[0]))
  1102. y_min = points[0][1]
  1103. y_max = points[-1][1]
  1104. # print("加上收尾分割线", y_min, y_max)
  1105. if y_min-threshold < 0:
  1106. split_line_y.append(0)
  1107. else:
  1108. split_line_y.append(y_min-threshold)
  1109. if y_max+threshold > image_np.shape[0]:
  1110. split_line_y.append(image_np.shape[0])
  1111. else:
  1112. split_line_y.append(y_max+threshold)
  1113. split_line_y = list(set(split_line_y))
  1114. # 剔除两条相隔太近分割线
  1115. temp_split_line_y = []
  1116. split_line_y.sort(key=lambda x: x)
  1117. last_y = -20
  1118. for y in split_line_y:
  1119. # print(y)
  1120. if y - last_y >= 20:
  1121. # print(y, last_y)
  1122. temp_split_line_y.append(y)
  1123. last_y = y
  1124. split_line_y = temp_split_line_y
  1125. # print("split_line_y", split_line_y)
  1126. # 生成分割线
  1127. split_line = []
  1128. last_y = 0
  1129. for y in split_line_y:
  1130. # if y - last_y <= 15:
  1131. # continue
  1132. split_line.append([(0, y), (image_np.shape[1], y)])
  1133. last_y = y
  1134. split_line.append([(0, 0), (image_np.shape[1], 0)])
  1135. split_line.append([(0, image_np.shape[0]), (image_np.shape[1], image_np.shape[0])])
  1136. split_line.sort(key=lambda x: x[0][1])
  1137. # print("split_line", split_line)
  1138. # 画图画线
  1139. # split_line_img = np.copy(image_np)
  1140. # for y in split_line_y:
  1141. # cv2.line(split_line_img, (0, y), (image_np.shape[1], y), (0, 0, 0), 1)
  1142. # cv2.imshow("split_line_img", split_line_img)
  1143. # cv2.waitKey(0)
  1144. return split_line, split_line_y
  1145. def get_points(row_lines, col_lines, image_size):
  1146. # 创建空图
  1147. row_img = np.zeros(image_size, np.uint8)
  1148. col_img = np.zeros(image_size, np.uint8)
  1149. # 画线
  1150. thresh = 3
  1151. for row in row_lines:
  1152. cv2.line(row_img, (int(row[0]-thresh), int(row[1])), (int(row[2]+thresh), int(row[3])), (255, 255, 255), 1)
  1153. for col in col_lines:
  1154. cv2.line(col_img, (int(col[0]), int(col[1]-thresh)), (int(col[2]), int(col[3]+thresh)), (255, 255, 255), 1)
  1155. # 求出交点
  1156. point_img = np.bitwise_and(row_img, col_img)
  1157. # cv2.imwrite("get_points.jpg", row_img+col_img)
  1158. # cv2.imshow("get_points", row_img+col_img)
  1159. # cv2.waitKey(0)
  1160. # 识别黑白图中的白色交叉点,将横纵坐标取出
  1161. ys, xs = np.where(point_img > 0)
  1162. points = []
  1163. for i in range(len(xs)):
  1164. points.append((xs[i], ys[i]))
  1165. points.sort(key=lambda x: (x[0], x[1]))
  1166. return points
  1167. def get_minAreaRect(image):
  1168. gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  1169. gray = cv2.bitwise_not(gray)
  1170. thresh = cv2.threshold(gray, 0, 255,
  1171. cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
  1172. coords = np.column_stack(np.where(thresh > 0))
  1173. return cv2.minAreaRect(coords)
  1174. def get_contours(image):
  1175. gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  1176. ret, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
  1177. contours, hierarchy = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  1178. cv2.drawContours(image, contours, -1, (0, 0, 255), 3)
  1179. cv2.imshow("get contours", image)
  1180. cv2.waitKey(0)
  1181. def merge_line(lines, axis, threshold=5):
  1182. """
  1183. 解决模型预测一条直线错开成多条直线,合并成一条直线
  1184. :param lines: 线条列表
  1185. :param axis: 0:横线 1:竖线
  1186. :param threshold: 两条线间像素差阈值
  1187. :return: 合并后的线条列表
  1188. """
  1189. # 任意一条line获取该合并的line,横线往下找,竖线往右找
  1190. lines.sort(key=lambda x: (x[axis], x[1-axis]))
  1191. merged_lines = []
  1192. used_lines = []
  1193. for line1 in lines:
  1194. if line1 in used_lines:
  1195. continue
  1196. merged_line = [line1]
  1197. used_lines.append(line1)
  1198. for line2 in lines:
  1199. if line2 in used_lines:
  1200. continue
  1201. if line1[1-axis]-threshold <= line2[1-axis] <= line1[1-axis]+threshold:
  1202. # 计算基准长度
  1203. min_axis = 10000
  1204. max_axis = 0
  1205. for line3 in merged_line:
  1206. if line3[axis] < min_axis:
  1207. min_axis = line3[axis]
  1208. if line3[axis+2] > max_axis:
  1209. max_axis = line3[axis+2]
  1210. # 判断两条线有无交集
  1211. if min_axis <= line2[axis] <= max_axis \
  1212. or min_axis <= line2[axis+2] <= max_axis:
  1213. merged_line.append(line2)
  1214. used_lines.append(line2)
  1215. if merged_line:
  1216. merged_lines.append(merged_line)
  1217. # 合并line
  1218. result_lines = []
  1219. for merged_line in merged_lines:
  1220. # 获取line宽的平均值
  1221. axis_average = 0
  1222. for line in merged_line:
  1223. axis_average += line[1-axis]
  1224. axis_average = int(axis_average/len(merged_line))
  1225. # 获取最长line两端
  1226. merged_line.sort(key=lambda x: (x[axis]))
  1227. axis_start = merged_line[0][axis]
  1228. merged_line.sort(key=lambda x: (x[axis+2]))
  1229. axis_end = merged_line[-1][axis+2]
  1230. if axis:
  1231. result_lines.append([axis_average, axis_start, axis_average, axis_end])
  1232. else:
  1233. result_lines.append([axis_start, axis_average, axis_end, axis_average])
  1234. return result_lines
  1235. def fix_inner2(row_points, col_points, row_lines, col_lines, threshold=3):
  1236. for i in range(len(row_points)):
  1237. row = row_points[i]
  1238. row.sort(key=lambda x: (x[1], x[0]))
  1239. for j in range(len(row)):
  1240. # 当前点
  1241. point = row[j]
  1242. # 获取当前点在所在行的下个点
  1243. if j >= len(row) - 1:
  1244. next_row_point = []
  1245. else:
  1246. next_row_point = row[j+1]
  1247. if next_row_point:
  1248. for k in range(len(row_lines)):
  1249. line = row_lines[k]
  1250. if line[1] - threshold <= point[1] <= line[1] + threshold:
  1251. if not line[0] <= point[0] <= next_row_point[0] <= line[2]:
  1252. if point[0] <= line[2] < next_row_point[0]:
  1253. if line[2] - point[0] >= 1/3 * (next_row_point[0] - point[0]):
  1254. row_lines[k][2] = next_row_point[0]
  1255. if point[0] < line[0] <= next_row_point[0]:
  1256. if next_row_point[0] - line[0] >= 1/3 * (next_row_point[0] - point[0]):
  1257. row_lines[k][0] = point[0]
  1258. # 获取当前点所在列的下个点
  1259. next_col_point = []
  1260. for col in col_points:
  1261. if point in col:
  1262. col.sort(key=lambda x: (x[0], x[1]))
  1263. if col.index(point) < len(col) - 1:
  1264. next_col_point = col[col.index(point)+1]
  1265. break
  1266. # 获取当前点的对角线点,通过该列下个点所在行的下个点获得
  1267. next_row_next_col_point = []
  1268. if next_col_point:
  1269. for row2 in row_points:
  1270. if next_col_point in row2:
  1271. row2.sort(key=lambda x: (x[1], x[0]))
  1272. if row2.index(next_col_point) < len(row2) - 1:
  1273. next_row_next_col_point = row2[row2.index(next_col_point)+1]
  1274. break
  1275. # 有该列下一点但没有该列下一点所在行的下个点
  1276. if not next_row_next_col_point:
  1277. # 如果有该行下个点
  1278. if next_row_point:
  1279. next_row_next_col_point = [next_row_point[0], next_col_point[1]]
  1280. if next_col_point:
  1281. for k in range(len(col_lines)):
  1282. line = col_lines[k]
  1283. if line[0] - threshold <= point[0] <= line[0] + threshold:
  1284. if not line[1] <= point[1] <= next_col_point[1] <= line[3]:
  1285. if point[1] <= line[3] < next_col_point[1]:
  1286. if line[3] - point[1] >= 1/3 * (next_col_point[1] - point[1]):
  1287. col_lines[k][3] = next_col_point[1]
  1288. if point[1] < line[1] <= next_col_point[1]:
  1289. if next_col_point[1] - line[1] >= 1/3 * (next_col_point[1] - point[1]):
  1290. col_lines[k][1] = point[1]
  1291. if next_row_next_col_point:
  1292. for k in range(len(col_lines)):
  1293. line = col_lines[k]
  1294. if line[0] - threshold <= next_row_next_col_point[0] <= line[0] + threshold:
  1295. if not line[1] <= point[1] <= next_row_next_col_point[1] <= line[3]:
  1296. if point[1] < line[1] <= next_row_next_col_point[1]:
  1297. if next_row_next_col_point[1] - line[1] >= 1/3 * (next_row_next_col_point[1] - point[1]):
  1298. col_lines[k][1] = point[1]
  1299. return row_lines, col_lines
  1300. def fix_inner1(row_lines, col_lines, points, split_y):
  1301. def fix(fix_lines, assist_lines, split_points, axis):
  1302. new_points = []
  1303. for line1 in fix_lines:
  1304. min_assist_line = [[], []]
  1305. min_distance = [1000, 1000]
  1306. if_find = [0, 0]
  1307. # 获取fix_line中的所有col point,里面可能不包括两个顶点,col point是交点,顶点可能不是交点
  1308. fix_line_points = []
  1309. for point in split_points:
  1310. if abs(point[1-axis] - line1[1-axis]) <= 2:
  1311. if line1[axis] <= point[axis] <= line1[axis+2]:
  1312. fix_line_points.append(point)
  1313. # 找出离两个顶点最近的assist_line, 并且assist_line与fix_line不相交
  1314. line1_point = [line1[:2], line1[2:]]
  1315. for i in range(2):
  1316. point = line1_point[i]
  1317. for line2 in assist_lines:
  1318. if not if_find[i] and abs(point[axis] - line2[axis]) <= 2:
  1319. if line1[1-axis] <= point[1-axis] <= line2[1-axis+2]:
  1320. # print("line1, match line2", line1, line2)
  1321. if_find[i] = 1
  1322. break
  1323. else:
  1324. if abs(point[axis] - line2[axis]) < min_distance[i] and line2[1-axis] <= point[1-axis] <= line2[1-axis+2]:
  1325. if line1[axis] <= line2[axis] <= line1[axis+2]:
  1326. continue
  1327. min_distance[i] = abs(line1[axis] - line2[axis])
  1328. min_assist_line[i] = line2
  1329. # 找出离assist_line最近的交点
  1330. # 顶点到交点的距离(多出来的线)需大于assist_line到交点的距离(bbox的边)的1/3
  1331. min_distance = [1000, 1000]
  1332. min_col_point = [[], []]
  1333. for i in range(2):
  1334. # print("顶点", i, line1_point[i])
  1335. if min_assist_line[i]:
  1336. for point in fix_line_points:
  1337. if abs(point[axis] - min_assist_line[i][axis]) < min_distance[i]:
  1338. min_distance[i] = abs(point[axis] - min_assist_line[i][axis])
  1339. min_col_point[i] = point
  1340. if min_col_point[i]:
  1341. bbox_len = abs(min_col_point[i][axis] - min_assist_line[i][axis])
  1342. line_distance = abs(min_col_point[i][axis] - line1_point[i][axis])
  1343. if bbox_len/3 <= line_distance <= bbox_len:
  1344. add_point = (line1_point[i][1-axis], min_assist_line[i][axis])
  1345. print("============================table line==")
  1346. print("fix_inner add point", add_point)
  1347. print(min_col_point[i][axis], line1_point[i][axis], min_col_point[i][axis], min_assist_line[i][axis])
  1348. print(abs(min_col_point[i][axis] - line1_point[i][axis]), abs(min_col_point[i][axis] - min_assist_line[i][axis])/3)
  1349. print("line1, line2", line1, min_assist_line[i])
  1350. new_points.append(add_point)
  1351. return new_points
  1352. new_points = []
  1353. for i in range(1, len(split_y)):
  1354. last_y = split_y[i-1]
  1355. y = split_y[i]
  1356. # 先对点线进行分区
  1357. split_row_lines = []
  1358. split_col_lines = []
  1359. split_points = []
  1360. for row in row_lines:
  1361. if last_y <= row[1] <= y:
  1362. split_row_lines.append(row)
  1363. for col in col_lines:
  1364. if last_y <= col[1] <= y:
  1365. split_col_lines.append(col)
  1366. for point in points:
  1367. if last_y <= point[1] <= y:
  1368. split_points.append(point)
  1369. new_points += fix(split_col_lines, split_row_lines, split_points, axis=1)
  1370. new_points += fix(split_row_lines, split_col_lines, split_points, axis=0)
  1371. # 找出所有col的顶点不在row上的、row的顶点不在col上的
  1372. # for col in split_col_lines:
  1373. # print("*"*30)
  1374. #
  1375. # # 获取该line中的所有point
  1376. # col_points = []
  1377. # for point in split_points:
  1378. # if abs(point[0] - col[0]) <= 2:
  1379. # if col[1] <= point[1] <= col[3]:
  1380. # col_points.append(point)
  1381. #
  1382. # # 比较顶点
  1383. # min_row_1 = []
  1384. # min_row_2 = []
  1385. # min_distance_1 = 1000
  1386. # min_distance_2 = 1000
  1387. # if_find_1 = 0
  1388. # if_find_2 = 0
  1389. # for row in split_row_lines:
  1390. # # 第一个顶点
  1391. # if not if_find_1 and abs(col[1] - row[1]) <= 2:
  1392. # if row[0] <= col[0] <= row[2]:
  1393. # print("col, match row", col, row)
  1394. # if_find_1 = 1
  1395. # break
  1396. # else:
  1397. # if abs(col[1] - row[1]) < min_distance_1 and row[0] <= col[0] <= row[2]:
  1398. # if col[1] <= row[1] <= col[3]:
  1399. # continue
  1400. # min_distance_1 = abs(col[1] - row[1])
  1401. # min_row_1 = row
  1402. #
  1403. # # 第二个顶点
  1404. # if not if_find_2 and abs(col[3] - row[1]) <= 2:
  1405. # if row[0] <= col[2] <= row[2]:
  1406. # if_find_2 = 1
  1407. # break
  1408. # else:
  1409. # if abs(col[3] - row[1]) < min_distance_2 and row[0] <= col[2] <= row[2]:
  1410. # min_distance_2 = abs(col[3] - row[1])
  1411. # min_row_2 = row
  1412. #
  1413. # if not if_find_1:
  1414. # print("col", col)
  1415. # print("min_row_1", min_row_1)
  1416. # if min_row_1:
  1417. # min_distance_1 = 1000
  1418. # min_col_point = []
  1419. # for point in col_points:
  1420. # if abs(point[1] - min_row_1[1]) < min_distance_1:
  1421. # min_distance_1 = abs(point[1] - min_row_1[1])
  1422. # min_col_point = point
  1423. #
  1424. # if abs(min_col_point[1] - col[1]) >= abs(min_col_point[1] - min_row_1[1])/3:
  1425. #
  1426. # add_point = (col[0], min_row_1[1])
  1427. # print("fix_inner add point", add_point)
  1428. # new_points.append(add_point)
  1429. # else:
  1430. # print("distance too long", min_col_point, min_row_1)
  1431. # print(abs(min_col_point[1] - col[1]), abs(min_col_point[1] - min_row_1[1])/3)
  1432. return points+new_points
  1433. def fix_inner(row_lines, col_lines, points, split_y):
  1434. def fix(fix_lines, assist_lines, split_points, axis):
  1435. new_points = []
  1436. for line1 in fix_lines:
  1437. min_assist_line = [[], []]
  1438. min_distance = [1000, 1000]
  1439. if_find = [0, 0]
  1440. # 获取fix_line中的所有col point,里面可能不包括两个顶点,col point是交点,顶点可能不是交点
  1441. fix_line_points = []
  1442. for point in split_points:
  1443. if abs(point[1-axis] - line1[1-axis]) <= 2:
  1444. if line1[axis] <= point[axis] <= line1[axis+2]:
  1445. fix_line_points.append(point)
  1446. # 找出离两个顶点最近的assist_line, 并且assist_line与fix_line不相交
  1447. line1_point = [line1[:2], line1[2:]]
  1448. for i in range(2):
  1449. point = line1_point[i]
  1450. for line2 in assist_lines:
  1451. if not if_find[i] and abs(point[axis] - line2[axis]) <= 2:
  1452. if line1[1-axis] <= point[1-axis] <= line2[1-axis+2]:
  1453. # print("line1, match line2", line1, line2)
  1454. if_find[i] = 1
  1455. break
  1456. else:
  1457. if abs(point[axis] - line2[axis]) < min_distance[i] and line2[1-axis] <= point[1-axis] <= line2[1-axis+2]:
  1458. if line1[axis] <= line2[axis] <= line1[axis+2]:
  1459. continue
  1460. min_distance[i] = abs(line1[axis] - line2[axis])
  1461. min_assist_line[i] = line2
  1462. # 找出离assist_line最近的交点
  1463. # 顶点到交点的距离(多出来的线)需大于assist_line到交点的距离(bbox的边)的1/3
  1464. min_distance = [1000, 1000]
  1465. min_col_point = [[], []]
  1466. for i in range(2):
  1467. # print("顶点", i, line1_point[i])
  1468. if min_assist_line[i]:
  1469. for point in fix_line_points:
  1470. if abs(point[axis] - min_assist_line[i][axis]) < min_distance[i]:
  1471. min_distance[i] = abs(point[axis] - min_assist_line[i][axis])
  1472. min_col_point[i] = point
  1473. # print("min_col_point", min_col_point)
  1474. # print("min_assist_line", min_assist_line)
  1475. # print("line1_point", line1_point)
  1476. if min_assist_line[0] and min_assist_line[0] == min_assist_line[1]:
  1477. if min_assist_line[0][axis] < line1_point[0][axis]:
  1478. bbox_len = abs(min_col_point[0][axis] - min_assist_line[0][axis])
  1479. line_distance = abs(min_col_point[0][axis] - line1_point[0][axis])
  1480. if bbox_len/3 <= line_distance <= bbox_len:
  1481. if axis == 1:
  1482. add_point = (line1_point[0][1-axis], min_assist_line[0][axis])
  1483. else:
  1484. add_point = (min_assist_line[0][axis], line1_point[0][1-axis])
  1485. new_points.append([line1, add_point])
  1486. elif min_assist_line[1][axis] > line1_point[1][axis]:
  1487. bbox_len = abs(min_col_point[1][axis] - min_assist_line[1][axis])
  1488. line_distance = abs(min_col_point[1][axis] - line1_point[1][axis])
  1489. if bbox_len/3 <= line_distance <= bbox_len:
  1490. if axis == 1:
  1491. add_point = (line1_point[1][1-axis], min_assist_line[1][axis])
  1492. else:
  1493. add_point = (min_assist_line[1][axis], line1_point[1][1-axis])
  1494. new_points.append([line1, add_point])
  1495. else:
  1496. for i in range(2):
  1497. if min_col_point[i]:
  1498. bbox_len = abs(min_col_point[i][axis] - min_assist_line[i][axis])
  1499. line_distance = abs(min_col_point[i][axis] - line1_point[i][axis])
  1500. # print("bbox_len, line_distance", bbox_len, line_distance)
  1501. if bbox_len/3 <= line_distance <= bbox_len:
  1502. if axis == 1:
  1503. add_point = (line1_point[i][1-axis], min_assist_line[i][axis])
  1504. else:
  1505. add_point = (min_assist_line[i][axis], line1_point[i][1-axis])
  1506. # print("============================table line==")
  1507. # print("fix_inner add point", add_point)
  1508. # print(min_col_point[i][axis], line1_point[i][axis], min_col_point[i][axis], min_assist_line[i][axis])
  1509. # print(abs(min_col_point[i][axis] - line1_point[i][axis]), abs(min_col_point[i][axis] - min_assist_line[i][axis])/3)
  1510. # print("line1, line2", line1, min_assist_line[i])
  1511. # print("line1, add_point", [line1, add_point])
  1512. new_points.append([line1, add_point])
  1513. return new_points
  1514. row_lines_copy = copy.deepcopy(row_lines)
  1515. col_lines_copy = copy.deepcopy(col_lines)
  1516. try:
  1517. new_points = []
  1518. for i in range(1, len(split_y)):
  1519. last_y = split_y[i-1]
  1520. y = split_y[i]
  1521. # 先对点线进行分区
  1522. split_row_lines = []
  1523. split_col_lines = []
  1524. split_points = []
  1525. for row in row_lines:
  1526. if last_y <= row[1] <= y:
  1527. split_row_lines.append(row)
  1528. for col in col_lines:
  1529. if last_y <= col[1] <= y:
  1530. split_col_lines.append(col)
  1531. for point in points:
  1532. if last_y <= point[1] <= y:
  1533. split_points.append(point)
  1534. new_point_list = fix(split_col_lines, split_row_lines, split_points, axis=1)
  1535. for line, new_point in new_point_list:
  1536. if line in col_lines:
  1537. index = col_lines.index(line)
  1538. point1 = line[:2]
  1539. point2 = line[2:]
  1540. if new_point[1] >= point2[1]:
  1541. col_lines[index] = [point1[0], point1[1], new_point[0], new_point[1]]
  1542. elif new_point[1] <= point1[1]:
  1543. col_lines[index] = [new_point[0], new_point[1], point2[0], point2[1]]
  1544. new_point_list = fix(split_row_lines, split_col_lines, split_points, axis=0)
  1545. for line, new_point in new_point_list:
  1546. if line in row_lines:
  1547. index = row_lines.index(line)
  1548. point1 = line[:2]
  1549. point2 = line[2:]
  1550. if new_point[0] >= point2[0]:
  1551. row_lines[index] = [point1[0], point1[1], new_point[0], new_point[1]]
  1552. elif new_point[0] <= point1[0]:
  1553. row_lines[index] = [new_point[0], new_point[1], point2[0], point2[1]]
  1554. return row_lines, col_lines
  1555. except:
  1556. traceback.print_exc()
  1557. return row_lines_copy, col_lines_copy
  1558. def fix_corner(row_lines, col_lines, split_y, threshold=0):
  1559. new_row_lines = []
  1560. new_col_lines = []
  1561. last_y = split_y[0]
  1562. for y in split_y:
  1563. if y == last_y:
  1564. continue
  1565. split_row_lines = []
  1566. split_col_lines = []
  1567. for row in row_lines:
  1568. if last_y-threshold <= row[1] <= y+threshold or last_y-threshold <= row[3] <= y+threshold:
  1569. split_row_lines.append(row)
  1570. for col in col_lines:
  1571. # fix corner 容易因split line 漏掉线
  1572. if last_y-threshold <= col[1] <= y+threshold or last_y-threshold <= col[3] <= y+threshold:
  1573. split_col_lines.append(col)
  1574. if not split_row_lines or not split_col_lines:
  1575. last_y = y
  1576. continue
  1577. split_row_lines.sort(key=lambda x: (x[1], x[0]))
  1578. split_col_lines.sort(key=lambda x: (x[0], x[1]))
  1579. up_line = split_row_lines[0]
  1580. bottom_line = split_row_lines[-1]
  1581. left_line = split_col_lines[0]
  1582. right_line = split_col_lines[-1]
  1583. # 左上角
  1584. if up_line[0:2] != left_line[0:2]:
  1585. # print("up_line, left_line", up_line, left_line)
  1586. add_corner = [left_line[0], up_line[1]]
  1587. split_row_lines[0][0] = add_corner[0]
  1588. split_col_lines[0][1] = add_corner[1]
  1589. # 右上角
  1590. if up_line[2:] != right_line[:2]:
  1591. # print("up_line, right_line", up_line, right_line)
  1592. add_corner = [right_line[0], up_line[1]]
  1593. split_row_lines[0][2] = add_corner[0]
  1594. split_col_lines[-1][1] = add_corner[1]
  1595. new_row_lines = new_row_lines + split_row_lines
  1596. new_col_lines = new_col_lines + split_col_lines
  1597. last_y = y
  1598. return new_row_lines, new_col_lines
  1599. def delete_outline(row_lines, col_lines, points):
  1600. row_lines.sort(key=lambda x: (x[1], x[0]))
  1601. col_lines.sort(key=lambda x: (x[0], x[1]))
  1602. line = [row_lines[0], row_lines[-1], col_lines[0], col_lines[-1]]
  1603. threshold = 2
  1604. point_cnt = [0, 0, 0, 0]
  1605. for point in points:
  1606. for i in range(4):
  1607. if i < 2:
  1608. if line[i][1]-threshold <= point[1] <= line[i][1]+threshold:
  1609. if line[i][0] <= point[0] <= line[i][2]:
  1610. point_cnt[i] += 1
  1611. else:
  1612. if line[i][0]-threshold <= point[0] <= line[i][0]+threshold:
  1613. if line[i][1] <= point[1] <= line[i][3]:
  1614. point_cnt[i] += 1
  1615. # if line[0][1]-threshold <= point[1] <= line[0][1]+threshold:
  1616. # if line[0][0] <= point[0] <= line[0][2]:
  1617. # point_cnt[0] += 1
  1618. # elif line[1][1]-threshold <= point[1] <= line[1][1]+threshold:
  1619. # if line[1][0] <= point[0] <= line[1][2]:
  1620. # point_cnt[1] += 1
  1621. # elif line[2][0]-threshold <= point[0] <= line[2][0]+threshold:
  1622. # if line[2][1] <= point[1] <= line[2][3]:
  1623. # point_cnt[2] += 1
  1624. # elif line[3][0]-threshold <= point[0] <= line[3][0]+threshold:
  1625. # if line[3][1] <= point[1] <= line[3][3]:
  1626. # point_cnt[3] += 1
  1627. # 轮廓line至少包含3个交点
  1628. for i in range(4):
  1629. if point_cnt[i] < 3:
  1630. if i < 2:
  1631. if line[i] in row_lines:
  1632. row_lines.remove(line[i])
  1633. else:
  1634. if line[i] in col_lines:
  1635. col_lines.remove(line[i])
  1636. return row_lines, col_lines
  1637. def fix_outline2(image, row_lines, col_lines, points, split_y):
  1638. print("split_y", split_y)
  1639. # 分割线纵坐标
  1640. if len(split_y) < 2:
  1641. return [], [], [], []
  1642. # elif len(split_y) == 2:
  1643. # split_y = [2000., 2000., 2000., 2000.]
  1644. split_y.sort(key=lambda x: x)
  1645. new_split_y = []
  1646. for i in range(1, len(split_y), 2):
  1647. new_split_y.append(int((split_y[i]+split_y[i-1])/2))
  1648. # # 查看是否正确输出区域分割线
  1649. # for line in split_y:
  1650. # cv2.line(image, (0, int(line)), (int(image.shape[1]), int(line)), (0, 0, 255), 2)
  1651. # cv2.imshow("split_y", image)
  1652. # cv2.waitKey(0)
  1653. # 预测线根据分割线纵坐标分为多个分割区域
  1654. # row_lines.sort(key=lambda x: (x[3], x[2], x[1], x[0]))
  1655. # col_lines.sort(key=lambda x: (x[3], x[2], x[1], x[0]))
  1656. # points.sort(key=lambda x: (x[1], x[0]))
  1657. # row_count = 0
  1658. # col_count = 0
  1659. # point_count = 0
  1660. split_row_list = []
  1661. split_col_list = []
  1662. split_point_list = []
  1663. # for i in range(1, len(split_y)):
  1664. # y = split_y[i]
  1665. # last_y = split_y[i-1]
  1666. # row_lines = row_lines[row_count:]
  1667. # col_lines = col_lines[col_count:]
  1668. # points = points[point_count:]
  1669. # row_count = 0
  1670. # col_count = 0
  1671. # point_count = 0
  1672. #
  1673. # if not row_lines:
  1674. # split_row_list.append([])
  1675. # for row in row_lines:
  1676. # if last_y <= row[3] <= y:
  1677. # row_count += 1
  1678. # else:
  1679. # split_row_list.append(row_lines[:row_count])
  1680. # break
  1681. # if row_count == len(row_lines):
  1682. # split_row_list.append(row_lines[:row_count])
  1683. # break
  1684. #
  1685. # if not col_lines:
  1686. # split_col_list.append([])
  1687. #
  1688. # for col in col_lines:
  1689. # # if last_y <= col[3] <= y:
  1690. # if col[1] <= last_y <= y <= col[3] or last_y <= col[3] <= y:
  1691. # # if last_y <= col[1] <= y or last_y <= col[3] <= y:
  1692. # col_count += 1
  1693. # else:
  1694. # split_col_list.append(col_lines[:col_count])
  1695. # break
  1696. # if col_count == len(col_lines):
  1697. # split_col_list.append(col_lines[:col_count])
  1698. # break
  1699. #
  1700. # if not points:
  1701. # split_point_list.append([])
  1702. # for point in points:
  1703. # if last_y <= point[1] <= y:
  1704. # point_count += 1
  1705. # else:
  1706. # split_point_list.append(points[:point_count])
  1707. # break
  1708. # if point_count == len(points):
  1709. # split_point_list.append(points[:point_count])
  1710. # break
  1711. #
  1712. # # print("len(split_row_list)", len(split_row_list))
  1713. # # print("len(split_col_list)", len(split_col_list))
  1714. # if row_count < len(row_lines) - 1 and col_count < len(col_lines) - 1:
  1715. # row_lines = row_lines[row_count:]
  1716. # split_row_list.append(row_lines)
  1717. # col_lines = col_lines[col_count:]
  1718. # split_col_list.append(col_lines)
  1719. #
  1720. # if point_count < len(points) - 1:
  1721. # points = points[point_count:len(points)]
  1722. # split_point_list.append(points)
  1723. for i in range(1, len(split_y)):
  1724. y = split_y[i]
  1725. last_y = split_y[i-1]
  1726. split_row = []
  1727. for row in row_lines:
  1728. if last_y <= row[3] <= y:
  1729. split_row.append(row)
  1730. split_row_list.append(split_row)
  1731. split_col = []
  1732. for col in col_lines:
  1733. if last_y <= col[1] <= y or last_y <= col[3] <= y or col[1] < last_y < y < col[3]:
  1734. split_col.append(col)
  1735. split_col_list.append(split_col)
  1736. split_point = []
  1737. for point in points:
  1738. if last_y <= point[1] <= y:
  1739. split_point.append(point)
  1740. split_point_list.append(split_point)
  1741. # 预测线取上下左右4个边(会有超出表格部分) [(), ()]
  1742. area_row_line = []
  1743. area_col_line = []
  1744. for area in split_row_list:
  1745. if not area:
  1746. area_row_line.append([])
  1747. continue
  1748. area.sort(key=lambda x: (x[1], x[0]))
  1749. up_line = area[0]
  1750. bottom_line = area[-1]
  1751. area_row_line.append([up_line, bottom_line])
  1752. for area in split_col_list:
  1753. if not area:
  1754. area_col_line.append([])
  1755. continue
  1756. area.sort(key=lambda x: x[0])
  1757. left_line = area[0]
  1758. right_line = area[-1]
  1759. area_col_line.append([left_line, right_line])
  1760. # 线交点根据分割线纵坐标分为多个分割区域
  1761. # points.sort(key=lambda x: (x[1], x[0]))
  1762. # point_count = 0
  1763. # split_point_list = []
  1764. # for y in new_split_y:
  1765. # points = points[point_count:len(points)]
  1766. # point_count = 0
  1767. # for point in points:
  1768. # if point[1] <= y:
  1769. # point_count += 1
  1770. # else:
  1771. # split_point_list.append(points[:point_count])
  1772. # break
  1773. # if point_count == len(points):
  1774. # split_point_list.append(points[:point_count])
  1775. # break
  1776. # if point_count < len(points) - 1:
  1777. # points = points[point_count:len(points)]
  1778. # split_point_list.append(points)
  1779. # print("len(split_point_list)", len(split_point_list))
  1780. # 取每个分割区域的4条线(无超出表格部分)
  1781. area_row_line2 = []
  1782. area_col_line2 = []
  1783. for area in split_point_list:
  1784. if not area:
  1785. area_row_line2.append([])
  1786. area_col_line2.append([])
  1787. continue
  1788. area.sort(key=lambda x: (x[0], x[1]))
  1789. left_up = area[0]
  1790. right_bottom = area[-1]
  1791. up_line = [left_up[0], left_up[1], right_bottom[0], left_up[1]]
  1792. bottom_line = [left_up[0], right_bottom[1], right_bottom[0], right_bottom[1]]
  1793. left_line = [left_up[0], left_up[1], left_up[0], right_bottom[1]]
  1794. right_line = [right_bottom[0], left_up[1], right_bottom[0], right_bottom[1]]
  1795. area_row_line2.append([up_line, bottom_line])
  1796. area_col_line2.append([left_line, right_line])
  1797. # 判断超出部分的长度,超出一定长度就补线
  1798. new_row_lines = []
  1799. new_col_lines = []
  1800. longer_row_lines = []
  1801. longer_col_lines = []
  1802. all_longer_row_lines = []
  1803. all_longer_col_lines = []
  1804. # print("split_y", split_y)
  1805. # print("split_row_list", split_row_list, len(split_row_list))
  1806. # print("split_row_list", split_col_list, len(split_col_list))
  1807. # print("area_row_line", area_row_line, len(area_row_line))
  1808. # print("area_col_line", area_col_line, len(area_col_line))
  1809. for i in range(len(area_row_line)):
  1810. if not area_row_line[i] or not area_col_line[i]:
  1811. continue
  1812. up_line = area_row_line[i][0]
  1813. up_line2 = area_row_line2[i][0]
  1814. bottom_line = area_row_line[i][1]
  1815. bottom_line2 = area_row_line2[i][1]
  1816. left_line = area_col_line[i][0]
  1817. left_line2 = area_col_line2[i][0]
  1818. right_line = area_col_line[i][1]
  1819. right_line2 = area_col_line2[i][1]
  1820. # 计算单格高度宽度
  1821. if len(split_row_list[i]) > 1:
  1822. height_dict = {}
  1823. for j in range(len(split_row_list[i])):
  1824. if j + 1 > len(split_row_list[i]) - 1:
  1825. break
  1826. height = abs(int(split_row_list[i][j][3] - split_row_list[i][j+1][3]))
  1827. if height in height_dict.keys():
  1828. height_dict[height] = height_dict[height] + 1
  1829. else:
  1830. height_dict[height] = 1
  1831. height_list = [[x, height_dict[x]] for x in height_dict.keys()]
  1832. height_list.sort(key=lambda x: (x[1], -x[0]), reverse=True)
  1833. # print("height_list", height_list)
  1834. box_height = height_list[0][0]
  1835. else:
  1836. box_height = 10
  1837. if len(split_col_list[i]) > 1:
  1838. box_width = abs(split_col_list[i][1][2] - split_col_list[i][0][2])
  1839. else:
  1840. box_width = 10
  1841. print("box_height", box_height, "box_width", box_width)
  1842. # cv2.line(image, (int(up_line[0]), int(up_line[1])),
  1843. # (int(up_line[2]), int(up_line[3])),
  1844. # (255, 255, 0), 2)
  1845. # cv2.line(image, (int(right_line[0]), int(right_line[1])),
  1846. # (int(right_line[2]), int(right_line[3])),
  1847. # (0, 255, 255), 2)
  1848. # cv2.imshow("right_line", image)
  1849. # cv2.waitKey(0)
  1850. # 补左右两条竖线超出来的线的row
  1851. if (up_line[1] - left_line[1] >= 10 and up_line[1] - right_line[1] >= 2) or \
  1852. (up_line[1] - left_line[1] >= 2 and up_line[1] - right_line[1] >= 10):
  1853. if up_line[1] - left_line[1] >= up_line[1] - right_line[1]:
  1854. new_row_lines.append([left_line[0], left_line[1], right_line[0], left_line[1]])
  1855. new_col_y = left_line[1]
  1856. # 补了row,要将其他短的col连到row上
  1857. for j in range(len(split_col_list[i])):
  1858. col = split_col_list[i][j]
  1859. # 且距离不能相差大于一格
  1860. # print("abs(new_col_y - col[1])", abs(new_col_y - col[1]))
  1861. if abs(new_col_y - col[1]) <= box_height:
  1862. split_col_list[i][j][1] = min([new_col_y, col[1]])
  1863. longer_col_lines.append([col[0], min([new_col_y, col[1]]), col[2], col[3]])
  1864. else:
  1865. new_row_lines.append([left_line[0], right_line[1], right_line[0], right_line[1]])
  1866. new_col_y = right_line[1]
  1867. # 补了row,要将其他短的col连到row上
  1868. for j in range(len(split_col_list[i])):
  1869. # 需判断该线在这个区域中
  1870. # if up_line2[1]-3 <= col[1] <= col[3] <= bottom_line2[1]+3:
  1871. col = split_col_list[i][j]
  1872. # 且距离不能相差太大
  1873. # print("abs(new_col_y - col[1])", abs(new_col_y - col[1]))
  1874. if abs(new_col_y - col[1]) <= box_height:
  1875. split_col_list[i][j][1] = min([new_col_y, col[1]])
  1876. if (left_line[3] - bottom_line[3] >= 10 and right_line[3] - bottom_line[3] >= 2) or \
  1877. (left_line[3] - bottom_line[3] >= 2 and right_line[3] - bottom_line[3] >= 10):
  1878. if left_line[3] - bottom_line[3] >= right_line[3] - bottom_line[3]:
  1879. new_row_lines.append([left_line[2], left_line[3], right_line[2], left_line[3]])
  1880. new_col_y = left_line[3]
  1881. # 补了row,要将其他短的col连到row上
  1882. for j in range(len(split_col_list[i])):
  1883. col = split_col_list[i][j]
  1884. # 且距离不能相差太大
  1885. if abs(new_col_y - col[3]) <= box_height:
  1886. split_col_list[i][j][3] = max([new_col_y, col[3]])
  1887. else:
  1888. new_row_lines.append([left_line[2], right_line[3], right_line[2], right_line[3]])
  1889. new_col_y = right_line[3]
  1890. # 补了row,要将其他短的col连到row上
  1891. for j in range(len(split_col_list[i])):
  1892. col = split_col_list[i][j]
  1893. # 且距离不能相差太大
  1894. if abs(new_col_y - col[3]) <= box_height:
  1895. split_col_list[i][j][3] = max([new_col_y, col[3]])
  1896. # 补上下两条横线超出来的线的col
  1897. if (left_line[0] - up_line[0] >= 10 and left_line[0] - bottom_line[0] >= 2) or \
  1898. (left_line[0] - up_line[0] >= 2 and left_line[0] - bottom_line[0] >= 10):
  1899. if left_line[0] - up_line[0] >= left_line[0] - bottom_line[0]:
  1900. new_col_lines.append([up_line[0], up_line[1], up_line[0], bottom_line[1]])
  1901. new_row_x = up_line[0]
  1902. # 补了col,要将其他短的row连到col上
  1903. for j in range(len(split_row_list[i])):
  1904. row = split_row_list[i][j]
  1905. # 且距离不能相差太大
  1906. if abs(new_row_x - row[0]) <= box_width:
  1907. split_row_list[i][j][0] = min([new_row_x, row[0]])
  1908. else:
  1909. new_col_lines.append([bottom_line[0], up_line[1], bottom_line[0], bottom_line[1]])
  1910. new_row_x = bottom_line[0]
  1911. # 补了col,要将其他短的row连到col上
  1912. for j in range(len(split_row_list[i])):
  1913. row = split_row_list[i][j]
  1914. # 且距离不能相差太大
  1915. if abs(new_row_x - row[0]) <= box_width:
  1916. split_row_list[i][j][0] = min([new_row_x, row[0]])
  1917. if (up_line[2] - right_line[2] >= 10 and bottom_line[2] - right_line[2] >= 2) or \
  1918. (up_line[2] - right_line[2] >= 2 and bottom_line[2] - right_line[2] >= 10):
  1919. if up_line[2] - right_line[2] >= bottom_line[2] - right_line[2]:
  1920. new_col_lines.append([up_line[2], up_line[3], up_line[2], bottom_line[3]])
  1921. new_row_x = up_line[2]
  1922. # 补了col,要将其他短的row连到col上
  1923. for j in range(len(split_row_list[i])):
  1924. row = split_row_list[i][j]
  1925. # 且距离不能相差太大
  1926. if abs(new_row_x - row[2]) <= box_width:
  1927. split_row_list[i][j][2] = max([new_row_x, row[2]])
  1928. else:
  1929. new_col_lines.append([bottom_line[2], up_line[3], bottom_line[2], bottom_line[3]])
  1930. new_row_x = bottom_line[2]
  1931. # 补了col,要将其他短的row连到col上
  1932. for j in range(len(split_row_list[i])):
  1933. # 需判断该线在这个区域中
  1934. # if up_line2[1]-3 <= row[1] <= bottom_line2[1]+3:
  1935. row = split_row_list[i][j]
  1936. # 且距离不能相差太大
  1937. if abs(new_row_x - row[2]) <= box_width:
  1938. split_row_list[i][j][2] = max([new_row_x, row[2]])
  1939. all_longer_row_lines += split_row_list[i]
  1940. all_longer_col_lines += split_col_list[i]
  1941. # print("all_longer_row_lines", len(all_longer_row_lines), i)
  1942. # print("all_longer_col_lines", len(all_longer_col_lines), i)
  1943. # print("new_row_lines", len(new_row_lines), i)
  1944. # print("new_col_lines", len(new_col_lines), i)
  1945. # 删除表格内部的补线
  1946. # temp_list = []
  1947. # for row in new_row_lines:
  1948. # if up_line[1]-5 <= row[1] <= bottom_line[1]+5:
  1949. # continue
  1950. # temp_list.append(row)
  1951. # print("fix_outline", new_row_lines)
  1952. # new_row_lines = temp_list
  1953. # print("fix_outline", new_row_lines)
  1954. # temp_list = []
  1955. # for col in new_col_lines:
  1956. # if left_line[0]-5 <= col[0] <= right_line[0]+5:
  1957. # continue
  1958. # temp_list.append(col)
  1959. #
  1960. # new_col_lines = temp_list
  1961. # print("fix_outline", new_col_lines)
  1962. # print("fix_outline", new_row_lines)
  1963. # 删除重复包含的补线
  1964. # temp_list = []
  1965. # for row in new_row_lines:
  1966. # if up_line[1]-5 <= row[1] <= bottom_line[1]+5:
  1967. # continue
  1968. # temp_list.append(row)
  1969. # new_row_lines = temp_list
  1970. # 展示上下左右边框线
  1971. # for i in range(len(area_row_line)):
  1972. # print("row1", area_row_line[i])
  1973. # print("row2", area_row_line2[i])
  1974. # print("col1", area_col_line[i])
  1975. # print("col2", area_col_line2[i])
  1976. # cv2.line(image, (int(area_row_line[i][0][0]), int(area_row_line[i][0][1])),
  1977. # (int(area_row_line[i][0][2]), int(area_row_line[i][0][3])), (0, 255, 0), 2)
  1978. # cv2.line(image, (int(area_row_line2[i][1][0]), int(area_row_line2[i][1][1])),
  1979. # (int(area_row_line2[i][1][2]), int(area_row_line2[i][1][3])), (0, 0, 255), 2)
  1980. # cv2.imshow("fix_outline", image)
  1981. # cv2.waitKey(0)
  1982. # 展示所有线
  1983. # for line in all_longer_col_lines:
  1984. # cv2.line(image, (int(line[0]), int(line[1])),
  1985. # (int(line[2]), int(line[3])),
  1986. # (0, 255, 0), 2)
  1987. # cv2.imshow("fix_outline", image)
  1988. # cv2.waitKey(0)
  1989. # for line in all_longer_row_lines:
  1990. # cv2.line(image, (int(line[0]), int(line[1])),
  1991. # (int(line[2]), int(line[3])),
  1992. # (0, 0, 255), 2)
  1993. # cv2.imshow("fix_outline", image)
  1994. # cv2.waitKey(0)
  1995. return new_row_lines, new_col_lines, all_longer_row_lines, all_longer_col_lines
  1996. def fix_outline(image, row_lines, col_lines, points, split_y, scale=25):
  1997. log("into fix_outline")
  1998. x_min_len = max(10, int(image.shape[0] / scale))
  1999. y_min_len = max(10, int(image.shape[1] / scale))
  2000. # print("x_min_len", x_min_len, "y_min_len", y_min_len)
  2001. # print("split_y", split_y)
  2002. # 分割线纵坐标
  2003. if len(split_y) < 2:
  2004. return [], [], [], []
  2005. split_y.sort(key=lambda x: x)
  2006. new_split_y = []
  2007. for i in range(1, len(split_y), 2):
  2008. new_split_y.append(int((split_y[i]+split_y[i-1])/2))
  2009. split_row_list = []
  2010. split_col_list = []
  2011. split_point_list = []
  2012. for i in range(1, len(split_y)):
  2013. y = split_y[i]
  2014. last_y = split_y[i-1]
  2015. split_row = []
  2016. for row in row_lines:
  2017. if last_y <= row[3] <= y:
  2018. split_row.append(row)
  2019. split_row_list.append(split_row)
  2020. split_col = []
  2021. for col in col_lines:
  2022. if last_y <= col[1] <= y or last_y <= col[3] <= y or col[1] < last_y < y < col[3]:
  2023. split_col.append(col)
  2024. split_col_list.append(split_col)
  2025. split_point = []
  2026. for point in points:
  2027. if last_y <= point[1] <= y:
  2028. split_point.append(point)
  2029. split_point_list.append(split_point)
  2030. # 预测线取上下左右4个边(会有超出表格部分) [(), ()]
  2031. area_row_line = []
  2032. area_col_line = []
  2033. for area in split_row_list:
  2034. if not area:
  2035. area_row_line.append([])
  2036. continue
  2037. area.sort(key=lambda x: (x[1], x[0]))
  2038. up_line = area[0]
  2039. bottom_line = area[-1]
  2040. area_row_line.append([up_line, bottom_line])
  2041. for area in split_col_list:
  2042. if not area:
  2043. area_col_line.append([])
  2044. continue
  2045. area.sort(key=lambda x: x[0])
  2046. left_line = area[0]
  2047. right_line = area[-1]
  2048. area_col_line.append([left_line, right_line])
  2049. # 取每个分割区域的4条线(无超出表格部分)
  2050. area_row_line2 = []
  2051. area_col_line2 = []
  2052. for area in split_point_list:
  2053. if not area:
  2054. area_row_line2.append([])
  2055. area_col_line2.append([])
  2056. continue
  2057. area.sort(key=lambda x: (x[0], x[1]))
  2058. left_up = area[0]
  2059. right_bottom = area[-1]
  2060. up_line = [left_up[0], left_up[1], right_bottom[0], left_up[1]]
  2061. bottom_line = [left_up[0], right_bottom[1], right_bottom[0], right_bottom[1]]
  2062. left_line = [left_up[0], left_up[1], left_up[0], right_bottom[1]]
  2063. right_line = [right_bottom[0], left_up[1], right_bottom[0], right_bottom[1]]
  2064. area_row_line2.append([up_line, bottom_line])
  2065. area_col_line2.append([left_line, right_line])
  2066. # 判断超出部分的长度,超出一定长度就补线
  2067. new_row_lines = []
  2068. new_col_lines = []
  2069. longer_row_lines = []
  2070. longer_col_lines = []
  2071. all_longer_row_lines = []
  2072. all_longer_col_lines = []
  2073. for i in range(len(area_row_line)):
  2074. if not area_row_line[i] or not area_col_line[i]:
  2075. continue
  2076. up_line = area_row_line[i][0]
  2077. up_line2 = area_row_line2[i][0]
  2078. bottom_line = area_row_line[i][1]
  2079. bottom_line2 = area_row_line2[i][1]
  2080. left_line = area_col_line[i][0]
  2081. left_line2 = area_col_line2[i][0]
  2082. right_line = area_col_line[i][1]
  2083. right_line2 = area_col_line2[i][1]
  2084. # 计算单格高度宽度
  2085. if len(split_row_list[i]) > 1:
  2086. height_dict = {}
  2087. for j in range(len(split_row_list[i])):
  2088. if j + 1 > len(split_row_list[i]) - 1:
  2089. break
  2090. # print("height_dict", split_row_list[i][j], split_row_list[i][j+1])
  2091. height = abs(int(split_row_list[i][j][3] - split_row_list[i][j+1][3]))
  2092. if height >= 10:
  2093. if height in height_dict.keys():
  2094. height_dict[height] = height_dict[height] + 1
  2095. else:
  2096. height_dict[height] = 1
  2097. height_list = [[x, height_dict[x]] for x in height_dict.keys()]
  2098. height_list.sort(key=lambda x: (x[1], -x[0]), reverse=True)
  2099. # print("box_height", height_list)
  2100. box_height = height_list[0][0]
  2101. else:
  2102. box_height = y_min_len
  2103. if len(split_col_list[i]) > 1:
  2104. box_width = abs(split_col_list[i][1][2] - split_col_list[i][0][2])
  2105. else:
  2106. box_width = x_min_len
  2107. # print("box_height", box_height, "box_width", box_width)
  2108. # 设置轮廓线需超出阈值
  2109. if box_height >= 2*y_min_len:
  2110. fix_h_len = y_min_len
  2111. else:
  2112. fix_h_len = box_height * 2/3
  2113. if box_width >= 2*x_min_len:
  2114. fix_w_len = x_min_len
  2115. else:
  2116. fix_w_len = box_width * 2/3
  2117. # 补左右两条竖线超出来的线的row
  2118. if up_line[1] - left_line[1] >= fix_h_len and up_line[1] - right_line[1] >= fix_h_len:
  2119. if up_line[1] - left_line[1] >= up_line[1] - right_line[1]:
  2120. new_row_lines.append([left_line[0], left_line[1], right_line[0], left_line[1]])
  2121. new_col_y = left_line[1]
  2122. # 补了row,要将其他短的col连到row上
  2123. for j in range(len(split_col_list[i])):
  2124. col = split_col_list[i][j]
  2125. # 且距离不能相差大于一格
  2126. # print("abs(new_col_y - col[1])", abs(new_col_y - col[1]))
  2127. if abs(new_col_y - col[1]) <= box_height:
  2128. split_col_list[i][j][1] = min([new_col_y, col[1]])
  2129. longer_col_lines.append([col[0], min([new_col_y, col[1]]), col[2], col[3]])
  2130. else:
  2131. new_row_lines.append([left_line[0], right_line[1], right_line[0], right_line[1]])
  2132. new_col_y = right_line[1]
  2133. # 补了row,要将其他短的col连到row上
  2134. for j in range(len(split_col_list[i])):
  2135. # 需判断该线在这个区域中
  2136. # if up_line2[1]-3 <= col[1] <= col[3] <= bottom_line2[1]+3:
  2137. col = split_col_list[i][j]
  2138. # 且距离不能相差太大
  2139. # print("abs(new_col_y - col[1])", abs(new_col_y - col[1]))
  2140. if abs(new_col_y - col[1]) <= box_height:
  2141. split_col_list[i][j][1] = min([new_col_y, col[1]])
  2142. if left_line[3] - bottom_line[3] >= fix_h_len and right_line[3] - bottom_line[3] >= fix_h_len:
  2143. if left_line[3] - bottom_line[3] >= right_line[3] - bottom_line[3]:
  2144. new_row_lines.append([left_line[2], left_line[3], right_line[2], left_line[3]])
  2145. new_col_y = left_line[3]
  2146. # 补了row,要将其他短的col连到row上
  2147. for j in range(len(split_col_list[i])):
  2148. col = split_col_list[i][j]
  2149. # 且距离不能相差太大
  2150. if abs(new_col_y - col[3]) <= box_height:
  2151. split_col_list[i][j][3] = max([new_col_y, col[3]])
  2152. else:
  2153. new_row_lines.append([left_line[2], right_line[3], right_line[2], right_line[3]])
  2154. new_col_y = right_line[3]
  2155. # 补了row,要将其他短的col连到row上
  2156. for j in range(len(split_col_list[i])):
  2157. col = split_col_list[i][j]
  2158. # 且距离不能相差太大
  2159. if abs(new_col_y - col[3]) <= box_height:
  2160. split_col_list[i][j][3] = max([new_col_y, col[3]])
  2161. # 补上下两条横线超出来的线的col
  2162. if left_line[0] - up_line[0] >= fix_w_len and left_line[0] - bottom_line[0] >= fix_w_len:
  2163. if left_line[0] - up_line[0] >= left_line[0] - bottom_line[0]:
  2164. new_col_lines.append([up_line[0], up_line[1], up_line[0], bottom_line[1]])
  2165. new_row_x = up_line[0]
  2166. # 补了col,要将其他短的row连到col上
  2167. for j in range(len(split_row_list[i])):
  2168. row = split_row_list[i][j]
  2169. # 且距离不能相差太大
  2170. if abs(new_row_x - row[0]) <= box_width:
  2171. split_row_list[i][j][0] = min([new_row_x, row[0]])
  2172. else:
  2173. new_col_lines.append([bottom_line[0], up_line[1], bottom_line[0], bottom_line[1]])
  2174. new_row_x = bottom_line[0]
  2175. # 补了col,要将其他短的row连到col上
  2176. for j in range(len(split_row_list[i])):
  2177. row = split_row_list[i][j]
  2178. # 且距离不能相差太大
  2179. if abs(new_row_x - row[0]) <= box_width:
  2180. split_row_list[i][j][0] = min([new_row_x, row[0]])
  2181. if up_line[2] - right_line[2] >= fix_w_len and bottom_line[2] - right_line[2] >= fix_w_len:
  2182. if up_line[2] - right_line[2] >= bottom_line[2] - right_line[2]:
  2183. new_col_lines.append([up_line[2], up_line[3], up_line[2], bottom_line[3]])
  2184. new_row_x = up_line[2]
  2185. # 补了col,要将其他短的row连到col上
  2186. for j in range(len(split_row_list[i])):
  2187. row = split_row_list[i][j]
  2188. # 且距离不能相差太大
  2189. if abs(new_row_x - row[2]) <= box_width:
  2190. split_row_list[i][j][2] = max([new_row_x, row[2]])
  2191. else:
  2192. new_col_lines.append([bottom_line[2], up_line[3], bottom_line[2], bottom_line[3]])
  2193. new_row_x = bottom_line[2]
  2194. # 补了col,要将其他短的row连到col上
  2195. for j in range(len(split_row_list[i])):
  2196. # 需判断该线在这个区域中
  2197. # if up_line2[1]-3 <= row[1] <= bottom_line2[1]+3:
  2198. row = split_row_list[i][j]
  2199. # 且距离不能相差太大
  2200. if abs(new_row_x - row[2]) <= box_width:
  2201. split_row_list[i][j][2] = max([new_row_x, row[2]])
  2202. all_longer_row_lines += split_row_list[i]
  2203. all_longer_col_lines += split_col_list[i]
  2204. return new_row_lines, new_col_lines, all_longer_row_lines, all_longer_col_lines
  2205. def fix_table(row_point_list, col_point_list, split_y, row_lines, col_lines):
  2206. # 分割线纵坐标
  2207. if len(split_y) < 2:
  2208. return []
  2209. # 获取bbox
  2210. bbox = []
  2211. # 每个点获取与其x最相近和y最相近的点
  2212. for i in range(1, len(split_y)):
  2213. # 循环每行
  2214. for row in row_point_list:
  2215. row.sort(key=lambda x: (x[0], x[1]))
  2216. # 行不在该区域跳过
  2217. if row[0][1] <= split_y[i-1] or row[0][1] >= split_y[i]:
  2218. continue
  2219. # print("len(row)", len(row))
  2220. # print("row", row)
  2221. # 循环行中的点
  2222. for j in range(len(row)):
  2223. if j == len(row) - 1:
  2224. break
  2225. current_point = row[j]
  2226. next_point_in_row_list = row[j+1:]
  2227. # 循环这一行的下一个点
  2228. for next_point_in_row in next_point_in_row_list:
  2229. # 是否在这一行点找到,找不到就这一行的下个点
  2230. not_found = 1
  2231. # 查询下个点所在列
  2232. next_col = []
  2233. for col in col_point_list:
  2234. col.sort(key=lambda x: (x[1], x[0]))
  2235. # 列不在该区域跳过
  2236. if col[0][1] <= split_y[i-1] or col[-1][1] >= split_y[i]:
  2237. continue
  2238. if col[0][0]-3 <= next_point_in_row[0] <= col[0][0]+3:
  2239. next_col = col
  2240. break
  2241. # 循环匹配当前点和下一列点
  2242. next_col.sort(key=lambda x: (x[1], x[0]))
  2243. for point1 in next_col:
  2244. # 同一行的就跳过
  2245. if current_point[1]-3 <= point1[1] <= current_point[1]+3:
  2246. continue
  2247. if point1[1] <= current_point[1]-3:
  2248. continue
  2249. # 候选bbox
  2250. candidate_bbox = [current_point[0], current_point[1], point1[0], point1[1]]
  2251. # print("candidate_bbox", candidate_bbox)
  2252. # 判断该bbox是否存在,判断bbox的上下两条边是否有包含在row中
  2253. contain_flag1 = 0
  2254. contain_flag2 = 0
  2255. for row1 in row_lines:
  2256. # 行不在该区域跳过
  2257. if row1[1] <= split_y[i-1] or row1[1] >= split_y[i]:
  2258. continue
  2259. # bbox上边框 y一样
  2260. if not contain_flag1:
  2261. if row1[1]-3 <= candidate_bbox[1] <= row1[1]+3:
  2262. # 格子里的断开线段
  2263. row1_break = (max([row1[0], candidate_bbox[0]]),
  2264. row1[1],
  2265. min([row1[2], candidate_bbox[2]]),
  2266. row1[3])
  2267. if row1_break[2] - row1_break[0] >= (candidate_bbox[2] - candidate_bbox[0])/3:
  2268. contain_flag1 = 1
  2269. # bbox下边框 y一样
  2270. if not contain_flag2:
  2271. if row1[1]-3 <= candidate_bbox[3] <= row1[1]+3:
  2272. # 格子里的断开线段
  2273. row1_break = (max([row1[0], candidate_bbox[0]]),
  2274. row1[1],
  2275. min([row1[2], candidate_bbox[2]]),
  2276. row1[3])
  2277. if row1_break[2] - row1_break[0] >= (candidate_bbox[2] - candidate_bbox[0])/3:
  2278. contain_flag2 = 1
  2279. # 判断该bbox是否存在,判断bbox的左右两条边是否有包含在col中
  2280. contain_flag3 = 0
  2281. contain_flag4 = 0
  2282. for col1 in col_lines:
  2283. # 列不在该区域跳过
  2284. if col1[1] <= split_y[i-1] or col1[3] >= split_y[i]:
  2285. continue
  2286. # bbox左边线 x一样
  2287. if not contain_flag3:
  2288. if col1[0]-3 <= candidate_bbox[0] <= col1[0]+3:
  2289. # 格子里的断开线段
  2290. col1_break = (col1[0],
  2291. max([col1[1], candidate_bbox[1]]),
  2292. col1[2],
  2293. min([col1[3], candidate_bbox[3]]))
  2294. if col1_break[3] - col1_break[1] >= (candidate_bbox[3] - candidate_bbox[1])/3:
  2295. contain_flag3 = 1
  2296. # bbox右边框 x一样
  2297. if not contain_flag4:
  2298. if col1[0]-3 <= candidate_bbox[2] <= col1[0]+3:
  2299. # 格子里的断开线段
  2300. col1_break = (col1[0],
  2301. max([col1[1], candidate_bbox[1]]),
  2302. col1[2],
  2303. min([col1[3], candidate_bbox[3]]))
  2304. if col1_break[3] - col1_break[1] >= (candidate_bbox[3] - candidate_bbox[1])/3:
  2305. contain_flag4 = 1
  2306. # 找到了该bbox,并且是存在的
  2307. if contain_flag1 and contain_flag2 and contain_flag3 and contain_flag4:
  2308. bbox.append([(candidate_bbox[0], candidate_bbox[1]),
  2309. (candidate_bbox[2], candidate_bbox[3])])
  2310. not_found = 0
  2311. break
  2312. if not not_found:
  2313. break
  2314. return bbox
  2315. def delete_close_points(point_list, row_point_list, col_point_list, threshold=5):
  2316. new_point_list = []
  2317. delete_point_list = []
  2318. point_list.sort(key=lambda x: (x[1], x[0]))
  2319. for i in range(len(point_list)):
  2320. point1 = point_list[i]
  2321. if point1 in delete_point_list:
  2322. continue
  2323. if i == len(point_list) - 1:
  2324. new_point_list.append(point1)
  2325. break
  2326. point2 = point_list[i+1]
  2327. # 判断坐标
  2328. if abs(point1[0] - point2[0]) > threshold or abs(point1[1] - point2[1]) > threshold:
  2329. new_point_list.append(point1)
  2330. else:
  2331. # 看两个点上的相同坐标点哪个多,就保留哪个
  2332. count1 = 0
  2333. count2 = 0
  2334. for col in col_point_list:
  2335. if point1[0] == col[0][0]:
  2336. count1 += len(col)
  2337. elif point2[0] == col[0][0]:
  2338. count2 += len(col)
  2339. if count1 >= count2:
  2340. new_point_list.append(point1)
  2341. delete_point_list.append(point2)
  2342. else:
  2343. new_point_list.append(point2)
  2344. delete_point_list.append(point1)
  2345. point_list = new_point_list
  2346. new_point_list = []
  2347. delete_point_list = []
  2348. point_list.sort(key=lambda x: (x[0], x[1]))
  2349. for i in range(len(point_list)):
  2350. point1 = point_list[i]
  2351. if point1 in delete_point_list:
  2352. continue
  2353. if i == len(point_list) - 1:
  2354. new_point_list.append(point1)
  2355. break
  2356. point2 = point_list[i+1]
  2357. # 判断坐标
  2358. if abs(point1[0] - point2[0]) > threshold or abs(point1[1] - point2[1]) > threshold:
  2359. new_point_list.append(point1)
  2360. else:
  2361. count1 = 0
  2362. count2 = 0
  2363. for row in row_point_list:
  2364. if point1[0] == row[0][0]:
  2365. count1 += len(row)
  2366. elif point2[0] == row[0][0]:
  2367. count2 += len(row)
  2368. if count1 >= count2:
  2369. new_point_list.append(point1)
  2370. delete_point_list.append(point2)
  2371. else:
  2372. new_point_list.append(point2)
  2373. delete_point_list.append(point1)
  2374. return new_point_list
  2375. def get_bbox2(image_np, points):
  2376. # # 坐标点按行分
  2377. # row_point_list = []
  2378. # row_point = []
  2379. # points.sort(key=lambda x: (x[0], x[1]))
  2380. # for p in points:
  2381. # if len(row_point) == 0:
  2382. # x = p[0]
  2383. # if x-5 <= p[0] <= x+5:
  2384. # row_point.append(p)
  2385. # else:
  2386. # row_point_list.append(row_point)
  2387. # row_point = []
  2388. # # 坐标点按列分
  2389. # col_point_list = []
  2390. # col_point = []
  2391. # points.sort(key=lambda x: (x[1], x[0]))
  2392. # for p in points:
  2393. # if len(col_point) == 0:
  2394. # y = p[1]
  2395. # if y-5 <= p[1] <= y+5:
  2396. # col_point.append(p)
  2397. # else:
  2398. # col_point_list.append(col_point)
  2399. # col_point = []
  2400. row_point_list = get_points_row(points)
  2401. col_point_list = get_points_col(points)
  2402. print("len(points)", len(points))
  2403. for point in points:
  2404. cv2.circle(image_np, point, 1, (0, 255, 0), 1)
  2405. cv2.imshow("points_deleted", image_np)
  2406. points = delete_close_points(points, row_point_list, col_point_list)
  2407. print("len(points)", len(points))
  2408. for point in points:
  2409. cv2.circle(image_np, point, 1, (255, 0, 0), 3)
  2410. cv2.imshow("points_deleted", image_np)
  2411. cv2.waitKey(0)
  2412. row_point_list = get_points_row(points, 5)
  2413. col_point_list = get_points_col(points, 5)
  2414. print("len(row_point_list)", len(row_point_list))
  2415. for row in row_point_list:
  2416. print("row", len(row))
  2417. print("col_point_list", len(col_point_list))
  2418. for col in col_point_list:
  2419. print("col", len(col))
  2420. bbox = []
  2421. for i in range(len(row_point_list)):
  2422. if i == len(row_point_list) - 1:
  2423. break
  2424. # 遍历每个row的point,找到其所在列的下一个点和所在行的下一个点
  2425. current_row = row_point_list[i]
  2426. for j in range(len(current_row)):
  2427. current_point = current_row[j]
  2428. if j == len(current_row) - 1:
  2429. break
  2430. next_row_point = current_row[j+1]
  2431. # 找出当前点所在的col,得到该列下一个point
  2432. current_col = col_point_list[j]
  2433. for k in range(len(current_col)):
  2434. if current_col[k][1] > current_point[1] + 10:
  2435. next_col_point = current_col[k]
  2436. break
  2437. next_row = row_point_list[k]
  2438. for k in range(len(next_row)):
  2439. if next_row[k][0] >= next_row_point[0] + 5:
  2440. next_point = next_row[k]
  2441. break
  2442. # 得到bbox
  2443. bbox.append([(current_point[0], current_point[1]), (next_point[0], next_point[1])])
  2444. # bbox = []
  2445. # for p in points:
  2446. # # print("p", p)
  2447. # p_row = []
  2448. # p_col = []
  2449. # for row in row_point_list:
  2450. # if p[0] == row[0][0]:
  2451. # for p1 in row:
  2452. # if abs(p[1]-p1[1]) <= 5:
  2453. # continue
  2454. # p_row.append([p1, abs(p[1]-p1[1])])
  2455. # p_row.sort(key=lambda x: x[1])
  2456. # for col in col_point_list:
  2457. # if p[1] == col[0][1]:
  2458. # for p2 in col:
  2459. # if abs(p[0]-p2[0]) <= 5:
  2460. # continue
  2461. # p_col.append([p2, abs(p[0]-p2[0])])
  2462. # p_col.sort(key=lambda x: x[1])
  2463. # if len(p_row) == 0 or len(p_col) == 0:
  2464. # continue
  2465. # break_flag = 0
  2466. # for i in range(len(p_row)):
  2467. # for j in range(len(p_col)):
  2468. # # print(p_row[i][0])
  2469. # # print(p_col[j][0])
  2470. # another_point = (p_col[j][0][0], p_row[i][0][1])
  2471. # # print("another_point", another_point)
  2472. # if abs(p[0]-another_point[0]) <= 5 or abs(p[1]-another_point[1]) <= 5:
  2473. # continue
  2474. # if p[0] >= another_point[0] or p[1] >= another_point[1]:
  2475. # continue
  2476. # if another_point in points:
  2477. # box = [p, another_point]
  2478. # box.sort(key=lambda x: x[0])
  2479. # if box not in bbox:
  2480. # bbox.append(box)
  2481. # break_flag = 1
  2482. # break
  2483. # if break_flag:
  2484. # break
  2485. #
  2486. # # delete duplicate
  2487. # delete_bbox = []
  2488. # for i in range(len(bbox)):
  2489. # for j in range(i+1, len(bbox)):
  2490. # if bbox[i][0] == bbox[j][0]:
  2491. # if bbox[i][1][0] - bbox[j][1][0] <= 3 \
  2492. # and bbox[i][1][1] - bbox[j][1][1] <= 3:
  2493. # delete_bbox.append(bbox[j])
  2494. # if bbox[i][1] == bbox[j][1]:
  2495. # if bbox[i][0][0] - bbox[j][0][0] <= 3 \
  2496. # and bbox[i][0][1] - bbox[j][0][1] <= 3:
  2497. # delete_bbox.append(bbox[j])
  2498. # # delete too small area
  2499. # # for box in bbox:
  2500. # # if box[1][0] - box[0][0] <=
  2501. # for d_box in delete_bbox:
  2502. # if d_box in bbox:
  2503. # bbox.remove(d_box)
  2504. # print bbox
  2505. bbox.sort(key=lambda x: (x[0][0], x[0][1], x[1][0], x[1][1]))
  2506. # origin bbox
  2507. # origin_bbox = []
  2508. # for box in bbox:
  2509. # origin_bbox.append([(box[0][0], box[0][1] - 40), (box[1][0], box[1][1] - 40)])
  2510. # for box in origin_bbox:
  2511. # cv2.rectangle(origin_image, box[0], box[1], (0, 0, 255), 2, 8)
  2512. # cv2.imshow('AlanWang', origin_image)
  2513. # cv2.waitKey(0)
  2514. for box in bbox:
  2515. cv2.rectangle(image_np, box[0], box[1], (0, 0, 255), 2, 8)
  2516. cv2.imshow('bboxes', image_np)
  2517. cv2.waitKey(0)
  2518. # for point in points:
  2519. # print(point)
  2520. # cv2.circle(image_np, point, 1, (0, 0, 255), 3)
  2521. # cv2.imshow('points', image_np)
  2522. # cv2.waitKey(0)
  2523. return bbox
  2524. def get_bbox1(image_np, points, split_y):
  2525. # 分割线纵坐标
  2526. # print("split_y", split_y)
  2527. if len(split_y) < 2:
  2528. return []
  2529. # 计算行列,剔除相近交点
  2530. row_point_list = get_points_row(points)
  2531. col_point_list = get_points_col(points)
  2532. print("len(row_point_list)", row_point_list)
  2533. print("len(col_point_list)", len(col_point_list))
  2534. # for point in points:
  2535. # cv2.circle(image_np, point, 1, (0, 255, 0), 1)
  2536. # cv2.imshow("points", image_np)
  2537. points = delete_close_points(points, row_point_list, col_point_list)
  2538. # print("len(points)", len(points))
  2539. # for point in points:
  2540. # cv2.circle(image_np, point, 1, (255, 0, 0), 3)
  2541. # cv2.imshow("points_deleted", image_np)
  2542. # cv2.waitKey(0)
  2543. # 获取bbox
  2544. bbox = []
  2545. # 每个点获取与其x最相近和y最相近的点
  2546. for i in range(1, len(split_y)):
  2547. for point1 in points:
  2548. if point1[1] <= split_y[i-1] or point1[1] >= split_y[i]:
  2549. continue
  2550. distance_x = 10000
  2551. distance_y = 10000
  2552. x = 0
  2553. y = 0
  2554. threshold = 10
  2555. for point2 in points:
  2556. if point2[1] <= split_y[i-1] or point2[1] >= split_y[i]:
  2557. continue
  2558. # 最近 x y
  2559. if 2 < point2[0] - point1[0] < distance_x and point2[1] - point1[1] <= threshold:
  2560. distance_x = point2[0] - point1[0]
  2561. x = point2[0]
  2562. if 2 < point2[1] - point1[1] < distance_y and point2[0] - point1[0] <= threshold:
  2563. distance_y = point2[1] - point1[1]
  2564. y = point2[1]
  2565. if not x or not y:
  2566. continue
  2567. bbox.append([(point1[0], point1[1]), (x, y)])
  2568. # 删除包含关系bbox
  2569. temp_list = []
  2570. for i in range(len(bbox)):
  2571. box1 = bbox[i]
  2572. for j in range(len(bbox)):
  2573. if i == j:
  2574. continue
  2575. box2 = bbox[j]
  2576. contain_flag = 0
  2577. if box2[0][0] <= box1[0][0] <= box1[1][0] <= box2[1][0] and \
  2578. box2[0][1] <= box1[0][1] <= box1[1][1] <= box2[1][1]:
  2579. contain_flag = 1
  2580. break
  2581. temp_list.append(box1)
  2582. bbox = temp_list
  2583. # 展示
  2584. for box in bbox:
  2585. # print(box[0], box[1])
  2586. # if abs(box[0][1] - box[1][1]) > abs(box[0][0] - box[1][0]):
  2587. # continue
  2588. cv2.rectangle(image_np, box[0], box[1], (0, 0, 255), 2, 8)
  2589. cv2.imshow('bboxes', image_np)
  2590. cv2.waitKey(0)
  2591. return bbox
  2592. def get_bbox0(image_np, row_point_list, col_point_list, split_y, row_lines, col_lines):
  2593. # 分割线纵坐标
  2594. if len(split_y) < 2:
  2595. return []
  2596. # 计算行列,剔除相近交点
  2597. # row_point_list = get_points_row(points)
  2598. # col_point_list = get_points_col(points)
  2599. # points = delete_close_points(points, row_point_list, col_point_list)
  2600. # row_point_list = get_points_row(points)
  2601. # col_point_list = get_points_col(points)
  2602. # 获取bbox
  2603. bbox = []
  2604. # print("get_bbox split_y", split_y)
  2605. # 每个点获取与其x最相近和y最相近的点
  2606. for i in range(1, len(split_y)):
  2607. # 循环每行
  2608. for row in row_point_list:
  2609. row.sort(key=lambda x: (x[0], x[1]))
  2610. # 行不在该区域跳过
  2611. if row[0][1] <= split_y[i-1] or row[0][1] >= split_y[i]:
  2612. continue
  2613. # 循环行中的点
  2614. for j in range(len(row)):
  2615. if j == len(row) - 1:
  2616. break
  2617. current_point = row[j]
  2618. next_point_in_row = row[j+1]
  2619. # 查询下个点所在列
  2620. next_col = []
  2621. for col in col_point_list:
  2622. col.sort(key=lambda x: (x[1], x[0]))
  2623. # 列不在该区域跳过
  2624. if col[0][1] <= split_y[i-1] or col[-1][1] >= split_y[i]:
  2625. continue
  2626. if col[0][0]-3 <= next_point_in_row[0] <= col[0][0]+3:
  2627. next_col = col
  2628. break
  2629. # 循环匹配当前点和下一列点
  2630. for point1 in next_col:
  2631. # 同一行的就跳过
  2632. if current_point[1]-3 <= point1[1] <= current_point[1]+3:
  2633. continue
  2634. if point1[1] <= current_point[1]-3:
  2635. continue
  2636. # 候选bbox
  2637. candidate_bbox = [current_point[0], current_point[1], point1[0], point1[1]]
  2638. # 判断该bbox是否存在,线条包含关系
  2639. contain_flag1 = 0
  2640. contain_flag2 = 0
  2641. for row1 in row_lines:
  2642. # 行不在该区域跳过
  2643. if row1[1] <= split_y[i-1] or row1[1] >= split_y[i]:
  2644. continue
  2645. # bbox上边框 y一样
  2646. if not contain_flag1:
  2647. if row1[1]-3 <= candidate_bbox[1] <= row1[1]+3:
  2648. # candidate的x1,x2需被包含在row线中
  2649. if row1[0]-3 <= candidate_bbox[0] <= candidate_bbox[2] <= row1[2]+3:
  2650. contain_flag1 = 1
  2651. # bbox下边框 y一样
  2652. if not contain_flag2:
  2653. if row1[1]-3 <= candidate_bbox[3] <= row1[1]+3:
  2654. # candidate的x1,x2需被包含在row线中
  2655. if row1[0]-3 <= candidate_bbox[0] <= candidate_bbox[2] <= row1[2]+3:
  2656. contain_flag2 = 1
  2657. # 找到了该bbox,并且是存在的
  2658. if contain_flag1 and contain_flag2:
  2659. bbox.append([(candidate_bbox[0], candidate_bbox[1]),
  2660. (candidate_bbox[2], candidate_bbox[3])])
  2661. break
  2662. return bbox
  2663. def get_bbox3(image_np, row_point_list, col_point_list, split_y, row_lines, col_lines):
  2664. # 分割线纵坐标
  2665. if len(split_y) < 2:
  2666. return []
  2667. # 获取bbox
  2668. bbox = []
  2669. # 每个点获取与其x最相近和y最相近的点
  2670. for i in range(1, len(split_y)):
  2671. # 循环每行
  2672. for row in row_point_list:
  2673. row.sort(key=lambda x: (x[0], x[1]))
  2674. # 行不在该区域跳过
  2675. if row[0][1] <= split_y[i-1] or row[0][1] >= split_y[i]:
  2676. continue
  2677. # print("len(row)", len(row))
  2678. # print("row", row)
  2679. # 循环行中的点
  2680. for j in range(len(row)):
  2681. if j == len(row) - 1:
  2682. break
  2683. current_point = row[j]
  2684. # print("current_point", current_point)
  2685. next_point_in_row_list = row[j+1:]
  2686. # 循环这一行的下一个点
  2687. for next_point_in_row in next_point_in_row_list:
  2688. # 是否在这一行点找到,找不到就这一行的下个点
  2689. not_found = 1
  2690. # 查询下个点所在列
  2691. next_col = []
  2692. for col in col_point_list:
  2693. col.sort(key=lambda x: (x[1], x[0]))
  2694. # 列不在该区域跳过
  2695. if col[0][1] <= split_y[i-1] or col[-1][1] >= split_y[i]:
  2696. continue
  2697. if col[0][0]-3 <= next_point_in_row[0] <= col[0][0]+3:
  2698. next_col = col
  2699. break
  2700. # 循环匹配当前点和下一列点
  2701. next_col.sort(key=lambda x: (x[1], x[0]))
  2702. for point1 in next_col:
  2703. # 同一行的就跳过
  2704. if current_point[1]-3 <= point1[1] <= current_point[1]+3:
  2705. continue
  2706. if point1[1] <= current_point[1]-3:
  2707. continue
  2708. # 候选bbox
  2709. candidate_bbox = [current_point[0], current_point[1], point1[0], point1[1]]
  2710. # print("candidate_bbox", candidate_bbox)
  2711. # 判断该bbox是否存在,判断bbox的上下两条边是否有包含在row中
  2712. contain_flag1 = 0
  2713. contain_flag2 = 0
  2714. for row1 in row_lines:
  2715. # 行不在该区域跳过
  2716. if row1[1] <= split_y[i-1] or row1[1] >= split_y[i]:
  2717. continue
  2718. # bbox上边框 y一样
  2719. if not contain_flag1:
  2720. if row1[1]-3 <= candidate_bbox[1] <= row1[1]+3:
  2721. # 格子里的断开线段
  2722. row1_break = (max([row1[0], candidate_bbox[0]]),
  2723. row1[1],
  2724. min([row1[2], candidate_bbox[2]]),
  2725. row1[3])
  2726. if row1_break[2] - row1_break[0] >= (candidate_bbox[2] - candidate_bbox[0])/3:
  2727. contain_flag1 = 1
  2728. # # candidate的x1,x2需被包含在row线中
  2729. # if row1[0]-3 <= candidate_bbox[0] <= candidate_bbox[2] <= row1[2]+3:
  2730. # contain_flag1 = 1
  2731. #
  2732. # # 判断线条有无端点在格子中
  2733. # elif candidate_bbox[0] < row1[0] < candidate_bbox[2] \
  2734. # or candidate_bbox[0] < row1[2] < candidate_bbox[2]:
  2735. # # 线条会有缺一点情况,判断长度超过格子一半
  2736. # if row1_break[2] - row1_break[0] >= (candidate_bbox[2] - candidate_bbox[0])/3:
  2737. # contain_flag1 = 1
  2738. # bbox下边框 y一样
  2739. if not contain_flag2:
  2740. if row1[1]-3 <= candidate_bbox[3] <= row1[1]+3:
  2741. # 格子里的断开线段
  2742. row1_break = (max([row1[0], candidate_bbox[0]]),
  2743. row1[1],
  2744. min([row1[2], candidate_bbox[2]]),
  2745. row1[3])
  2746. if row1_break[2] - row1_break[0] >= (candidate_bbox[2] - candidate_bbox[0])/3:
  2747. contain_flag2 = 1
  2748. # # candidate的x1,x2需被包含在row线中
  2749. # if row1[0]-3 <= candidate_bbox[0] <= candidate_bbox[2] <= row1[2]+3:
  2750. # contain_flag2 = 1
  2751. #
  2752. # # 判断线条有无端点在格子中
  2753. # elif candidate_bbox[0] < row1[0] < candidate_bbox[2] \
  2754. # or candidate_bbox[0] < row1[2] < candidate_bbox[2]:
  2755. # # 线条会有缺一点情况,判断长度超过格子一半
  2756. # if row1_break[2] - row1_break[0] >= (candidate_bbox[2] - candidate_bbox[0])/3:
  2757. # contain_flag2 = 1
  2758. # 判断该bbox是否存在,判断bbox的左右两条边是否有包含在col中
  2759. contain_flag3 = 0
  2760. contain_flag4 = 0
  2761. for col1 in col_lines:
  2762. # 列不在该区域跳过
  2763. if col1[1] <= split_y[i-1] or col1[3] >= split_y[i]:
  2764. continue
  2765. # bbox左边线 x一样
  2766. if not contain_flag3:
  2767. if col1[0]-3 <= candidate_bbox[0] <= col1[0]+3:
  2768. # 格子里的断开线段
  2769. col1_break = (col1[0],
  2770. max([col1[1], candidate_bbox[1]]),
  2771. col1[2],
  2772. min([col1[3], candidate_bbox[3]]))
  2773. if col1_break[3] - col1_break[1] >= (candidate_bbox[3] - candidate_bbox[1])/3:
  2774. contain_flag3 = 1
  2775. # # candidate的y1,y2需被包含在col线中
  2776. # if col1[1]-3 <= candidate_bbox[1] <= candidate_bbox[3] <= col1[3]+3:
  2777. # contain_flag3 = 1
  2778. #
  2779. # # 判断线条有无端点在格子中
  2780. # elif candidate_bbox[1] < col1[1] < candidate_bbox[3] \
  2781. # or candidate_bbox[1] < col1[3] < candidate_bbox[3]:
  2782. # # 线条会有缺一点情况,判断长度超过格子一半
  2783. # if col1_break[3] - col1_break[1] >= (candidate_bbox[3] - candidate_bbox[1])/3:
  2784. # contain_flag3 = 1
  2785. # bbox右边框 x一样
  2786. if not contain_flag4:
  2787. if col1[0]-3 <= candidate_bbox[2] <= col1[0]+3:
  2788. # 格子里的断开线段
  2789. # col1_break = (col1[0],
  2790. # max([col1[1], candidate_bbox[1]]),
  2791. # col1[2],
  2792. # min([col1[3], candidate_bbox[3]]))
  2793. # if col1_break[3] - col1_break[1] >= (candidate_bbox[3] - candidate_bbox[1])/3:
  2794. # contain_flag4 = 1
  2795. # 如果候选bbox的边的上1/3或下1/3包含在col中
  2796. candidate_bbox_line1 = [candidate_bbox[1],
  2797. candidate_bbox[1] + (candidate_bbox[3]-candidate_bbox[1])/3]
  2798. candidate_bbox_line2 = [candidate_bbox[3] - (candidate_bbox[3]-candidate_bbox[1])/3,
  2799. candidate_bbox[3]]
  2800. if col1[1] <= candidate_bbox_line1[0] <= candidate_bbox_line1[1] <= col1[3] \
  2801. or col1[1] <= candidate_bbox_line2[0] <= candidate_bbox_line2[1] <= col1[3]:
  2802. # print("candidate_bbox", candidate_bbox)
  2803. # print("col1", col1)
  2804. contain_flag4 = 1
  2805. # # candidate的y1,y2需被包含在col线中
  2806. # if col1[1]-3 <= candidate_bbox[1] <= candidate_bbox[3] <= col1[3]+3:
  2807. # contain_flag4 = 1
  2808. #
  2809. # # 判断线条有无端点在格子中
  2810. # elif candidate_bbox[1] < col1[1] < candidate_bbox[3] \
  2811. # or candidate_bbox[1] < col1[3] < candidate_bbox[3]:
  2812. # # 线条会有缺一点情况,判断长度超过格子一半
  2813. # if col1_break[3] - col1_break[1] >= (candidate_bbox[3] - candidate_bbox[1])/3:
  2814. # contain_flag4 = 1
  2815. # 找到了该bbox,并且是存在的
  2816. if contain_flag1 and contain_flag2 and contain_flag3 and contain_flag4:
  2817. bbox.append([(candidate_bbox[0], candidate_bbox[1]),
  2818. (candidate_bbox[2], candidate_bbox[3])])
  2819. not_found = 0
  2820. # print("exist candidate_bbox", candidate_bbox)
  2821. # print(contain_flag1, contain_flag2, contain_flag3, contain_flag4)
  2822. break
  2823. # else:
  2824. # print("candidate_bbox", candidate_bbox)
  2825. # print(contain_flag1, contain_flag2, contain_flag3, contain_flag4)
  2826. if not not_found:
  2827. break
  2828. return bbox
  2829. def get_bbox(image_np, row_point_list, col_point_list, split_y, row_lines, col_lines):
  2830. # 分割线纵坐标
  2831. if len(split_y) < 2:
  2832. return []
  2833. # 获取bbox
  2834. bbox_list = []
  2835. for i in range(1, len(split_y)):
  2836. last_y = split_y[i-1]
  2837. y = split_y[i]
  2838. # 先对点线进行分区
  2839. split_row_point_list = []
  2840. split_col_point_list = []
  2841. split_row_lines = []
  2842. split_col_lines = []
  2843. for row in row_point_list:
  2844. if last_y <= row[0][1] <= y:
  2845. row.sort(key=lambda x: (x[1], x[0]))
  2846. split_row_point_list.append(row)
  2847. for col in col_point_list:
  2848. if last_y <= col[0][1] <= y:
  2849. split_col_point_list.append(col)
  2850. for row in row_lines:
  2851. if last_y <= row[1] <= y:
  2852. split_row_lines.append(row)
  2853. for col in col_lines:
  2854. if last_y <= col[1] <= y:
  2855. split_col_lines.append(col)
  2856. # 每个点获取其对角线点,以便形成bbox,按行循环
  2857. for i in range(len(split_row_point_list)-1):
  2858. row = split_row_point_list[i]
  2859. # 循环该行的点
  2860. for k in range(len(row)-1):
  2861. point1 = row[k]
  2862. next_point1 = row[k+1]
  2863. # print("*"*30)
  2864. # print("point1", point1)
  2865. # 有三种对角线点
  2866. # 1. 该点下一行的下一列的点
  2867. # 2. 该点下一列的下一行的点
  2868. # 3. 上述两个点是同一个点
  2869. # 下一行没找到就循环后面的行
  2870. if_find = 0
  2871. for j in range(i+1, len(split_row_point_list)):
  2872. if if_find:
  2873. break
  2874. next_row = split_row_point_list[j]
  2875. # print("next_row", next_row)
  2876. # 循环下一行的点
  2877. for point2 in next_row:
  2878. if abs(point1[0] - point2[0]) <= 2:
  2879. continue
  2880. if point2[0] < point1[0]:
  2881. continue
  2882. bbox = [point1[0], point1[1], point2[0], point2[1]]
  2883. if abs(bbox[0] - bbox[2]) <= 10:
  2884. continue
  2885. if abs(bbox[1] - bbox[3]) <= 10:
  2886. continue
  2887. # bbox的四条边都需要验证是否在line上
  2888. if check_bbox(bbox, split_row_lines, split_col_lines):
  2889. bbox_list.append([(bbox[0], bbox[1]), (bbox[2], bbox[3])])
  2890. if_find = 1
  2891. # print("check bbox", bbox)
  2892. break
  2893. return bbox_list
  2894. def check_bbox(bbox, rows, cols, threshold=5):
  2895. def check(check_line, lines, limit_axis, axis):
  2896. # 需检查的线的1/2段,1/3段,2/3段,1/4段,3/4段
  2897. line_1_2 = [check_line[0], (check_line[0]+check_line[1])/2]
  2898. line_2_2 = [(check_line[0]+check_line[1])/2, check_line[1]]
  2899. line_1_3 = [check_line[0], check_line[0]+(check_line[1]-check_line[0])/3]
  2900. line_2_3 = [check_line[1]-(check_line[1]-check_line[0])/3, check_line[1]]
  2901. line_1_4 = [check_line[0], check_line[0]+(check_line[1]-check_line[0])/4]
  2902. line_3_4 = [check_line[1]-(check_line[1]-check_line[0])/4, check_line[1]]
  2903. # 限制row相同y,col相同x
  2904. if_line = 0
  2905. for line1 in lines:
  2906. if not if_line and abs(line1[1-axis] - limit_axis) <= threshold:
  2907. # check_line完全包含在line中
  2908. if line1[axis] <= check_line[0] <= check_line[1] <= line1[axis+2]:
  2909. if_line = 1
  2910. # check_line的1/2包含在line
  2911. elif line1[axis] <= line_1_2[0] <= line_1_2[1] <= line1[axis+2] \
  2912. or line1[axis] <= line_2_2[0] <= line_2_2[1] <= line1[axis+2]:
  2913. if_line = 1
  2914. # check_line两个1/3段被包含在不同line中
  2915. elif line1[axis] <= line_1_3[0] <= line_1_3[1] <= line1[axis+2]:
  2916. # check_line另一边的1/4被包含
  2917. for line2 in lines:
  2918. if abs(line1[1-axis] - limit_axis) <= threshold:
  2919. if line2[axis] <= line_3_4[0] <= line_3_4[1] <= line2[axis+2]:
  2920. if_line = 1
  2921. break
  2922. elif line1[axis] <= line_2_3[0] <= line_2_3[1] <= line1[axis+2]:
  2923. # check_line另一边的1/4被包含
  2924. for line2 in lines:
  2925. if abs(line1[1-axis] - limit_axis) <= threshold:
  2926. if line2[axis] <= line_1_4[0] <= line_1_4[1] <= line2[axis+2]:
  2927. if_line = 1
  2928. break
  2929. return if_line
  2930. up_down_line = [bbox[0], bbox[2]]
  2931. up_y, down_y = bbox[1], bbox[3]
  2932. left_right_line = [bbox[1], bbox[3]]
  2933. left_x, right_x = bbox[0], bbox[2]
  2934. # 检查bbox四条边是否存在
  2935. if_up = check(up_down_line, rows, up_y, 0)
  2936. if_down = check(up_down_line, rows, down_y, 0)
  2937. if_left = check(left_right_line, cols, left_x, 1)
  2938. if_right = check(left_right_line, cols, right_x, 1)
  2939. # 检查bbox内部除了四条边,是否有其它line在bbox内部
  2940. if_col = 0
  2941. if_row = 0
  2942. if if_up and if_down and if_left and if_right:
  2943. for col in cols:
  2944. if not if_col and left_x+threshold <= col[0] <= right_x-threshold:
  2945. if col[1] <= left_right_line[0] <= left_right_line[1] <= col[3]:
  2946. if_col = 1
  2947. elif left_right_line[0] <= col[1] <= left_right_line[1]:
  2948. if left_right_line[1] - col[1] >= (left_right_line[1] + left_right_line[0])/2:
  2949. if_col = 1
  2950. elif left_right_line[0] <= col[3] <= left_right_line[1]:
  2951. if col[3] - left_right_line[0] >= (left_right_line[1] + left_right_line[0])/2:
  2952. if_col = 1
  2953. for row in rows:
  2954. if not if_row and up_y+threshold <= row[1] <= down_y-threshold:
  2955. if row[0] <= up_down_line[0] <= up_down_line[1] <= row[2]:
  2956. if_row = 1
  2957. elif up_down_line[0] <= row[0] <= up_down_line[1]:
  2958. if up_down_line[1] - row[0] >= (up_down_line[1] + up_down_line[0])/2:
  2959. if_row = 1
  2960. elif up_down_line[0] <= row[2] <= up_down_line[1]:
  2961. if row[2] - up_down_line[0] >= (up_down_line[1] + up_down_line[0])/2:
  2962. if_row = 1
  2963. if if_up and if_down and if_left and if_right and not if_col and not if_row:
  2964. return True
  2965. else:
  2966. return False
  2967. # def add_continue_bbox(bboxes):
  2968. # add_bbox_list = []
  2969. # bboxes.sort(key=lambda x: (x[0][0], x[0][1]))
  2970. # last_bbox = bboxes[0]
  2971. # # 先对bbox分区
  2972. # for i in range(1, len(split_y)):
  2973. # y = split_y[i]
  2974. # last_y = split_y[i-1]
  2975. #
  2976. # split_bbox = []
  2977. # for bbox in bboxes:
  2978. # if last_y <= bbox[1][1] <= y:
  2979. # split_bbox.append(bbox)
  2980. # split_bbox.sort
  2981. #
  2982. # for i in range(1, len(bboxes)):
  2983. # bbox = bboxes[i]
  2984. # if last_y <= bbox[1][1] <= y and last_y <= last_bbox[1][1] <= y:
  2985. # if abs(last_bbox[1][1] - bbox[0][1]) <= 2:
  2986. # last_bbox = bbox
  2987. # else:
  2988. # if last_bbox[1][1] > bbox[0][1]:
  2989. # last_bbox = bbox
  2990. # else:
  2991. # add_bbox = [(last_bbox[0][0], last_bbox[1][1]),
  2992. # (last_bbox[1][0], bbox[0][1])]
  2993. # add_bbox_list.append(add_bbox)
  2994. # last_y = y
  2995. #
  2996. # print("add_bbox_list", add_bbox_list)
  2997. # if add_bbox_list:
  2998. # bboxes = [str(x) for x in bboxes + add_bbox_list]
  2999. # bboxes = list(set(bboxes))
  3000. # bboxes = [eval(x) for x in bboxes]
  3001. # bboxes.sort(key=lambda x: (x[0][1], x[0][0]))
  3002. #
  3003. # return bboxes
  3004. def points_to_line(points_lines, axis):
  3005. new_line_list = []
  3006. for line in points_lines:
  3007. average = 0
  3008. _min = _min = line[0][axis]
  3009. _max = line[-1][axis]
  3010. for point in line:
  3011. average += point[1-axis]
  3012. if point[axis] < _min:
  3013. _min = point[axis]
  3014. if point[axis] > _max:
  3015. _max = point[axis]
  3016. average = int(average / len(line))
  3017. if axis:
  3018. new_line = [average, _min, average, _max]
  3019. else:
  3020. new_line = [_min, average, _max, average]
  3021. new_line_list.append(new_line)
  3022. return new_line_list
  3023. def get_bbox_by_contours(image_np):
  3024. img_gray = cv2.cvtColor(image_np, cv2.COLOR_BGR2GRAY)
  3025. ret, img_bin = cv2.threshold(img_gray, 127, 255, cv2.THRESH_BINARY)
  3026. # 3.连通域分析
  3027. img_bin, contours, hierarchy = cv2.findContours(img_bin,
  3028. cv2.RETR_LIST,
  3029. cv2.CHAIN_APPROX_SIMPLE)
  3030. # 4.获取最小外接圆 圆心 半径
  3031. center, radius = cv2.minEnclosingTriangle(contours[0])
  3032. center = np.int0(center)
  3033. # 5.绘制最小外接圆
  3034. img_result = image_np.copy()
  3035. cv2.circle(img_result, tuple(center), int(radius), (255, 255, 255), 2)
  3036. # # 读入图片
  3037. # img = image_np
  3038. # cv2.imshow("get_bbox_by_contours ", image_np)
  3039. # # 中值滤波,去噪
  3040. # img = cv2.medianBlur(img, 3)
  3041. # gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  3042. # cv2.namedWindow('original', cv2.WINDOW_AUTOSIZE)
  3043. # cv2.imshow('original', gray)
  3044. #
  3045. # # 阈值分割得到二值化图片
  3046. # ret, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
  3047. #
  3048. # # 膨胀操作
  3049. # kernel2 = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
  3050. # bin_clo = cv2.dilate(binary, kernel2, iterations=2)
  3051. #
  3052. # # 连通域分析
  3053. # num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(bin_clo, connectivity=8)
  3054. #
  3055. # # 查看各个返回值
  3056. # # 连通域数量
  3057. # print('num_labels = ',num_labels)
  3058. # # 连通域的信息:对应各个轮廓的x、y、width、height和面积
  3059. # print('stats = ',stats)
  3060. # # 连通域的中心点
  3061. # print('centroids = ',centroids)
  3062. # # 每一个像素的标签1、2、3.。。,同一个连通域的标签是一致的
  3063. # print('labels = ',labels)
  3064. #
  3065. # # 不同的连通域赋予不同的颜色
  3066. # output = np.zeros((img.shape[0], img.shape[1], 3), np.uint8)
  3067. # for i in range(1, num_labels):
  3068. #
  3069. # mask = labels == i
  3070. # output[:, :, 0][mask] = np.random.randint(0, 255)
  3071. # output[:, :, 1][mask] = np.random.randint(0, 255)
  3072. # output[:, :, 2][mask] = np.random.randint(0, 255)
  3073. # cv2.imshow('oginal', output)
  3074. # cv2.waitKey()
  3075. # cv2.destroyAllWindows()
  3076. def get_points_col(points, split_y, threshold=5):
  3077. # 坐标点按行分
  3078. row_point_list = []
  3079. row_point = []
  3080. points.sort(key=lambda x: (x[0], x[1]))
  3081. # print("get_points_col points sort", points)
  3082. x = points[0][0]
  3083. for i in range(1, len(split_y)):
  3084. for p in points:
  3085. if p[1] <= split_y[i-1] or p[1] >= split_y[i]:
  3086. continue
  3087. if x-threshold <= p[0] <= x+threshold:
  3088. row_point.append(p)
  3089. else:
  3090. # print("row_point", row_point)
  3091. row_point.sort(key=lambda x: (x[1], x[0]))
  3092. if row_point:
  3093. row_point_list.append(row_point)
  3094. row_point = []
  3095. x = p[0]
  3096. row_point.append(p)
  3097. if row_point:
  3098. row_point_list.append(row_point)
  3099. return row_point_list
  3100. def get_points_row(points, split_y, threshold=5):
  3101. # 坐标点按列分
  3102. col_point_list = []
  3103. col_point = []
  3104. points.sort(key=lambda x: (x[1], x[0]))
  3105. y = points[0][1]
  3106. for i in range(len(split_y)):
  3107. for p in points:
  3108. if p[1] <= split_y[i-1] or p[1] >= split_y[i]:
  3109. continue
  3110. if y-threshold <= p[1] <= y+threshold:
  3111. col_point.append(p)
  3112. else:
  3113. col_point.sort(key=lambda x: (x[0], x[1]))
  3114. if col_point:
  3115. col_point_list.append(col_point)
  3116. col_point = []
  3117. y = p[1]
  3118. col_point.append(p)
  3119. if col_point:
  3120. col_point_list.append(col_point)
  3121. return col_point_list
  3122. def get_outline_point(points, split_y):
  3123. # 分割线纵坐标
  3124. # print("get_outline_point split_y", split_y)
  3125. if len(split_y) < 2:
  3126. return []
  3127. outline_2point = []
  3128. points.sort(key=lambda x: (x[1], x[0]))
  3129. for i in range(1, len(split_y)):
  3130. area_points = []
  3131. for point in points:
  3132. if point[1] <= split_y[i-1] or point[1] >= split_y[i]:
  3133. continue
  3134. area_points.append(point)
  3135. if area_points:
  3136. area_points.sort(key=lambda x: (x[1], x[0]))
  3137. outline_2point.append([area_points[0], area_points[-1]])
  3138. return outline_2point
  3139. # def merge_row(row_lines):
  3140. # for row in row_lines:
  3141. # for row1 in row_lines:
  3142. def get_best_predict_size(image_np):
  3143. sizes = [1280, 1152, 1024, 896, 768, 640, 512, 384, 256, 128]
  3144. min_len = 10000
  3145. best_height = sizes[0]
  3146. for height in sizes:
  3147. if abs(image_np.shape[0] - height) < min_len:
  3148. min_len = abs(image_np.shape[0] - height)
  3149. best_height = height
  3150. min_len = 10000
  3151. best_width = sizes[0]
  3152. for width in sizes:
  3153. if abs(image_np.shape[1] - width) < min_len:
  3154. min_len = abs(image_np.shape[1] - width)
  3155. best_width = width
  3156. return best_height, best_width
  3157. def choose_longer_row(lines):
  3158. new_row = []
  3159. jump_row = []
  3160. for i in range(len(lines)):
  3161. row1 = lines[i]
  3162. jump_flag = 0
  3163. if row1 in jump_row:
  3164. continue
  3165. for j in range(i+1, len(lines)):
  3166. row2 = lines[j]
  3167. if row2 in jump_row:
  3168. continue
  3169. if row2[1]-5 <= row1[1] <= row2[1]+5:
  3170. if row1[0] <= row2[0] and row1[2] >= row2[2]:
  3171. new_row.append(row1)
  3172. jump_row.append(row1)
  3173. jump_row.append(row2)
  3174. jump_flag = 1
  3175. break
  3176. elif row2[0] <= row1[0] and row2[2] >= row1[2]:
  3177. new_row.append(row2)
  3178. jump_row.append(row1)
  3179. jump_row.append(row2)
  3180. jump_flag = 1
  3181. break
  3182. if not jump_flag:
  3183. new_row.append(row1)
  3184. jump_row.append(row1)
  3185. return new_row
  3186. def choose_longer_col(lines):
  3187. new_col = []
  3188. jump_col = []
  3189. for i in range(len(lines)):
  3190. col1 = lines[i]
  3191. jump_flag = 0
  3192. if col1 in jump_col:
  3193. continue
  3194. for j in range(i+1, len(lines)):
  3195. col2 = lines[j]
  3196. if col2 in jump_col:
  3197. continue
  3198. if col2[0]-5 <= col1[0] <= col2[0]+5:
  3199. if col1[1] <= col2[1] and col1[3] >= col2[3]:
  3200. new_col.append(col1)
  3201. jump_col.append(col1)
  3202. jump_col.append(col2)
  3203. jump_flag = 1
  3204. break
  3205. elif col2[1] <= col1[1] and col2[3] >= col1[3]:
  3206. new_col.append(col2)
  3207. jump_col.append(col1)
  3208. jump_col.append(col2)
  3209. jump_flag = 1
  3210. break
  3211. if not jump_flag:
  3212. new_col.append(col1)
  3213. jump_col.append(col1)
  3214. return new_col
  3215. def delete_contain_bbox(bboxes):
  3216. # bbox互相包含,取小的bbox
  3217. delete_bbox = []
  3218. for i in range(len(bboxes)):
  3219. for j in range(i+1, len(bboxes)):
  3220. bbox1 = bboxes[i]
  3221. bbox2 = bboxes[j]
  3222. # 横坐标相等情况
  3223. if bbox1[0][0] == bbox2[0][0] and bbox1[1][0] == bbox2[1][0]:
  3224. if bbox1[0][1] <= bbox2[0][1] <= bbox2[1][1] <= bbox1[1][1]:
  3225. # print("1", bbox1, bbox2)
  3226. delete_bbox.append(bbox1)
  3227. elif bbox2[0][1] <= bbox1[0][1] <= bbox1[1][1] <= bbox2[1][1]:
  3228. # print("2", bbox1, bbox2)
  3229. delete_bbox.append(bbox2)
  3230. # 纵坐标相等情况
  3231. elif bbox1[0][1] == bbox2[0][1] and bbox1[1][1] == bbox2[1][1]:
  3232. if bbox1[0][0] <= bbox2[0][0] <= bbox2[1][0] <= bbox1[1][0]:
  3233. print("3", bbox1, bbox2)
  3234. delete_bbox.append(bbox1)
  3235. elif bbox2[0][0] <= bbox1[0][0] <= bbox1[1][0] <= bbox2[1][0]:
  3236. print("4", bbox1, bbox2)
  3237. delete_bbox.append(bbox2)
  3238. print("delete_contain_bbox len(bboxes)", len(bboxes))
  3239. print("delete_contain_bbox len(delete_bbox)", len(delete_bbox))
  3240. for bbox in delete_bbox:
  3241. if bbox in bboxes:
  3242. bboxes.remove(bbox)
  3243. print("delete_contain_bbox len(bboxes)", len(bboxes))
  3244. return bboxes
  3245. if __name__ == '__main__':
  3246. print()
  3247. path = "C:\\Users\\Administrator\\Downloads\\pred.pk"
  3248. with open(path, "rb") as f:
  3249. pred = pickle.load(f)
  3250. draw_pixel(pred, 0.2, True)
  3251. points2lines(pred, False, 0.5)