mcol_sort.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337
  1. """
  2. This is an advanced PyMuPDF utility for detecting multi-column pages.
  3. It can be used in a shell script, or its main function can be imported and
  4. invoked as descript below.
  5. Features
  6. ---------
  7. - Identify text belonging to (a variable number of) columns on the page.
  8. - Text with different background color is handled separately, allowing for
  9. easier treatment of side remarks, comment boxes, etc.
  10. - Uses text block detection capability to identify text blocks and
  11. uses the block bboxes as primary structuring principle.
  12. - Supports ignoring footers via a footer margin parameter.
  13. - Returns re-created text boundary boxes (integer coordinates), sorted ascending
  14. by the top, then by the left coordinates.
  15. Restrictions
  16. -------------
  17. - Only supporting horizontal, left-to-right text
  18. - Returns a list of text boundary boxes - not the text itself. The caller is
  19. expected to extract text from within the returned boxes.
  20. - Text written above images is ignored altogether (option).
  21. - This utility works as expected in most cases. The following situation cannot
  22. be handled correctly:
  23. * overlapping (non-disjoint) text blocks
  24. * image captions are not recognized and are handled like normal text
  25. Usage
  26. ------
  27. - As a CLI shell command use
  28. python multi_column.py input.pdf footer_margin
  29. Where footer margin is the height of the bottom stripe to ignore on each page.
  30. This code is intended to be modified according to your need.
  31. - Use in a Python script as follows:
  32. ----------------------------------------------------------------------------------
  33. from multi_column import column_boxes
  34. # for each page execute
  35. bboxes = column_boxes(page, footer_margin=50, no_image_text=True)
  36. # bboxes is a list of fitz.IRect objects, that are sort ascending by their y0,
  37. # then x0 coordinates. Their text content can be extracted by all PyMuPDF
  38. # get_text() variants, like for instance the following:
  39. for rect in bboxes:
  40. print(page.get_text(clip=rect, sort=True))
  41. ----------------------------------------------------------------------------------
  42. """
  43. import os
  44. import sys
  45. from libs.commons import fitz
  46. def column_boxes(page, footer_margin=50, header_margin=50, no_image_text=True):
  47. """Determine bboxes which wrap a column."""
  48. paths = page.get_drawings()
  49. bboxes = []
  50. # path rectangles
  51. path_rects = []
  52. # image bboxes
  53. img_bboxes = []
  54. # bboxes of non-horizontal text
  55. # avoid when expanding horizontal text boxes
  56. vert_bboxes = []
  57. # compute relevant page area
  58. clip = +page.rect
  59. clip.y1 -= footer_margin # Remove footer area
  60. clip.y0 += header_margin # Remove header area
  61. def can_extend(temp, bb, bboxlist):
  62. """Determines whether rectangle 'temp' can be extended by 'bb'
  63. without intersecting any of the rectangles contained in 'bboxlist'.
  64. Items of bboxlist may be None if they have been removed.
  65. Returns:
  66. True if 'temp' has no intersections with items of 'bboxlist'.
  67. """
  68. for b in bboxlist:
  69. if not intersects_bboxes(temp, vert_bboxes) and (
  70. b == None or b == bb or (temp & b).is_empty
  71. ):
  72. continue
  73. return False
  74. return True
  75. def in_bbox(bb, bboxes):
  76. """Return 1-based number if a bbox contains bb, else return 0."""
  77. for i, bbox in enumerate(bboxes):
  78. if bb in bbox:
  79. return i + 1
  80. return 0
  81. def intersects_bboxes(bb, bboxes):
  82. """Return True if a bbox intersects bb, else return False."""
  83. for bbox in bboxes:
  84. if not (bb & bbox).is_empty:
  85. return True
  86. return False
  87. def extend_right(bboxes, width, path_bboxes, vert_bboxes, img_bboxes):
  88. """Extend a bbox to the right page border.
  89. Whenever there is no text to the right of a bbox, enlarge it up
  90. to the right page border.
  91. Args:
  92. bboxes: (list[IRect]) bboxes to check
  93. width: (int) page width
  94. path_bboxes: (list[IRect]) bboxes with a background color
  95. vert_bboxes: (list[IRect]) bboxes with vertical text
  96. img_bboxes: (list[IRect]) bboxes of images
  97. Returns:
  98. Potentially modified bboxes.
  99. """
  100. for i, bb in enumerate(bboxes):
  101. # do not extend text with background color
  102. if in_bbox(bb, path_bboxes):
  103. continue
  104. # do not extend text in images
  105. if in_bbox(bb, img_bboxes):
  106. continue
  107. # temp extends bb to the right page border
  108. temp = +bb
  109. temp.x1 = width
  110. # do not cut through colored background or images
  111. if intersects_bboxes(temp, path_bboxes + vert_bboxes + img_bboxes):
  112. continue
  113. # also, do not intersect other text bboxes
  114. check = can_extend(temp, bb, bboxes)
  115. if check:
  116. bboxes[i] = temp # replace with enlarged bbox
  117. return [b for b in bboxes if b != None]
  118. def clean_nblocks(nblocks):
  119. """Do some elementary cleaning."""
  120. # 1. remove any duplicate blocks.
  121. blen = len(nblocks)
  122. if blen < 2:
  123. return nblocks
  124. start = blen - 1
  125. for i in range(start, -1, -1):
  126. bb1 = nblocks[i]
  127. bb0 = nblocks[i - 1]
  128. if bb0 == bb1:
  129. del nblocks[i]
  130. # 2. repair sequence in special cases:
  131. # consecutive bboxes with almost same bottom value are sorted ascending
  132. # by x-coordinate.
  133. y1 = nblocks[0].y1 # first bottom coordinate
  134. i0 = 0 # its index
  135. i1 = -1 # index of last bbox with same bottom
  136. # Iterate over bboxes, identifying segments with approx. same bottom value.
  137. # Replace every segment by its sorted version.
  138. for i in range(1, len(nblocks)):
  139. b1 = nblocks[i]
  140. if abs(b1.y1 - y1) > 10: # different bottom
  141. if i1 > i0: # segment length > 1? Sort it!
  142. nblocks[i0 : i1 + 1] = sorted(
  143. nblocks[i0 : i1 + 1], key=lambda b: b.x0
  144. )
  145. y1 = b1.y1 # store new bottom value
  146. i0 = i # store its start index
  147. i1 = i # store current index
  148. if i1 > i0: # segment waiting to be sorted
  149. nblocks[i0 : i1 + 1] = sorted(nblocks[i0 : i1 + 1], key=lambda b: b.x0)
  150. return nblocks
  151. # extract vector graphics
  152. for p in paths:
  153. path_rects.append(p["rect"].irect)
  154. path_bboxes = path_rects
  155. # sort path bboxes by ascending top, then left coordinates
  156. path_bboxes.sort(key=lambda b: (b.y0, b.x0))
  157. # bboxes of images on page, no need to sort them
  158. for item in page.get_images():
  159. img_bboxes.extend(page.get_image_rects(item[0]))
  160. # blocks of text on page
  161. blocks = page.get_text(
  162. "dict",
  163. flags=fitz.TEXTFLAGS_TEXT,
  164. clip=clip,
  165. )["blocks"]
  166. # Make block rectangles, ignoring non-horizontal text
  167. for b in blocks:
  168. bbox = fitz.IRect(b["bbox"]) # bbox of the block
  169. # ignore text written upon images
  170. if no_image_text and in_bbox(bbox, img_bboxes):
  171. continue
  172. # confirm first line to be horizontal
  173. line0 = b["lines"][0] # get first line
  174. if line0["dir"] != (1, 0): # only accept horizontal text
  175. vert_bboxes.append(bbox)
  176. continue
  177. srect = fitz.EMPTY_IRECT()
  178. for line in b["lines"]:
  179. lbbox = fitz.IRect(line["bbox"])
  180. text = "".join([s["text"].strip() for s in line["spans"]])
  181. if len(text) > 1:
  182. srect |= lbbox
  183. bbox = +srect
  184. if not bbox.is_empty:
  185. bboxes.append(bbox)
  186. # Sort text bboxes by ascending background, top, then left coordinates
  187. bboxes.sort(key=lambda k: (in_bbox(k, path_bboxes), k.y0, k.x0))
  188. # Extend bboxes to the right where possible
  189. bboxes = extend_right(
  190. bboxes, int(page.rect.width), path_bboxes, vert_bboxes, img_bboxes
  191. )
  192. # immediately return of no text found
  193. if bboxes == []:
  194. return []
  195. # --------------------------------------------------------------------
  196. # Join bboxes to establish some column structure
  197. # --------------------------------------------------------------------
  198. # the final block bboxes on page
  199. nblocks = [bboxes[0]] # pre-fill with first bbox
  200. bboxes = bboxes[1:] # remaining old bboxes
  201. for i, bb in enumerate(bboxes): # iterate old bboxes
  202. check = False # indicates unwanted joins
  203. # check if bb can extend one of the new blocks
  204. for j in range(len(nblocks)):
  205. nbb = nblocks[j] # a new block
  206. # never join across columns
  207. if bb == None or nbb.x1 < bb.x0 or bb.x1 < nbb.x0:
  208. continue
  209. # never join across different background colors
  210. if in_bbox(nbb, path_bboxes) != in_bbox(bb, path_bboxes):
  211. continue
  212. temp = bb | nbb # temporary extension of new block
  213. check = can_extend(temp, nbb, nblocks)
  214. if check == True:
  215. break
  216. if not check: # bb cannot be used to extend any of the new bboxes
  217. nblocks.append(bb) # so add it to the list
  218. j = len(nblocks) - 1 # index of it
  219. temp = nblocks[j] # new bbox added
  220. # check if some remaining bbox is contained in temp
  221. check = can_extend(temp, bb, bboxes)
  222. if check == False:
  223. nblocks.append(bb)
  224. else:
  225. nblocks[j] = temp
  226. bboxes[i] = None
  227. # do some elementary cleaning
  228. nblocks = clean_nblocks(nblocks)
  229. # return identified text bboxes
  230. return nblocks
  231. if __name__ == "__main__":
  232. """Only for debugging purposes, currently.
  233. Draw red borders around the returned text bboxes and insert
  234. the bbox number.
  235. Then save the file under the name "input-blocks.pdf".
  236. """
  237. # get the file name
  238. filename = sys.argv[1]
  239. # check if footer margin is given
  240. if len(sys.argv) > 2:
  241. footer_margin = int(sys.argv[2])
  242. else: # use default vaue
  243. footer_margin = 50
  244. # check if header margin is given
  245. if len(sys.argv) > 3:
  246. header_margin = int(sys.argv[3])
  247. else: # use default vaue
  248. header_margin = 50
  249. # open document
  250. doc = fitz.open(filename)
  251. # iterate over the pages
  252. for page in doc:
  253. # remove any geometry issues
  254. page.wrap_contents()
  255. # get the text bboxes
  256. bboxes = column_boxes(page, footer_margin=footer_margin, header_margin=header_margin)
  257. # prepare a canvas to draw rectangles and text
  258. shape = page.new_shape()
  259. # iterate over the bboxes
  260. for i, rect in enumerate(bboxes):
  261. shape.draw_rect(rect) # draw a border
  262. # write sequence number
  263. shape.insert_text(rect.tl + (5, 15), str(i), color=fitz.pdfcolor["red"])
  264. # finish drawing / text with color red
  265. shape.finish(color=fitz.pdfcolor["red"])
  266. shape.commit() # store to the page
  267. # save document with text bboxes
  268. doc.ez_save(filename.replace(".pdf", "-blocks.pdf"))