table_line.py 140 KB

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