Nevar pievienot vairāk kā 25 tēmas Tēmai ir jāsākas ar burtu vai ciparu, tā var saturēt domu zīmes ('-') un var būt līdz 35 simboliem gara.

pdf_parser.py 62KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650
  1. # -*- coding: utf-8 -*-
  2. import fitz
  3. import xgboost as xgb
  4. from io import BytesIO
  5. import torch
  6. import re
  7. import pdfplumber
  8. import logging
  9. from PIL import Image
  10. import numpy as np
  11. from rag.nlp import huqie
  12. from collections import Counter
  13. from copy import deepcopy
  14. from rag.cv.table_recognize import TableTransformer
  15. from rag.cv.ppdetection import PPDet
  16. from huggingface_hub import hf_hub_download
  17. logging.getLogger("pdfminer").setLevel(logging.WARNING)
  18. class HuParser:
  19. def __init__(self):
  20. from paddleocr import PaddleOCR
  21. logging.getLogger("ppocr").setLevel(logging.ERROR)
  22. self.ocr = PaddleOCR(use_angle_cls=False, lang="ch")
  23. self.layouter = PPDet()
  24. self.tbl_det = TableTransformer()
  25. self.updown_cnt_mdl = xgb.Booster()
  26. if torch.cuda.is_available():
  27. self.updown_cnt_mdl.set_param({"device": "cuda"})
  28. self.updown_cnt_mdl.load_model(hf_hub_download(repo_id="InfiniFlow/text_concat_xgb_v1.0",
  29. filename="updown_concat_xgb.model"))
  30. """
  31. If you have trouble downloading HuggingFace models, -_^ this might help!!
  32. For Linux:
  33. export HF_ENDPOINT=https://hf-mirror.com
  34. For Windows:
  35. Good luck
  36. ^_-
  37. """
  38. def __char_width(self, c):
  39. return (c["x1"] - c["x0"]) // len(c["text"])
  40. def __height(self, c):
  41. return c["bottom"] - c["top"]
  42. def _x_dis(self, a, b):
  43. return min(abs(a["x1"] - b["x0"]), abs(a["x0"] - b["x1"]),
  44. abs(a["x0"] + a["x1"] - b["x0"] - b["x1"]) / 2)
  45. def _y_dis(
  46. self, a, b):
  47. return (
  48. b["top"] + b["bottom"] - a["top"] - a["bottom"]) / 2
  49. def _match_proj(self, b):
  50. proj_patt = [
  51. r"第[零一二三四五六七八九十百]+章",
  52. r"第[零一二三四五六七八九十百]+[条节]",
  53. r"[零一二三四五六七八九十百]+[、是  ]",
  54. r"[\((][零一二三四五六七八九十百]+[)\)]",
  55. r"[\((][0-9]+[)\)]",
  56. r"[0-9]+(、|\.[  ]|)|\.[^0-9./a-zA-Z_%><-]{4,})",
  57. r"[0-9]+\.[0-9.]+(、|\.[  ])",
  58. r"[⚫•➢①② ]",
  59. ]
  60. return any([re.match(p, b["text"]) for p in proj_patt])
  61. def _updown_concat_features(self, up, down):
  62. w = max(self.__char_width(up), self.__char_width(down))
  63. h = max(self.__height(up), self.__height(down))
  64. y_dis = self._y_dis(up, down)
  65. LEN = 6
  66. tks_down = huqie.qie(down["text"][:LEN]).split(" ")
  67. tks_up = huqie.qie(up["text"][-LEN:]).split(" ")
  68. tks_all = up["text"][-LEN:].strip() \
  69. + (" " if re.match(r"[a-zA-Z0-9]+",
  70. up["text"][-1] + down["text"][0]) else "") \
  71. + down["text"][:LEN].strip()
  72. tks_all = huqie.qie(tks_all).split(" ")
  73. fea = [
  74. up.get("R", -1) == down.get("R", -1),
  75. y_dis / h,
  76. down["page_number"] - up["page_number"],
  77. up["layout_type"] == down["layout_type"],
  78. up["layout_type"] == "text",
  79. down["layout_type"] == "text",
  80. up["layout_type"] == "table",
  81. down["layout_type"] == "table",
  82. True if re.search(
  83. r"([。?!;!?;+))]|[a-z]\.)$",
  84. up["text"]) else False,
  85. True if re.search(r"[,:‘“、0-9(+-]$", up["text"]) else False,
  86. True if re.search(
  87. r"(^.?[/,?;:\],。;:’”?!》】)-])",
  88. down["text"]) else False,
  89. True if re.match(r"[\((][^\(\)()]+[)\)]$", up["text"]) else False,
  90. True if re.search(r"[,,][^。.]+$", up["text"]) else False,
  91. True if re.search(r"[,,][^。.]+$", up["text"]) else False,
  92. True if re.search(r"[\((][^\))]+$", up["text"])
  93. and re.search(r"[\))]", down["text"]) else False,
  94. self._match_proj(down),
  95. True if re.match(r"[A-Z]", down["text"]) else False,
  96. True if re.match(r"[A-Z]", up["text"][-1]) else False,
  97. True if re.match(r"[a-z0-9]", up["text"][-1]) else False,
  98. True if re.match(r"[0-9.%,-]+$", down["text"]) else False,
  99. up["text"].strip()[-2:] == down["text"].strip()[-2:] if len(up["text"].strip()
  100. ) > 1 and len(
  101. down["text"].strip()) > 1 else False,
  102. up["x0"] > down["x1"],
  103. abs(self.__height(up) - self.__height(down)) / min(self.__height(up),
  104. self.__height(down)),
  105. self._x_dis(up, down) / max(w, 0.000001),
  106. (len(up["text"]) - len(down["text"])) /
  107. max(len(up["text"]), len(down["text"])),
  108. len(tks_all) - len(tks_up) - len(tks_down),
  109. len(tks_down) - len(tks_up),
  110. tks_down[-1] == tks_up[-1],
  111. max(down["in_row"], up["in_row"]),
  112. abs(down["in_row"] - up["in_row"]),
  113. len(tks_down) == 1 and huqie.tag(tks_down[0]).find("n") >= 0,
  114. len(tks_up) == 1 and huqie.tag(tks_up[0]).find("n") >= 0
  115. ]
  116. return fea
  117. @staticmethod
  118. def sort_Y_firstly(arr, threashold):
  119. # sort using y1 first and then x1
  120. arr = sorted(arr, key=lambda r: (r["top"], r["x0"]))
  121. for i in range(len(arr) - 1):
  122. for j in range(i, -1, -1):
  123. # restore the order using th
  124. if abs(arr[j + 1]["top"] - arr[j]["top"]) < threashold \
  125. and arr[j + 1]["x0"] < arr[j]["x0"]:
  126. tmp = deepcopy(arr[j])
  127. arr[j] = deepcopy(arr[j + 1])
  128. arr[j + 1] = deepcopy(tmp)
  129. return arr
  130. @staticmethod
  131. def sort_R_firstly(arr, thr=0):
  132. # sort using y1 first and then x1
  133. # sorted(arr, key=lambda r: (r["top"], r["x0"]))
  134. arr = HuParser.sort_Y_firstly(arr, thr)
  135. for i in range(len(arr) - 1):
  136. for j in range(i, -1, -1):
  137. if "R" not in arr[j] or "R" not in arr[j + 1]:
  138. continue
  139. if arr[j + 1]["R"] < arr[j]["R"] \
  140. or (
  141. arr[j + 1]["R"] == arr[j]["R"]
  142. and arr[j + 1]["x0"] < arr[j]["x0"]
  143. ):
  144. tmp = arr[j]
  145. arr[j] = arr[j + 1]
  146. arr[j + 1] = tmp
  147. return arr
  148. @staticmethod
  149. def sort_X_firstly(arr, threashold, copy=True):
  150. # sort using y1 first and then x1
  151. arr = sorted(arr, key=lambda r: (r["x0"], r["top"]))
  152. for i in range(len(arr) - 1):
  153. for j in range(i, -1, -1):
  154. # restore the order using th
  155. if abs(arr[j + 1]["x0"] - arr[j]["x0"]) < threashold \
  156. and arr[j + 1]["top"] < arr[j]["top"]:
  157. tmp = deepcopy(arr[j]) if copy else arr[j]
  158. arr[j] = deepcopy(arr[j + 1]) if copy else arr[j + 1]
  159. arr[j + 1] = deepcopy(tmp) if copy else tmp
  160. return arr
  161. @staticmethod
  162. def sort_C_firstly(arr, thr=0):
  163. # sort using y1 first and then x1
  164. # sorted(arr, key=lambda r: (r["x0"], r["top"]))
  165. arr = HuParser.sort_X_firstly(arr, thr)
  166. for i in range(len(arr) - 1):
  167. for j in range(i, -1, -1):
  168. # restore the order using th
  169. if "C" not in arr[j] or "C" not in arr[j + 1]:
  170. continue
  171. if arr[j + 1]["C"] < arr[j]["C"] \
  172. or (
  173. arr[j + 1]["C"] == arr[j]["C"]
  174. and arr[j + 1]["top"] < arr[j]["top"]
  175. ):
  176. tmp = arr[j]
  177. arr[j] = arr[j + 1]
  178. arr[j + 1] = tmp
  179. return arr
  180. return sorted(arr, key=lambda r: (r.get("C", r["x0"]), r["top"]))
  181. def _has_color(self, o):
  182. if o.get("ncs", "") == "DeviceGray":
  183. if o["stroking_color"] and o["stroking_color"][0] == 1 and o["non_stroking_color"] and \
  184. o["non_stroking_color"][0] == 1:
  185. if re.match(r"[a-zT_\[\]\(\)-]+", o.get("text", "")):
  186. return False
  187. return True
  188. def __overlapped_area(self, a, b, ratio=True):
  189. tp, btm, x0, x1 = a["top"], a["bottom"], a["x0"], a["x1"]
  190. if b["x0"] > x1 or b["x1"] < x0:
  191. return 0
  192. if b["bottom"] < tp or b["top"] > btm:
  193. return 0
  194. x0_ = max(b["x0"], x0)
  195. x1_ = min(b["x1"], x1)
  196. assert x0_ <= x1_, "Fuckedup! T:{},B:{},X0:{},X1:{} ==> {}".format(
  197. tp, btm, x0, x1, b)
  198. tp_ = max(b["top"], tp)
  199. btm_ = min(b["bottom"], btm)
  200. assert tp_ <= btm_, "Fuckedup! T:{},B:{},X0:{},X1:{} => {}".format(
  201. tp, btm, x0, x1, b)
  202. ov = (btm_ - tp_) * (x1_ - x0_) if x1 - \
  203. x0 != 0 and btm - tp != 0 else 0
  204. if ov > 0 and ratio:
  205. ov /= (x1 - x0) * (btm - tp)
  206. return ov
  207. def __find_overlapped_with_threashold(self, box, boxes, thr=0.3):
  208. if not boxes:
  209. return
  210. max_overlaped_i, max_overlaped, _max_overlaped = None, thr, 0
  211. s, e = 0, len(boxes)
  212. for i in range(s, e):
  213. ov = self.__overlapped_area(box, boxes[i])
  214. _ov = self.__overlapped_area(boxes[i], box)
  215. if (ov, _ov) < (max_overlaped, _max_overlaped):
  216. continue
  217. max_overlaped_i = i
  218. max_overlaped = ov
  219. _max_overlaped = _ov
  220. return max_overlaped_i
  221. def __find_overlapped(self, box, boxes_sorted_by_y, naive=False):
  222. if not boxes_sorted_by_y:
  223. return
  224. bxs = boxes_sorted_by_y
  225. s, e, ii = 0, len(bxs), 0
  226. while s < e and not naive:
  227. ii = (e + s) // 2
  228. pv = bxs[ii]
  229. if box["bottom"] < pv["top"]:
  230. e = ii
  231. continue
  232. if box["top"] > pv["bottom"]:
  233. s = ii + 1
  234. continue
  235. break
  236. while s < ii:
  237. if box["top"] > bxs[s]["bottom"]:
  238. s += 1
  239. break
  240. while e - 1 > ii:
  241. if box["bottom"] < bxs[e - 1]["top"]:
  242. e -= 1
  243. break
  244. max_overlaped_i, max_overlaped = None, 0
  245. for i in range(s, e):
  246. ov = self.__overlapped_area(bxs[i], box)
  247. if ov <= max_overlaped:
  248. continue
  249. max_overlaped_i = i
  250. max_overlaped = ov
  251. return max_overlaped_i
  252. def _is_garbage(self, b):
  253. patt = [r"^•+$", r"(版权归©|免责条款|地址[::])", r"\.{3,}", "^[0-9]{1,2} / ?[0-9]{1,2}$",
  254. r"^[0-9]{1,2} of [0-9]{1,2}$", "^http://[^ ]{12,}",
  255. "(资料|数据)来源[::]", "[0-9a-z._-]+@[a-z0-9-]+\\.[a-z]{2,3}",
  256. "\\(cid *: *[0-9]+ *\\)"
  257. ]
  258. return any([re.search(p, b["text"]) for p in patt])
  259. def __layouts_cleanup(self, boxes, layouts, far=2, thr=0.7):
  260. def notOverlapped(a, b):
  261. return any([a["x1"] < b["x0"],
  262. a["x0"] > b["x1"],
  263. a["bottom"] < b["top"],
  264. a["top"] > b["bottom"]])
  265. i = 0
  266. while i + 1 < len(layouts):
  267. j = i + 1
  268. while j < min(i + far, len(layouts)) \
  269. and (layouts[i].get("type", "") != layouts[j].get("type", "")
  270. or notOverlapped(layouts[i], layouts[j])):
  271. j += 1
  272. if j >= min(i + far, len(layouts)):
  273. i += 1
  274. continue
  275. if self.__overlapped_area(layouts[i], layouts[j]) < thr \
  276. and self.__overlapped_area(layouts[j], layouts[i]) < thr:
  277. i += 1
  278. continue
  279. if layouts[i].get("score") and layouts[j].get("score"):
  280. if layouts[i]["score"] > layouts[j]["score"]:
  281. layouts.pop(j)
  282. else:
  283. layouts.pop(i)
  284. continue
  285. area_i, area_i_1 = 0, 0
  286. for b in boxes:
  287. if not notOverlapped(b, layouts[i]):
  288. area_i += self.__overlapped_area(b, layouts[i], False)
  289. if not notOverlapped(b, layouts[j]):
  290. area_i_1 += self.__overlapped_area(b, layouts[j], False)
  291. if area_i > area_i_1:
  292. layouts.pop(j)
  293. else:
  294. layouts.pop(i)
  295. return layouts
  296. def __table_paddle(self, images):
  297. tbls = self.tbl_det([img for img in images], threshold=0.5)
  298. res = []
  299. # align left&right for rows, align top&bottom for columns
  300. for tbl in tbls:
  301. lts = [{"label": b["type"],
  302. "score": b["score"],
  303. "x0": b["bbox"][0], "x1": b["bbox"][2],
  304. "top": b["bbox"][1], "bottom": b["bbox"][-1]
  305. } for b in tbl]
  306. if not lts:
  307. continue
  308. left = [b["x0"] for b in lts if b["label"].find(
  309. "row") > 0 or b["label"].find("header") > 0]
  310. right = [b["x1"] for b in lts if b["label"].find(
  311. "row") > 0 or b["label"].find("header") > 0]
  312. if not left:
  313. continue
  314. left = np.median(left) if len(left) > 4 else np.min(left)
  315. right = np.median(right) if len(right) > 4 else np.max(right)
  316. for b in lts:
  317. if b["label"].find("row") > 0 or b["label"].find("header") > 0:
  318. if b["x0"] > left:
  319. b["x0"] = left
  320. if b["x1"] < right:
  321. b["x1"] = right
  322. top = [b["top"] for b in lts if b["label"] == "table column"]
  323. bottom = [b["bottom"] for b in lts if b["label"] == "table column"]
  324. if not top:
  325. res.append(lts)
  326. continue
  327. top = np.median(top) if len(top) > 4 else np.min(top)
  328. bottom = np.median(bottom) if len(bottom) > 4 else np.max(bottom)
  329. for b in lts:
  330. if b["label"] == "table column":
  331. if b["top"] > top:
  332. b["top"] = top
  333. if b["bottom"] < bottom:
  334. b["bottom"] = bottom
  335. res.append(lts)
  336. return res
  337. def __table_transformer_job(self, ZM):
  338. logging.info("Table processing...")
  339. imgs, pos = [], []
  340. tbcnt = [0]
  341. MARGIN = 10
  342. self.tb_cpns = []
  343. assert len(self.page_layout) == len(self.page_images)
  344. for p, tbls in enumerate(self.page_layout): # for page
  345. tbls = [f for f in tbls if f["type"] == "table"]
  346. tbcnt.append(len(tbls))
  347. if not tbls:
  348. continue
  349. for tb in tbls: # for table
  350. left, top, right, bott = tb["x0"] - MARGIN, tb["top"] - MARGIN, \
  351. tb["x1"] + MARGIN, tb["bottom"] + MARGIN
  352. left *= ZM
  353. top *= ZM
  354. right *= ZM
  355. bott *= ZM
  356. pos.append((left, top))
  357. imgs.append(self.page_images[p].crop((left, top, right, bott)))
  358. assert len(self.page_images) == len(tbcnt) - 1
  359. if not imgs:
  360. return
  361. recos = self.__table_paddle(imgs)
  362. tbcnt = np.cumsum(tbcnt)
  363. for i in range(len(tbcnt) - 1): # for page
  364. pg = []
  365. for j, tb_items in enumerate(
  366. recos[tbcnt[i]: tbcnt[i + 1]]): # for table
  367. poss = pos[tbcnt[i]: tbcnt[i + 1]]
  368. for it in tb_items: # for table components
  369. it["x0"] = (it["x0"] + poss[j][0])
  370. it["x1"] = (it["x1"] + poss[j][0])
  371. it["top"] = (it["top"] + poss[j][1])
  372. it["bottom"] = (it["bottom"] + poss[j][1])
  373. for n in ["x0", "x1", "top", "bottom"]:
  374. it[n] /= ZM
  375. it["top"] += self.page_cum_height[i]
  376. it["bottom"] += self.page_cum_height[i]
  377. it["pn"] = i
  378. it["layoutno"] = j
  379. pg.append(it)
  380. self.tb_cpns.extend(pg)
  381. def __ocr_paddle(self, pagenum, img, chars, ZM=3):
  382. bxs = self.ocr.ocr(np.array(img), cls=True)[0]
  383. if not bxs:
  384. self.boxes.append([])
  385. return
  386. bxs = [(line[0], line[1][0]) for line in bxs]
  387. bxs = self.sort_Y_firstly(
  388. [{"x0": b[0][0] / ZM, "x1": b[1][0] / ZM,
  389. "top": b[0][1] / ZM, "text": "", "txt": t,
  390. "bottom": b[-1][1] / ZM,
  391. "page_number": pagenum} for b, t in bxs if b[0][0] <= b[1][0] and b[0][1] <= b[-1][1]],
  392. self.mean_height[-1] / 3
  393. )
  394. # merge chars in the same rect
  395. for c in self.sort_X_firstly(chars, self.mean_width[pagenum - 1] // 4):
  396. ii = self.__find_overlapped(c, bxs)
  397. if ii is None:
  398. self.lefted_chars.append(c)
  399. continue
  400. ch = c["bottom"] - c["top"]
  401. bh = bxs[ii]["bottom"] - bxs[ii]["top"]
  402. if abs(ch - bh) / max(ch, bh) >= 0.7:
  403. self.lefted_chars.append(c)
  404. continue
  405. bxs[ii]["text"] += c["text"]
  406. for b in bxs:
  407. if not b["text"]:
  408. b["text"] = b["txt"]
  409. del b["txt"]
  410. if self.mean_height[-1] == 0:
  411. self.mean_height[-1] = np.median([b["bottom"] - b["top"]
  412. for b in bxs])
  413. self.boxes.append(bxs)
  414. def __layouts_paddle(self, ZM):
  415. assert len(self.page_images) == len(self.boxes)
  416. # Tag layout type
  417. boxes = []
  418. layouts = self.layouter([np.array(img) for img in self.page_images])
  419. assert len(self.page_images) == len(layouts)
  420. for pn, lts in enumerate(layouts):
  421. bxs = self.boxes[pn]
  422. lts = [{"type": b["type"],
  423. "score": float(b["score"]),
  424. "x0": b["bbox"][0] / ZM, "x1": b["bbox"][2] / ZM,
  425. "top": b["bbox"][1] / ZM, "bottom": b["bbox"][-1] / ZM,
  426. "page_number": pn,
  427. } for b in lts]
  428. lts = self.sort_Y_firstly(lts, self.mean_height[pn] / 2)
  429. lts = self.__layouts_cleanup(bxs, lts)
  430. self.page_layout.append(lts)
  431. # Tag layout type, layouts are ready
  432. def findLayout(ty):
  433. nonlocal bxs, lts
  434. lts_ = [lt for lt in lts if lt["type"] == ty]
  435. i = 0
  436. while i < len(bxs):
  437. if bxs[i].get("layout_type"):
  438. i += 1
  439. continue
  440. if self._is_garbage(bxs[i]):
  441. logging.debug("GARBAGE: " + bxs[i]["text"])
  442. bxs.pop(i)
  443. continue
  444. ii = self.__find_overlapped_with_threashold(bxs[i], lts_,
  445. thr=0.4)
  446. if ii is None: # belong to nothing
  447. bxs[i]["layout_type"] = ""
  448. i += 1
  449. continue
  450. lts_[ii]["visited"] = True
  451. if lts_[ii]["type"] in ["footer", "header", "reference"]:
  452. if lts_[ii]["type"] not in self.garbages:
  453. self.garbages[lts_[ii]["type"]] = []
  454. self.garbages[lts_[ii]["type"]].append(bxs[i]["text"])
  455. logging.debug("GARBAGE: " + bxs[i]["text"])
  456. bxs.pop(i)
  457. continue
  458. bxs[i]["layoutno"] = f"{ty}-{ii}"
  459. bxs[i]["layout_type"] = lts_[ii]["type"]
  460. i += 1
  461. for lt in ["footer", "header", "reference", "figure caption",
  462. "table caption", "title", "text", "table", "figure"]:
  463. findLayout(lt)
  464. # add box to figure layouts which has not text box
  465. for i, lt in enumerate(
  466. [lt for lt in lts if lt["type"] == "figure"]):
  467. if lt.get("visited"):
  468. continue
  469. lt = deepcopy(lt)
  470. del lt["type"]
  471. lt["text"] = ""
  472. lt["layout_type"] = "figure"
  473. lt["layoutno"] = f"figure-{i}"
  474. bxs.append(lt)
  475. boxes.extend(bxs)
  476. self.boxes = boxes
  477. def __text_merge(self, garbage):
  478. # merge adjusted boxes
  479. bxs = self.boxes
  480. def end_with(b, txt):
  481. txt = txt.strip()
  482. tt = b.get("text", "").strip()
  483. return tt and tt.find(txt) == len(tt) - len(txt)
  484. def start_with(b, txts):
  485. tt = b.get("text", "").strip()
  486. return tt and any([tt.find(t.strip()) == 0 for t in txts])
  487. i = 0
  488. while i < len(bxs) - 1:
  489. b = bxs[i]
  490. b_ = bxs[i + 1]
  491. if b.get("layoutno", "0") != b_.get("layoutno", "1"):
  492. i += 1
  493. continue
  494. dis_thr = 1
  495. dis = b["x1"] - b_["x0"]
  496. if b.get("layout_type", "") != "text" or b_.get(
  497. "layout_type", "") != "text":
  498. if end_with(b, ",") or start_with(b_, "(,"):
  499. dis_thr = -8
  500. else:
  501. i += 1
  502. continue
  503. if abs(self._y_dis(b, b_)) < self.mean_height[bxs[i]["page_number"] - 1] / 5 \
  504. and dis >= dis_thr and b["x1"] < b_["x1"]:
  505. # merge
  506. bxs[i]["x1"] = b_["x1"]
  507. bxs[i]["top"] = (b["top"] + b_["top"]) / 2
  508. bxs[i]["bottom"] = (b["bottom"] + b_["bottom"]) / 2
  509. bxs[i]["text"] += b_["text"]
  510. bxs.pop(i + 1)
  511. continue
  512. i += 1
  513. self.boxes = bxs
  514. # count boxes in the same row
  515. for i in range(len(self.boxes)):
  516. mh = self.mean_height[self.boxes[i]["page_number"] - 1]
  517. self.boxes[i]["in_row"] = 0
  518. j = max(0, i - 12)
  519. while j < min(i + 12, len(self.boxes)):
  520. if j == i:
  521. j += 1
  522. continue
  523. ydis = self._y_dis(self.boxes[i], self.boxes[j]) / mh
  524. if abs(ydis) < 1:
  525. self.boxes[i]["in_row"] += 1
  526. elif ydis > 0:
  527. break
  528. j += 1
  529. def gather(kwd, fzy=10, ption=0.6):
  530. eles = self.sort_Y_firstly(
  531. [r for r in self.tb_cpns if re.match(kwd, r["label"])], fzy)
  532. eles = self.__layouts_cleanup(self.boxes, eles, 5, ption)
  533. return self.sort_Y_firstly(eles, 0)
  534. headers = gather(r".*header$")
  535. rows = gather(r".* (row|header)")
  536. spans = gather(r".*spanning")
  537. clmns = sorted([r for r in self.tb_cpns if re.match(
  538. r"table column$", r["label"])], key=lambda x: (x["pn"], x["layoutno"], x["x0"]))
  539. clmns = self.__layouts_cleanup(self.boxes, clmns, 5, 0.5)
  540. for b in self.boxes:
  541. if b.get("layout_type", "") != "table":
  542. continue
  543. ii = self.__find_overlapped_with_threashold(b, rows, thr=0.3)
  544. if ii is not None:
  545. b["R"] = ii
  546. b["R_top"] = rows[ii]["top"]
  547. b["R_bott"] = rows[ii]["bottom"]
  548. ii = self.__find_overlapped_with_threashold(b, headers, thr=0.3)
  549. if ii is not None:
  550. b["H_top"] = headers[ii]["top"]
  551. b["H_bott"] = headers[ii]["bottom"]
  552. b["H_left"] = headers[ii]["x0"]
  553. b["H_right"] = headers[ii]["x1"]
  554. b["H"] = ii
  555. ii = self.__find_overlapped_with_threashold(b, clmns, thr=0.3)
  556. if ii is not None:
  557. b["C"] = ii
  558. b["C_left"] = clmns[ii]["x0"]
  559. b["C_right"] = clmns[ii]["x1"]
  560. ii = self.__find_overlapped_with_threashold(b, spans, thr=0.3)
  561. if ii is not None:
  562. b["H_top"] = spans[ii]["top"]
  563. b["H_bott"] = spans[ii]["bottom"]
  564. b["H_left"] = spans[ii]["x0"]
  565. b["H_right"] = spans[ii]["x1"]
  566. b["SP"] = ii
  567. # concat between rows
  568. boxes = deepcopy(self.boxes)
  569. blocks = []
  570. while boxes:
  571. chunks = []
  572. def dfs(up, dp):
  573. if not up["text"].strip() or up["text"].strip() in garbage:
  574. return
  575. chunks.append(up)
  576. i = dp
  577. while i < min(dp + 12, len(boxes)):
  578. ydis = self._y_dis(up, boxes[i])
  579. smpg = up["page_number"] == boxes[i]["page_number"]
  580. mh = self.mean_height[up["page_number"] - 1]
  581. mw = self.mean_width[up["page_number"] - 1]
  582. if smpg and ydis > mh * 4:
  583. break
  584. if not smpg and ydis > mh * 16:
  585. break
  586. down = boxes[i]
  587. if up.get("R", "") != down.get(
  588. "R", "") and up["text"][-1] != ",":
  589. i += 1
  590. continue
  591. if re.match(r"[0-9]{2,3}/[0-9]{3}$", up["text"]) \
  592. or re.match(r"[0-9]{2,3}/[0-9]{3}$", down["text"]):
  593. i += 1
  594. continue
  595. if not down["text"].strip() \
  596. or down["text"].strip() in garbage:
  597. i += 1
  598. continue
  599. if up["x1"] < down["x0"] - 10 * \
  600. mw or up["x0"] > down["x1"] + 10 * mw:
  601. i += 1
  602. continue
  603. if i - dp < 5 and up.get("layout_type") == "text":
  604. if up.get("layoutno", "1") == down.get(
  605. "layoutno", "2"):
  606. dfs(down, i + 1)
  607. boxes.pop(i)
  608. return
  609. i += 1
  610. continue
  611. fea = self._updown_concat_features(up, down)
  612. if self.updown_cnt_mdl.predict(
  613. xgb.DMatrix([fea]))[0] <= 0.5:
  614. i += 1
  615. continue
  616. dfs(down, i + 1)
  617. boxes.pop(i)
  618. return
  619. dfs(boxes[0], 1)
  620. boxes.pop(0)
  621. if chunks:
  622. blocks.append(chunks)
  623. # concat within each block
  624. boxes = []
  625. for b in blocks:
  626. if len(b) == 1:
  627. boxes.append(b[0])
  628. continue
  629. t = b[0]
  630. for c in b[1:]:
  631. t["text"] = t["text"].strip()
  632. c["text"] = c["text"].strip()
  633. if not c["text"]:
  634. continue
  635. if t["text"] and re.match(
  636. r"[0-9\.a-zA-Z]+$", t["text"][-1] + c["text"][-1]):
  637. t["text"] += " "
  638. t["text"] += c["text"]
  639. t["x0"] = min(t["x0"], c["x0"])
  640. t["x1"] = max(t["x1"], c["x1"])
  641. t["page_number"] = min(t["page_number"], c["page_number"])
  642. t["bottom"] = c["bottom"]
  643. if not t["layout_type"] \
  644. and c["layout_type"]:
  645. t["layout_type"] = c["layout_type"]
  646. boxes.append(t)
  647. self.boxes = self.sort_Y_firstly(boxes, 0)
  648. def __filter_forpages(self):
  649. if not self.boxes:
  650. return
  651. to = min(7, len(self.page_images) // 5)
  652. pg_hits = [0 for _ in range(to)]
  653. def possible(c):
  654. if c.get("layout_type", "") == "reference":
  655. return True
  656. if c["bottom"] - c["top"] >= 2 * \
  657. self.mean_height[c["page_number"] - 1]:
  658. return False
  659. if c["text"].find("....") >= 0 \
  660. or (c["x1"] - c["x0"] > 250 and re.search(r"[0-9]+$",
  661. c["text"].strip())):
  662. return True
  663. return self.is_caption(c) and re.search(
  664. r"[0-9]+$", c["text"].strip())
  665. for c in self.boxes:
  666. if c["page_number"] >= to:
  667. break
  668. if possible(c):
  669. pg_hits[c["page_number"] - 1] += 1
  670. st, ed = -1, -1
  671. for i in range(len(self.boxes)):
  672. c = self.boxes[i]
  673. if c["page_number"] >= to:
  674. break
  675. if pg_hits[c["page_number"] - 1] >= 3 and possible(c):
  676. if st < 0:
  677. st = i
  678. else:
  679. ed = i
  680. for _ in range(st, ed + 1):
  681. self.boxes.pop(st)
  682. def _blockType(self, b):
  683. patt = [
  684. ("^(20|19)[0-9]{2}[年/-][0-9]{1,2}[月/-][0-9]{1,2}日*$", "Dt"),
  685. (r"^(20|19)[0-9]{2}年$", "Dt"),
  686. (r"^(20|19)[0-9]{2}[年-][0-9]{1,2}月*$", "Dt"),
  687. ("^[0-9]{1,2}[月-][0-9]{1,2}日*$", "Dt"),
  688. (r"^第*[一二三四1-4]季度$", "Dt"),
  689. (r"^(20|19)[0-9]{2}年*[一二三四1-4]季度$", "Dt"),
  690. (r"^(20|19)[0-9]{2}[ABCDE]$", "Dt"),
  691. ("^[0-9.,+%/ -]+$", "Nu"),
  692. (r"^[0-9A-Z/\._~-]+$", "Ca"),
  693. (r"^[A-Z]*[a-z' -]+$", "En"),
  694. (r"^[0-9.,+-]+[0-9A-Za-z/$¥%<>()()' -]+$", "NE"),
  695. (r"^.{1}$", "Sg")
  696. ]
  697. for p, n in patt:
  698. if re.search(p, b["text"].strip()):
  699. return n
  700. tks = [t for t in huqie.qie(b["text"]).split(" ") if len(t) > 1]
  701. if len(tks) > 3:
  702. if len(tks) < 12:
  703. return "Tx"
  704. else:
  705. return "Lx"
  706. if len(tks) == 1 and huqie.tag(tks[0]) == "nr":
  707. return "Nr"
  708. return "Ot"
  709. def __cal_spans(self, boxes, rows, cols, tbl, html=True):
  710. # caculate span
  711. clft = [np.mean([c.get("C_left", c["x0"]) for c in cln])
  712. for cln in cols]
  713. crgt = [np.mean([c.get("C_right", c["x1"]) for c in cln])
  714. for cln in cols]
  715. rtop = [np.mean([c.get("R_top", c["top"]) for c in row])
  716. for row in rows]
  717. rbtm = [np.mean([c.get("R_btm", c["bottom"])
  718. for c in row]) for row in rows]
  719. for b in boxes:
  720. if "SP" not in b:
  721. continue
  722. b["colspan"] = [b["cn"]]
  723. b["rowspan"] = [b["rn"]]
  724. # col span
  725. for j in range(0, len(clft)):
  726. if j == b["cn"]:
  727. continue
  728. if clft[j] + (crgt[j] - clft[j]) / 2 < b["H_left"]:
  729. continue
  730. if crgt[j] - (crgt[j] - clft[j]) / 2 > b["H_right"]:
  731. continue
  732. b["colspan"].append(j)
  733. # row span
  734. for j in range(0, len(rtop)):
  735. if j == b["rn"]:
  736. continue
  737. if rtop[j] + (rbtm[j] - rtop[j]) / 2 < b["H_top"]:
  738. continue
  739. if rbtm[j] - (rbtm[j] - rtop[j]) / 2 > b["H_bott"]:
  740. continue
  741. b["rowspan"].append(j)
  742. def join(arr):
  743. if not arr:
  744. return ""
  745. return "".join([t["text"] for t in arr])
  746. # rm the spaning cells
  747. for i in range(len(tbl)):
  748. for j, arr in enumerate(tbl[i]):
  749. if not arr:
  750. continue
  751. if all(["rowspan" not in a and "colspan" not in a for a in arr]):
  752. continue
  753. rowspan, colspan = [], []
  754. for a in arr:
  755. if isinstance(a.get("rowspan", 0), list):
  756. rowspan.extend(a["rowspan"])
  757. if isinstance(a.get("colspan", 0), list):
  758. colspan.extend(a["colspan"])
  759. rowspan, colspan = set(rowspan), set(colspan)
  760. if len(rowspan) < 2 and len(colspan) < 2:
  761. for a in arr:
  762. if "rowspan" in a:
  763. del a["rowspan"]
  764. if "colspan" in a:
  765. del a["colspan"]
  766. continue
  767. rowspan, colspan = sorted(rowspan), sorted(colspan)
  768. rowspan = list(range(rowspan[0], rowspan[-1] + 1))
  769. colspan = list(range(colspan[0], colspan[-1] + 1))
  770. assert i in rowspan, rowspan
  771. assert j in colspan, colspan
  772. arr = []
  773. for r in rowspan:
  774. for c in colspan:
  775. arr_txt = join(arr)
  776. if tbl[r][c] and join(tbl[r][c]) != arr_txt:
  777. arr.extend(tbl[r][c])
  778. tbl[r][c] = None if html else arr
  779. for a in arr:
  780. if len(rowspan) > 1:
  781. a["rowspan"] = len(rowspan)
  782. elif "rowspan" in a:
  783. del a["rowspan"]
  784. if len(colspan) > 1:
  785. a["colspan"] = len(colspan)
  786. elif "colspan" in a:
  787. del a["colspan"]
  788. tbl[rowspan[0]][colspan[0]] = arr
  789. return tbl
  790. def __construct_table(self, boxes, html=False):
  791. cap = ""
  792. i = 0
  793. while i < len(boxes):
  794. if self.is_caption(boxes[i]):
  795. cap += boxes[i]["text"]
  796. boxes.pop(i)
  797. i -= 1
  798. i += 1
  799. if not boxes:
  800. return []
  801. for b in boxes:
  802. b["btype"] = self._blockType(b)
  803. max_type = Counter([b["btype"] for b in boxes]).items()
  804. max_type = max(max_type, key=lambda x: x[1])[0] if max_type else ""
  805. logging.debug("MAXTYPE: " + max_type)
  806. rowh = [b["R_bott"] - b["R_top"] for b in boxes if "R" in b]
  807. rowh = np.min(rowh) if rowh else 0
  808. # boxes = self.sort_Y_firstly(boxes, rowh/5)
  809. boxes = self.sort_R_firstly(boxes, rowh / 2)
  810. boxes[0]["rn"] = 0
  811. rows = [[boxes[0]]]
  812. btm = boxes[0]["bottom"]
  813. for b in boxes[1:]:
  814. b["rn"] = len(rows) - 1
  815. lst_r = rows[-1]
  816. if lst_r[-1].get("R", "") != b.get("R", "") \
  817. or (b["top"] >= btm - 3 and lst_r[-1].get("R", "-1") != b.get("R", "-2")
  818. ): # new row
  819. btm = b["bottom"]
  820. b["rn"] += 1
  821. rows.append([b])
  822. continue
  823. btm = (btm + b["bottom"]) / 2.
  824. rows[-1].append(b)
  825. colwm = [b["C_right"] - b["C_left"] for b in boxes if "C" in b]
  826. colwm = np.min(colwm) if colwm else 0
  827. crosspage = len(set([b["page_number"] for b in boxes])) > 1
  828. if crosspage:
  829. boxes = self.sort_X_firstly(boxes, colwm / 2, False)
  830. else:
  831. boxes = self.sort_C_firstly(boxes, colwm / 2)
  832. boxes[0]["cn"] = 0
  833. cols = [[boxes[0]]]
  834. right = boxes[0]["x1"]
  835. for b in boxes[1:]:
  836. b["cn"] = len(cols) - 1
  837. lst_c = cols[-1]
  838. if (int(b.get("C", "1")) - int(lst_c[-1].get("C", "1")) == 1 and b["page_number"] == lst_c[-1][
  839. "page_number"]) \
  840. or (b["x0"] >= right and lst_c[-1].get("C", "-1") != b.get("C", "-2")): # new col
  841. right = b["x1"]
  842. b["cn"] += 1
  843. cols.append([b])
  844. continue
  845. right = (right + b["x1"]) / 2.
  846. cols[-1].append(b)
  847. tbl = [[[] for _ in range(len(cols))] for _ in range(len(rows))]
  848. for b in boxes:
  849. tbl[b["rn"]][b["cn"]].append(b)
  850. if len(rows) >= 4:
  851. # remove single in column
  852. j = 0
  853. while j < len(tbl[0]):
  854. e, ii = 0, 0
  855. for i in range(len(tbl)):
  856. if tbl[i][j]:
  857. e += 1
  858. ii = i
  859. if e > 1:
  860. break
  861. if e > 1:
  862. j += 1
  863. continue
  864. f = (j > 0 and tbl[ii][j - 1] and tbl[ii]
  865. [j - 1][0].get("text")) or j == 0
  866. ff = (j + 1 < len(tbl[ii]) and tbl[ii][j + 1] and tbl[ii]
  867. [j + 1][0].get("text")) or j + 1 >= len(tbl[ii])
  868. if f and ff:
  869. j += 1
  870. continue
  871. bx = tbl[ii][j][0]
  872. logging.debug("Relocate column single: " + bx["text"])
  873. # j column only has one value
  874. left, right = 100000, 100000
  875. if j > 0 and not f:
  876. for i in range(len(tbl)):
  877. if tbl[i][j - 1]:
  878. left = min(left, np.min(
  879. [bx["x0"] - a["x1"] for a in tbl[i][j - 1]]))
  880. if j + 1 < len(tbl[0]) and not ff:
  881. for i in range(len(tbl)):
  882. if tbl[i][j + 1]:
  883. right = min(right, np.min(
  884. [a["x0"] - bx["x1"] for a in tbl[i][j + 1]]))
  885. assert left < 100000 or right < 100000
  886. if left < right:
  887. for jj in range(j, len(tbl[0])):
  888. for i in range(len(tbl)):
  889. for a in tbl[i][jj]:
  890. a["cn"] -= 1
  891. if tbl[ii][j - 1]:
  892. tbl[ii][j - 1].extend(tbl[ii][j])
  893. else:
  894. tbl[ii][j - 1] = tbl[ii][j]
  895. for i in range(len(tbl)):
  896. tbl[i].pop(j)
  897. else:
  898. for jj in range(j + 1, len(tbl[0])):
  899. for i in range(len(tbl)):
  900. for a in tbl[i][jj]:
  901. a["cn"] -= 1
  902. if tbl[ii][j + 1]:
  903. tbl[ii][j + 1].extend(tbl[ii][j])
  904. else:
  905. tbl[ii][j + 1] = tbl[ii][j]
  906. for i in range(len(tbl)):
  907. tbl[i].pop(j)
  908. cols.pop(j)
  909. assert len(cols) == len(tbl[0]), "Column NO. miss matched: %d vs %d" % (
  910. len(cols), len(tbl[0]))
  911. if len(cols) >= 4:
  912. # remove single in row
  913. i = 0
  914. while i < len(tbl):
  915. e, jj = 0, 0
  916. for j in range(len(tbl[i])):
  917. if tbl[i][j]:
  918. e += 1
  919. jj = j
  920. if e > 1:
  921. break
  922. if e > 1:
  923. i += 1
  924. continue
  925. f = (i > 0 and tbl[i - 1][jj] and tbl[i - 1]
  926. [jj][0].get("text")) or i == 0
  927. ff = (i + 1 < len(tbl) and tbl[i + 1][jj] and tbl[i + 1]
  928. [jj][0].get("text")) or i + 1 >= len(tbl)
  929. if f and ff:
  930. i += 1
  931. continue
  932. bx = tbl[i][jj][0]
  933. logging.debug("Relocate row single: " + bx["text"])
  934. # i row only has one value
  935. up, down = 100000, 100000
  936. if i > 0 and not f:
  937. for j in range(len(tbl[i - 1])):
  938. if tbl[i - 1][j]:
  939. up = min(up, np.min(
  940. [bx["top"] - a["bottom"] for a in tbl[i - 1][j]]))
  941. if i + 1 < len(tbl) and not ff:
  942. for j in range(len(tbl[i + 1])):
  943. if tbl[i + 1][j]:
  944. down = min(down, np.min(
  945. [a["top"] - bx["bottom"] for a in tbl[i + 1][j]]))
  946. assert up < 100000 or down < 100000
  947. if up < down:
  948. for ii in range(i, len(tbl)):
  949. for j in range(len(tbl[ii])):
  950. for a in tbl[ii][j]:
  951. a["rn"] -= 1
  952. if tbl[i - 1][jj]:
  953. tbl[i - 1][jj].extend(tbl[i][jj])
  954. else:
  955. tbl[i - 1][jj] = tbl[i][jj]
  956. tbl.pop(i)
  957. else:
  958. for ii in range(i + 1, len(tbl)):
  959. for j in range(len(tbl[ii])):
  960. for a in tbl[ii][j]:
  961. a["rn"] -= 1
  962. if tbl[i + 1][jj]:
  963. tbl[i + 1][jj].extend(tbl[i][jj])
  964. else:
  965. tbl[i + 1][jj] = tbl[i][jj]
  966. tbl.pop(i)
  967. rows.pop(i)
  968. # which rows are headers
  969. hdset = set([])
  970. for i in range(len(tbl)):
  971. cnt, h = 0, 0
  972. for j, arr in enumerate(tbl[i]):
  973. if not arr:
  974. continue
  975. cnt += 1
  976. if max_type == "Nu" and arr[0]["btype"] == "Nu":
  977. continue
  978. if any([a.get("H") for a in arr]) \
  979. or (max_type == "Nu" and arr[0]["btype"] != "Nu"):
  980. h += 1
  981. if h / cnt > 0.5:
  982. hdset.add(i)
  983. if html:
  984. return [self.__html_table(cap, hdset,
  985. self.__cal_spans(boxes, rows,
  986. cols, tbl, True)
  987. )]
  988. return self.__desc_table(cap, hdset,
  989. self.__cal_spans(boxes, rows, cols, tbl, False))
  990. def __html_table(self, cap, hdset, tbl):
  991. # constrcut HTML
  992. html = "<table>"
  993. if cap:
  994. html += f"<caption>{cap}</caption>"
  995. for i in range(len(tbl)):
  996. row = "<tr>"
  997. txts = []
  998. for j, arr in enumerate(tbl[i]):
  999. if arr is None:
  1000. continue
  1001. if not arr:
  1002. row += "<td></td>" if i not in hdset else "<th></th>"
  1003. continue
  1004. txt = ""
  1005. if arr:
  1006. h = min(np.min([c["bottom"] - c["top"] for c in arr]) / 2,
  1007. self.mean_height[arr[0]["page_number"] - 1] / 2)
  1008. txt = "".join([c["text"]
  1009. for c in self.sort_Y_firstly(arr, h)])
  1010. txts.append(txt)
  1011. sp = ""
  1012. if arr[0].get("colspan"):
  1013. sp = "colspan={}".format(arr[0]["colspan"])
  1014. if arr[0].get("rowspan"):
  1015. sp += " rowspan={}".format(arr[0]["rowspan"])
  1016. if i in hdset:
  1017. row += f"<th {sp} >" + txt + "</th>"
  1018. else:
  1019. row += f"<td {sp} >" + txt + "</td>"
  1020. if i in hdset:
  1021. if all([t in hdset for t in txts]):
  1022. continue
  1023. for t in txts:
  1024. hdset.add(t)
  1025. if row != "<tr>":
  1026. row += "</tr>"
  1027. else:
  1028. row = ""
  1029. html += "\n" + row
  1030. html += "\n</table>"
  1031. return html
  1032. def __desc_table(self, cap, hdr_rowno, tbl):
  1033. # get text of every colomn in header row to become header text
  1034. clmno = len(tbl[0])
  1035. rowno = len(tbl)
  1036. headers = {}
  1037. hdrset = set()
  1038. lst_hdr = []
  1039. for r in sorted(list(hdr_rowno)):
  1040. headers[r] = ["" for _ in range(clmno)]
  1041. for i in range(clmno):
  1042. if not tbl[r][i]:
  1043. continue
  1044. txt = "".join([a["text"].strip() for a in tbl[r][i]])
  1045. headers[r][i] = txt
  1046. hdrset.add(txt)
  1047. if all([not t for t in headers[r]]):
  1048. del headers[r]
  1049. hdr_rowno.remove(r)
  1050. continue
  1051. for j in range(clmno):
  1052. if headers[r][j]:
  1053. continue
  1054. if j >= len(lst_hdr):
  1055. break
  1056. headers[r][j] = lst_hdr[j]
  1057. lst_hdr = headers[r]
  1058. for i in range(rowno):
  1059. if i not in hdr_rowno:
  1060. continue
  1061. for j in range(i + 1, rowno):
  1062. if j not in hdr_rowno:
  1063. break
  1064. for k in range(clmno):
  1065. if not headers[j - 1][k]:
  1066. continue
  1067. if headers[j][k].find(headers[j - 1][k]) >= 0:
  1068. continue
  1069. if len(headers[j][k]) > len(headers[j - 1][k]):
  1070. headers[j][k] += ("的" if headers[j][k]
  1071. else "") + headers[j - 1][k]
  1072. else:
  1073. headers[j][k] = headers[j - 1][k] \
  1074. + ("的" if headers[j - 1][k] else "") \
  1075. + headers[j][k]
  1076. logging.debug(
  1077. f">>>>>>>>>>>>>>>>>{cap}:SIZE:{rowno}X{clmno} Header: {hdr_rowno}")
  1078. row_txt = []
  1079. for i in range(rowno):
  1080. if i in hdr_rowno:
  1081. continue
  1082. rtxt = []
  1083. def append(delimer):
  1084. nonlocal rtxt, row_txt
  1085. rtxt = delimer.join(rtxt)
  1086. if row_txt and len(row_txt[-1]) + len(rtxt) < 64:
  1087. row_txt[-1] += "\n" + rtxt
  1088. else:
  1089. row_txt.append(rtxt)
  1090. r = 0
  1091. if len(headers.items()):
  1092. _arr = [(i - r, r) for r, _ in headers.items() if r < i]
  1093. if _arr:
  1094. _, r = min(_arr, key=lambda x: x[0])
  1095. if r not in headers and clmno <= 2:
  1096. for j in range(clmno):
  1097. if not tbl[i][j]:
  1098. continue
  1099. txt = "".join([a["text"].strip() for a in tbl[i][j]])
  1100. if txt:
  1101. rtxt.append(txt)
  1102. if rtxt:
  1103. append(":")
  1104. continue
  1105. for j in range(clmno):
  1106. if not tbl[i][j]:
  1107. continue
  1108. txt = "".join([a["text"].strip() for a in tbl[i][j]])
  1109. if not txt:
  1110. continue
  1111. ctt = headers[r][j] if r in headers else ""
  1112. if ctt:
  1113. ctt += ":"
  1114. ctt += txt
  1115. if ctt:
  1116. rtxt.append(ctt)
  1117. if rtxt:
  1118. row_txt.append("; ".join(rtxt))
  1119. if cap:
  1120. row_txt = [t + f"\t——来自“{cap}”" for t in row_txt]
  1121. return row_txt
  1122. @staticmethod
  1123. def is_caption(bx):
  1124. patt = [
  1125. r"[图表]+[ 0-9::]{2,}"
  1126. ]
  1127. if any([re.match(p, bx["text"].strip()) for p in patt]) \
  1128. or bx["layout_type"].find("caption") >= 0:
  1129. return True
  1130. return False
  1131. def __extract_table_figure(self, need_image, ZM, return_html):
  1132. tables = {}
  1133. figures = {}
  1134. # extract figure and table boxes
  1135. i = 0
  1136. lst_lout_no = ""
  1137. nomerge_lout_no = []
  1138. while i < len(self.boxes):
  1139. if "layoutno" not in self.boxes[i]:
  1140. i += 1
  1141. continue
  1142. lout_no = str(self.boxes[i]["page_number"]) + \
  1143. "-" + str(self.boxes[i]["layoutno"])
  1144. if self.is_caption(self.boxes[i]) or self.boxes[i]["layout_type"] in ["table caption", "title",
  1145. "figure caption", "reference"]:
  1146. nomerge_lout_no.append(lst_lout_no)
  1147. if self.boxes[i]["layout_type"] == "table":
  1148. if re.match(r"(数据|资料|图表)*来源[:: ]", self.boxes[i]["text"]):
  1149. self.boxes.pop(i)
  1150. continue
  1151. if lout_no not in tables:
  1152. tables[lout_no] = []
  1153. tables[lout_no].append(self.boxes[i])
  1154. self.boxes.pop(i)
  1155. lst_lout_no = lout_no
  1156. continue
  1157. if need_image and self.boxes[i]["layout_type"] == "figure":
  1158. if re.match(r"(数据|资料|图表)*来源[:: ]", self.boxes[i]["text"]):
  1159. self.boxes.pop(i)
  1160. continue
  1161. if lout_no not in figures:
  1162. figures[lout_no] = []
  1163. figures[lout_no].append(self.boxes[i])
  1164. self.boxes.pop(i)
  1165. lst_lout_no = lout_no
  1166. continue
  1167. i += 1
  1168. # merge table on different pages
  1169. nomerge_lout_no = set(nomerge_lout_no)
  1170. tbls = sorted([(k, bxs) for k, bxs in tables.items()],
  1171. key=lambda x: (x[1][0]["top"], x[1][0]["x0"]))
  1172. i = len(tbls) - 1
  1173. while i - 1 >= 0:
  1174. k0, bxs0 = tbls[i - 1]
  1175. k, bxs = tbls[i]
  1176. i -= 1
  1177. if k0 in nomerge_lout_no:
  1178. continue
  1179. if bxs[0]["page_number"] == bxs0[0]["page_number"]:
  1180. continue
  1181. if bxs[0]["page_number"] - bxs0[0]["page_number"] > 1:
  1182. continue
  1183. mh = self.mean_height[bxs[0]["page_number"] - 1]
  1184. if self._y_dis(bxs0[-1], bxs[0]) > mh * 23:
  1185. continue
  1186. tables[k0].extend(tables[k])
  1187. del tables[k]
  1188. def x_overlapped(a, b):
  1189. return not any([a["x1"] < b["x0"], a["x0"] > b["x1"]])
  1190. # find captions and pop out
  1191. i = 0
  1192. while i < len(self.boxes):
  1193. c = self.boxes[i]
  1194. # mh = self.mean_height[c["page_number"]-1]
  1195. if not self.is_caption(c):
  1196. i += 1
  1197. continue
  1198. # find the nearest layouts
  1199. def nearest(tbls):
  1200. nonlocal c
  1201. mink = ""
  1202. minv = 1000000000
  1203. for k, bxs in tbls.items():
  1204. for b in bxs[:10]:
  1205. if b.get("layout_type", "").find("caption") >= 0:
  1206. continue
  1207. y_dis = self._y_dis(c, b)
  1208. x_dis = self._x_dis(
  1209. c, b) if not x_overlapped(
  1210. c, b) else 0
  1211. dis = y_dis * y_dis + x_dis * x_dis
  1212. if dis < minv:
  1213. mink = k
  1214. minv = dis
  1215. return mink, minv
  1216. tk, tv = nearest(tables)
  1217. fk, fv = nearest(figures)
  1218. if min(tv, fv) > 2000:
  1219. i += 1
  1220. continue
  1221. if tv < fv:
  1222. tables[tk].insert(0, c)
  1223. logging.debug(
  1224. "TABLE:" +
  1225. self.boxes[i]["text"] +
  1226. "; Cap: " +
  1227. tk)
  1228. else:
  1229. figures[fk].insert(0, c)
  1230. logging.debug(
  1231. "FIGURE:" +
  1232. self.boxes[i]["text"] +
  1233. "; Cap: " +
  1234. tk)
  1235. self.boxes.pop(i)
  1236. res = []
  1237. def cropout(bxs, ltype):
  1238. nonlocal ZM
  1239. pn = set([b["page_number"] - 1 for b in bxs])
  1240. if len(pn) < 2:
  1241. pn = list(pn)[0]
  1242. ht = self.page_cum_height[pn]
  1243. b = {
  1244. "x0": np.min([b["x0"] for b in bxs]),
  1245. "top": np.min([b["top"] for b in bxs]) - ht,
  1246. "x1": np.max([b["x1"] for b in bxs]),
  1247. "bottom": np.max([b["bottom"] for b in bxs]) - ht
  1248. }
  1249. louts = [l for l in self.page_layout[pn] if l["type"] == ltype]
  1250. ii = self.__find_overlapped(b, louts, naive=True)
  1251. if ii is not None:
  1252. b = louts[ii]
  1253. else:
  1254. logging.warn(
  1255. f"Missing layout match: {pn + 1},%s" %
  1256. (bxs[0].get(
  1257. "layoutno", "")))
  1258. left, top, right, bott = b["x0"], b["top"], b["x1"], b["bottom"]
  1259. return self.page_images[pn] \
  1260. .crop((left * ZM, top * ZM,
  1261. right * ZM, bott * ZM))
  1262. pn = {}
  1263. for b in bxs:
  1264. p = b["page_number"] - 1
  1265. if p not in pn:
  1266. pn[p] = []
  1267. pn[p].append(b)
  1268. pn = sorted(pn.items(), key=lambda x: x[0])
  1269. imgs = [cropout(arr, ltype) for p, arr in pn]
  1270. pic = Image.new("RGB",
  1271. (int(np.max([i.size[0] for i in imgs])),
  1272. int(np.sum([m.size[1] for m in imgs]))),
  1273. (245, 245, 245))
  1274. height = 0
  1275. for img in imgs:
  1276. pic.paste(img, (0, int(height)))
  1277. height += img.size[1]
  1278. return pic
  1279. # crop figure out and add caption
  1280. for k, bxs in figures.items():
  1281. txt = "\n".join(
  1282. [b["text"] for b in bxs
  1283. if not re.match(r"[0-9a-z.\+%-]", b["text"].strip())
  1284. and len(b["text"].strip()) >= 4
  1285. ]
  1286. )
  1287. if not txt:
  1288. continue
  1289. res.append(
  1290. (cropout(
  1291. bxs,
  1292. "figure"),
  1293. [txt] if not return_html else [f"<p>{txt}</p>"]))
  1294. for k, bxs in tables.items():
  1295. if not bxs:
  1296. continue
  1297. res.append((cropout(bxs, "table"),
  1298. self.__construct_table(bxs, html=return_html)))
  1299. return res
  1300. def proj_match(self, line):
  1301. if len(line) <= 2:
  1302. return
  1303. if re.match(r"[0-9 ().,%%+/-]+$", line):
  1304. return False
  1305. for p, j in [
  1306. (r"第[零一二三四五六七八九十百]+章", 1),
  1307. (r"第[零一二三四五六七八九十百]+[条节]", 2),
  1308. (r"[零一二三四五六七八九十百]+[、  ]", 3),
  1309. (r"[\((][零一二三四五六七八九十百]+[)\)]", 4),
  1310. (r"[0-9]+(、|\.[  ]|\.[^0-9])", 5),
  1311. (r"[0-9]+\.[0-9]+(、|[.  ]|[^0-9])", 6),
  1312. (r"[0-9]+\.[0-9]+\.[0-9]+(、|[  ]|[^0-9])", 7),
  1313. (r"[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(、|[  ]|[^0-9])", 8),
  1314. (r".{,48}[::??]$", 9),
  1315. (r"[0-9]+)", 10),
  1316. (r"[\((][0-9]+[)\)]", 11),
  1317. (r"[零一二三四五六七八九十百]+是", 12),
  1318. (r"[⚫•➢✓]", 12)
  1319. ]:
  1320. if re.match(p, line):
  1321. return j
  1322. return
  1323. def __filterout_scraps(self, boxes, ZM):
  1324. def line_tag(bx):
  1325. pn = [bx["page_number"]]
  1326. top = bx["top"] - self.page_cum_height[pn[0] - 1]
  1327. bott = bx["bottom"] - self.page_cum_height[pn[0] - 1]
  1328. while bott * ZM > self.page_images[pn[-1] - 1].size[1]:
  1329. bott -= self.page_images[pn[-1] - 1].size[1] / ZM
  1330. pn.append(pn[-1] + 1)
  1331. return "@@{}\t{:.1f}\t{:.1f}\t{:.1f}\t{:.1f}##" \
  1332. .format("-".join([str(p) for p in pn]),
  1333. bx["x0"], bx["x1"], top, bott)
  1334. def width(b):
  1335. return b["x1"] - b["x0"]
  1336. def height(b):
  1337. return b["bottom"] - b["top"]
  1338. def usefull(b):
  1339. if b.get("layout_type"):
  1340. return True
  1341. if width(
  1342. b) > self.page_images[b["page_number"] - 1].size[0] / ZM / 3:
  1343. return True
  1344. if b["bottom"] - b["top"] > self.mean_height[b["page_number"] - 1]:
  1345. return True
  1346. return False
  1347. res = []
  1348. while boxes:
  1349. lines = []
  1350. widths = []
  1351. pw = self.page_images[boxes[0]["page_number"] - 1].size[0] / ZM
  1352. mh = self.mean_height[boxes[0]["page_number"] - 1]
  1353. mj = self.proj_match(
  1354. boxes[0]["text"]) or boxes[0].get(
  1355. "layout_type",
  1356. "") == "title"
  1357. def dfs(line, st):
  1358. nonlocal mh, pw, lines, widths
  1359. lines.append(line)
  1360. widths.append(width(line))
  1361. width_mean = np.mean(widths)
  1362. mmj = self.proj_match(
  1363. line["text"]) or line.get(
  1364. "layout_type",
  1365. "") == "title"
  1366. for i in range(st + 1, min(st + 20, len(boxes))):
  1367. if (boxes[i]["page_number"] - line["page_number"]) > 0:
  1368. break
  1369. if not mmj and self._y_dis(
  1370. line, boxes[i]) >= 3 * mh and height(line) < 1.5 * mh:
  1371. break
  1372. if not usefull(boxes[i]):
  1373. continue
  1374. if mmj or \
  1375. (self._x_dis(boxes[i], line) < pw / 10): \
  1376. # and abs(width(boxes[i])-width_mean)/max(width(boxes[i]),width_mean)<0.5):
  1377. # concat following
  1378. dfs(boxes[i], i)
  1379. boxes.pop(i)
  1380. break
  1381. try:
  1382. if usefull(boxes[0]):
  1383. dfs(boxes[0], 0)
  1384. else:
  1385. logging.debug("WASTE: " + boxes[0]["text"])
  1386. except Exception as e:
  1387. pass
  1388. boxes.pop(0)
  1389. mw = np.mean(widths)
  1390. if mj or mw / pw >= 0.35 or mw > 200:
  1391. res.append("\n".join([c["text"] + line_tag(c) for c in lines]))
  1392. else:
  1393. logging.debug("REMOVED: " +
  1394. "<<".join([c["text"] for c in lines]))
  1395. return "\n\n".join(res)
  1396. def __call__(self, fnm, need_image=True, zoomin=3, return_html=False):
  1397. self.lefted_chars = []
  1398. self.mean_height = []
  1399. self.mean_width = []
  1400. self.boxes = []
  1401. self.garbages = {}
  1402. self.page_cum_height = [0]
  1403. self.page_layout = []
  1404. try:
  1405. self.pdf = pdfplumber.open(fnm) if isinstance(fnm, str) else pdfplumber.open(BytesIO(fnm))
  1406. self.page_images = [p.to_image(resolution=72*zoomin).annotated for i,p in enumerate(self.pdf.pages[:299])]
  1407. self.page_chars = [[c for c in self.pdf.pages[i].chars if self._has_color(c)] for i in range(len(self.page_images))]
  1408. except Exception as e:
  1409. self.pdf = fitz.open(fnm) if isinstance(fnm, str) else fitz.open(stream=fnm, filetype="pdf")
  1410. self.page_images = []
  1411. self.page_chars = []
  1412. mat = fitz.Matrix(zoomin, zoomin)
  1413. for page in self.pdf:
  1414. pix = page.getPixmap(matrix = mat)
  1415. img = Image.frombytes("RGB", [pix.width, pix.height],
  1416. pix.samples)
  1417. self.page_images.append(img)
  1418. self.page_chars.append([])
  1419. logging.info("Images converted.")
  1420. for i, img in enumerate(self.page_images):
  1421. chars = self.page_chars[i]
  1422. self.mean_height.append(
  1423. np.median(sorted([c["height"] for c in chars])) if chars else 0
  1424. )
  1425. self.mean_width.append(
  1426. np.median(sorted([c["width"] for c in chars])) if chars else 8
  1427. )
  1428. if i > 0:
  1429. if not chars:
  1430. self.page_cum_height.append(img.size[1] / zoomin)
  1431. else:
  1432. self.page_cum_height.append(
  1433. np.max([c["bottom"] for c in chars]))
  1434. self.__ocr_paddle(i + 1, img, chars, zoomin)
  1435. self.__layouts_paddle(zoomin)
  1436. self.page_cum_height = np.cumsum(self.page_cum_height)
  1437. assert len(self.page_cum_height) == len(self.page_images)
  1438. garbage = set()
  1439. for k in self.garbages.keys():
  1440. self.garbages[k] = Counter(self.garbages[k])
  1441. for g, c in self.garbages[k].items():
  1442. if c > 1:
  1443. garbage.add(g)
  1444. logging.debug("GARBAGE:" + ",".join(garbage))
  1445. self.boxes = [b for b in self.boxes if b["text"] not in garbage]
  1446. # cumlative Y
  1447. for i in range(len(self.boxes)):
  1448. self.boxes[i]["top"] += \
  1449. self.page_cum_height[self.boxes[i]["page_number"] - 1]
  1450. self.boxes[i]["bottom"] += \
  1451. self.page_cum_height[self.boxes[i]["page_number"] - 1]
  1452. self.__table_transformer_job(zoomin)
  1453. self.__text_merge(garbage)
  1454. self.__filter_forpages()
  1455. tbls = self.__extract_table_figure(need_image, zoomin, return_html)
  1456. return self.__filterout_scraps(deepcopy(self.boxes), zoomin), tbls
  1457. def remove_tag(self, txt):
  1458. return re.sub(r"@@[\t0-9.-]+?##", "", txt)
  1459. def crop(self, text, ZM=3):
  1460. imgs = []
  1461. for tag in re.findall(r"@@[0-9-]+\t[0-9.\t]+##", text):
  1462. pn, left, right, top, bottom = tag.strip(
  1463. "#").strip("@").split("\t")
  1464. left, right, top, bottom = float(left), float(
  1465. right), float(top), float(bottom)
  1466. bottom *= ZM
  1467. pns = [int(p) - 1 for p in pn.split("-")]
  1468. for pn in pns[1:]:
  1469. bottom += self.page_images[pn - 1].size[1]
  1470. imgs.append(
  1471. self.page_images[pns[0]].crop((left * ZM, top * ZM,
  1472. right *
  1473. ZM, min(
  1474. bottom, self.page_images[pns[0]].size[1])
  1475. ))
  1476. )
  1477. bottom -= self.page_images[pns[0]].size[1]
  1478. for pn in pns[1:]:
  1479. imgs.append(
  1480. self.page_images[pn].crop((left * ZM, 0,
  1481. right * ZM,
  1482. min(bottom,
  1483. self.page_images[pn].size[1])
  1484. ))
  1485. )
  1486. bottom -= self.page_images[pn].size[1]
  1487. if not imgs:
  1488. return
  1489. GAP = 2
  1490. height = 0
  1491. for img in imgs:
  1492. height += img.size[1] + GAP
  1493. height = int(height)
  1494. pic = Image.new("RGB",
  1495. (int(np.max([i.size[0] for i in imgs])), height),
  1496. (245, 245, 245))
  1497. height = 0
  1498. for img in imgs:
  1499. pic.paste(img, (0, int(height)))
  1500. height += img.size[1] + GAP
  1501. return pic
  1502. if __name__ == "__main__":
  1503. pass