pdf_meta_scan.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371
  1. """
  2. 输入: s3路径,每行一个
  3. 输出: pdf文件元信息,包括每一页上的所有图片的长宽高,bbox位置
  4. """
  5. import math
  6. import sys
  7. import click
  8. from libs.commons import read_file, mymax, get_top_percent_list
  9. import json
  10. from libs.commons import fitz
  11. from loguru import logger
  12. from collections import Counter
  13. from libs.drop_reason import DropReason
  14. from libs.language import detect_lang
  15. scan_max_page = 50
  16. junk_limit_min = 10
  17. def calculate_max_image_area_per_page(result:list, page_width_pts, page_height_pts):
  18. max_image_area_per_page = [mymax([(x1 - x0) * (y1 - y0) for x0, y0, x1, y1, _ in page_img_sz]) for page_img_sz in
  19. result]
  20. page_area = int(page_width_pts) * int(page_height_pts)
  21. max_image_area_per_page = [area / page_area for area in max_image_area_per_page]
  22. max_image_area_per_page = [area for area in max_image_area_per_page if area > 0.6]
  23. return max_image_area_per_page
  24. def process_image(page, junk_img_bojids=[]):
  25. page_result = []# 存每个页面里的多张图四元组信息
  26. items = page.get_images()
  27. dedup = set()
  28. for img in items:
  29. # 这里返回的是图片在page上的实际展示的大小。返回一个数组,每个元素第一部分是
  30. img_bojid = img[0]# 在pdf文件中是全局唯一的,如果这个图反复出现在pdf里那么就可能是垃圾信息,例如水印、页眉页脚等
  31. if img_bojid in junk_img_bojids:# 如果是垃圾图像,就跳过
  32. continue
  33. recs = page.get_image_rects(img, transform=True)
  34. if recs:
  35. rec = recs[0][0]
  36. x0, y0, x1, y1 = map(int, rec)
  37. width = x1 - x0
  38. height = y1 - y0
  39. if (x0, y0, x1, y1, img_bojid) in dedup: # 这里面会出现一些重复的bbox,无需重复出现,需要去掉
  40. continue
  41. if not all([width, height]): # 长和宽任何一个都不能是0,否则这个图片不可见,没有实际意义
  42. continue
  43. dedup.add((x0, y0, x1, y1, img_bojid))
  44. page_result.append([x0, y0, x1, y1, img_bojid])
  45. return page_result
  46. def get_image_info(doc: fitz.Document, page_width_pts, page_height_pts) -> list:
  47. """
  48. 返回每个页面里的图片的四元组,每个页面多个图片。
  49. :param doc:
  50. :return:
  51. """
  52. # 使用 Counter 计数 img_bojid 的出现次数
  53. img_bojid_counter = Counter(img[0] for page in doc for img in page.get_images())
  54. # 找出出现次数超过 len(doc) 半数的 img_bojid
  55. junk_limit = max(len(doc)*0.5, junk_limit_min)# 对一些页数比较少的进行豁免
  56. junk_img_bojids = [img_bojid for img_bojid, count in img_bojid_counter.items() if count >= junk_limit]
  57. #todo 加个判断,用前十页就行,这些垃圾图片需要满足两个条件,不止出现的次数要足够多,而且图片占书页面积的比例要足够大,且图与图大小都差不多
  58. #有两种扫描版,一种文字版,这里可能会有误判
  59. #扫描版1:每页都有所有扫描页图片,特点是图占比大,每页展示1张
  60. #扫描版2,每页存储的扫描页图片数量递增,特点是图占比大,每页展示1张,需要清空junklist跑前50页图片信息用于分类判断
  61. #文字版1.每页存储所有图片,特点是图片占页面比例不大,每页展示可能为0也可能不止1张 这种pdf需要拿前10页抽样检测img大小和个数,如果符合需要清空junklist
  62. imgs_len_list = [len(page.get_images()) for page in doc]
  63. special_limit_pages = 10
  64. # 统一用前十页结果做判断
  65. result = []
  66. break_loop = False
  67. for i, page in enumerate(doc):
  68. if break_loop:
  69. break
  70. if i >= special_limit_pages:
  71. break
  72. page_result = process_image(page) # 这里不传junk_img_bojids,拿前十页所有图片信息用于后续分析
  73. result.append(page_result)
  74. for item in result:
  75. if not any(item): # 如果任何一页没有图片,说明是个文字版,需要判断是否为特殊文字版
  76. if max(imgs_len_list) == min(imgs_len_list) and max(imgs_len_list) >= junk_limit_min:# 如果是特殊文字版,就把junklist置空并break
  77. junk_img_bojids = []
  78. else:# 不是特殊文字版,是个普通文字版,但是存在垃圾图片,不置空junklist
  79. pass
  80. break_loop = True
  81. break
  82. if not break_loop:
  83. # 获取前80%的元素
  84. top_eighty_percent = get_top_percent_list(imgs_len_list, 0.8)
  85. # 检查前80%的元素是否都相等
  86. if len(set(top_eighty_percent)) == 1 and max(imgs_len_list) >= junk_limit_min:
  87. # # 如果前10页跑完都有图,根据每页图片数量是否相等判断是否需要清除junklist
  88. # if max(imgs_len_list) == min(imgs_len_list) and max(imgs_len_list) >= junk_limit_min:
  89. #前10页都有图,且每页数量一致,需要检测图片大小占页面的比例判断是否需要清除junklist
  90. max_image_area_per_page = calculate_max_image_area_per_page(result, page_width_pts, page_height_pts)
  91. if len(max_image_area_per_page) < 0.8 * special_limit_pages: # 前10页不全是大图,说明可能是个文字版pdf,把垃圾图片list置空
  92. junk_img_bojids = []
  93. else:# 前10页都有图,而且80%都是大图,且每页图片数量一致并都很多,说明是扫描版1,不需要清空junklist
  94. pass
  95. else:# 每页图片数量不一致,需要清掉junklist全量跑前50页图片
  96. junk_img_bojids = []
  97. #正式进入取前50页图片的信息流程
  98. result = []
  99. for i, page in enumerate(doc):
  100. if i >= scan_max_page:
  101. break
  102. page_result = process_image(page, junk_img_bojids)
  103. # logger.info(f"page {i} img_len: {len(page_result)}")
  104. result.append(page_result)
  105. return result, junk_img_bojids
  106. def get_pdf_page_size_pts(doc: fitz.Document):
  107. page_cnt = len(doc)
  108. l: int = min(page_cnt, 50)
  109. #把所有宽度和高度塞到两个list 分别取中位数(中间遇到了个在纵页里塞横页的pdf,导致宽高互换了)
  110. page_width_list = []
  111. page_height_list = []
  112. for i in range(l):
  113. page = doc[i]
  114. page_rect = page.rect
  115. page_width_list.append(page_rect.width)
  116. page_height_list.append(page_rect.height)
  117. page_width_list.sort()
  118. page_height_list.sort()
  119. median_width = page_width_list[len(page_width_list) // 2]
  120. median_height = page_height_list[len(page_height_list) // 2]
  121. return median_width, median_height
  122. def get_pdf_textlen_per_page(doc: fitz.Document):
  123. text_len_lst = []
  124. for page in doc:
  125. # 拿包含img和text的所有blocks
  126. # text_block = page.get_text("blocks")
  127. # 拿所有text的blocks
  128. # text_block = page.get_text("words")
  129. # text_block_len = sum([len(t[4]) for t in text_block])
  130. #拿所有text的str
  131. text_block = page.get_text("text")
  132. text_block_len = len(text_block)
  133. # logger.info(f"page {page.number} text_block_len: {text_block_len}")
  134. text_len_lst.append(text_block_len)
  135. return text_len_lst
  136. def get_pdf_text_layout_per_page(doc: fitz.Document):
  137. """
  138. 根据PDF文档的每一页文本布局,判断该页的文本布局是横向、纵向还是未知。
  139. Args:
  140. doc (fitz.Document): PDF文档对象。
  141. Returns:
  142. List[str]: 每一页的文本布局(横向、纵向、未知)。
  143. """
  144. text_layout_list = []
  145. for page_id, page in enumerate(doc):
  146. if page_id >= scan_max_page:
  147. break
  148. # 创建每一页的纵向和横向的文本行数计数器
  149. vertical_count = 0
  150. horizontal_count = 0
  151. text_dict = page.get_text("dict")
  152. if "blocks" in text_dict:
  153. for block in text_dict["blocks"]:
  154. if 'lines' in block:
  155. for line in block["lines"]:
  156. # 获取line的bbox顶点坐标
  157. x0, y0, x1, y1 = line['bbox']
  158. # 计算bbox的宽高
  159. width = x1 - x0
  160. height = y1 - y0
  161. # 计算bbox的面积
  162. area = width * height
  163. font_sizes = []
  164. for span in line['spans']:
  165. if 'size' in span:
  166. font_sizes.append(span['size'])
  167. if len(font_sizes) > 0:
  168. average_font_size = sum(font_sizes) / len(font_sizes)
  169. else:
  170. average_font_size = 10 # 有的line拿不到font_size,先定一个阈值100
  171. if area <= average_font_size ** 2: # 判断bbox的面积是否小于平均字体大小的平方,单字无法计算是横向还是纵向
  172. continue
  173. else:
  174. if 'wmode' in line: # 通过wmode判断文本方向
  175. if line['wmode'] == 1: # 判断是否为竖向文本
  176. vertical_count += 1
  177. elif line['wmode'] == 0: # 判断是否为横向文本
  178. horizontal_count += 1
  179. # if 'dir' in line: # 通过旋转角度计算判断文本方向
  180. # # 获取行的 "dir" 值
  181. # dir_value = line['dir']
  182. # cosine, sine = dir_value
  183. # # 计算角度
  184. # angle = math.degrees(math.acos(cosine))
  185. #
  186. # # 判断是否为横向文本
  187. # if abs(angle - 0) < 0.01 or abs(angle - 180) < 0.01:
  188. # # line_text = ' '.join(span['text'] for span in line['spans'])
  189. # # print('This line is horizontal:', line_text)
  190. # horizontal_count += 1
  191. # # 判断是否为纵向文本
  192. # elif abs(angle - 90) < 0.01 or abs(angle - 270) < 0.01:
  193. # # line_text = ' '.join(span['text'] for span in line['spans'])
  194. # # print('This line is vertical:', line_text)
  195. # vertical_count += 1
  196. # print(f"page_id: {page_id}, vertical_count: {vertical_count}, horizontal_count: {horizontal_count}")
  197. # 判断每一页的文本布局
  198. if vertical_count == 0 and horizontal_count == 0: # 该页没有文本,无法判断
  199. text_layout_list.append("unknow")
  200. continue
  201. else:
  202. if vertical_count > horizontal_count: # 该页的文本纵向行数大于横向的
  203. text_layout_list.append("vertical")
  204. else: # 该页的文本横向行数大于纵向的
  205. text_layout_list.append("horizontal")
  206. # logger.info(f"page_id: {page_id}, vertical_count: {vertical_count}, horizontal_count: {horizontal_count}")
  207. return text_layout_list
  208. '''定义一个自定义异常用来抛出单页svg太多的pdf'''
  209. class PageSvgsTooManyError(Exception):
  210. def __init__(self, message="Page SVGs are too many"):
  211. self.message = message
  212. super().__init__(self.message)
  213. def get_svgs_per_page(doc: fitz.Document):
  214. svgs_len_list = []
  215. for page_id, page in enumerate(doc):
  216. # svgs = page.get_drawings()
  217. svgs = page.get_cdrawings() # 切换成get_cdrawings,效率更高
  218. len_svgs = len(svgs)
  219. if len_svgs >= 3000:
  220. raise PageSvgsTooManyError()
  221. else:
  222. svgs_len_list.append(len_svgs)
  223. # logger.info(f"page_id: {page_id}, svgs_len: {len(svgs)}")
  224. return svgs_len_list
  225. def get_imgs_per_page(doc: fitz.Document):
  226. imgs_len_list = []
  227. for page_id, page in enumerate(doc):
  228. imgs = page.get_images()
  229. imgs_len_list.append(len(imgs))
  230. # logger.info(f"page_id: {page}, imgs_len: {len(imgs)}")
  231. return imgs_len_list
  232. def get_language(doc: fitz.Document):
  233. """
  234. 获取PDF文档的语言。
  235. Args:
  236. doc (fitz.Document): PDF文档对象。
  237. Returns:
  238. str: 文档语言,如 "en-US"。
  239. """
  240. language_lst = []
  241. for page_id, page in enumerate(doc):
  242. if page_id >= scan_max_page:
  243. break
  244. # 拿所有text的str
  245. text_block = page.get_text("text")
  246. page_language = detect_lang(text_block)
  247. language_lst.append(page_language)
  248. # logger.info(f"page_id: {page_id}, page_language: {page_language}")
  249. # 统计text_language_list中每种语言的个数
  250. count_dict = Counter(language_lst)
  251. # 输出text_language_list中出现的次数最多的语言
  252. language = max(count_dict, key=count_dict.get)
  253. return language
  254. def pdf_meta_scan(s3_pdf_path: str, pdf_bytes: bytes):
  255. """
  256. :param s3_pdf_path:
  257. :param pdf_bytes: pdf文件的二进制数据
  258. 几个维度来评价:是否加密,是否需要密码,纸张大小,总页数,是否文字可提取
  259. """
  260. doc = fitz.open("pdf", pdf_bytes)
  261. is_needs_password = doc.needs_pass
  262. is_encrypted = doc.is_encrypted
  263. total_page = len(doc)
  264. if total_page == 0:
  265. logger.warning(f"drop this pdf: {s3_pdf_path}, drop_reason: {DropReason.EMPTY_PDF}")
  266. result = {"need_drop": True, "drop_reason": DropReason.EMPTY_PDF}
  267. return result
  268. else:
  269. page_width_pts, page_height_pts = get_pdf_page_size_pts(doc)
  270. # logger.info(f"page_width_pts: {page_width_pts}, page_height_pts: {page_height_pts}")
  271. svgs_per_page = get_svgs_per_page(doc)
  272. # logger.info(f"svgs_per_page: {svgs_per_page}")
  273. imgs_per_page = get_imgs_per_page(doc)
  274. # logger.info(f"imgs_per_page: {imgs_per_page}")
  275. image_info_per_page, junk_img_bojids = get_image_info(doc, page_width_pts, page_height_pts)
  276. # logger.info(f"image_info_per_page: {image_info_per_page}, junk_img_bojids: {junk_img_bojids}")
  277. text_len_per_page = get_pdf_textlen_per_page(doc)
  278. # logger.info(f"text_len_per_page: {text_len_per_page}")
  279. text_layout_per_page = get_pdf_text_layout_per_page(doc)
  280. # logger.info(f"text_layout_per_page: {text_layout_per_page}")
  281. text_language = get_language(doc)
  282. # logger.info(f"text_language: {text_language}")
  283. # 最后输出一条json
  284. res = {
  285. "pdf_path": s3_pdf_path,
  286. "is_needs_password": is_needs_password,
  287. "is_encrypted": is_encrypted,
  288. "total_page": total_page,
  289. "page_width_pts": int(page_width_pts),
  290. "page_height_pts": int(page_height_pts),
  291. "image_info_per_page": image_info_per_page,
  292. "text_len_per_page": text_len_per_page,
  293. "text_layout_per_page": text_layout_per_page,
  294. "text_language": text_language,
  295. "svgs_per_page": svgs_per_page,
  296. "imgs_per_page": imgs_per_page, # 增加每页img数量list
  297. "junk_img_bojids": junk_img_bojids, # 增加垃圾图片的bojid list
  298. "metadata": doc.metadata
  299. }
  300. # logger.info(json.dumps(res, ensure_ascii=False))
  301. return res
  302. @click.command()
  303. @click.option('--s3-pdf-path', help='s3上pdf文件的路径')
  304. @click.option('--s3-profile', help='s3上的profile')
  305. def main(s3_pdf_path: str, s3_profile: str):
  306. """
  307. """
  308. try:
  309. file_content = read_file(s3_pdf_path, s3_profile)
  310. pdf_meta_scan(s3_pdf_path, file_content)
  311. except Exception as e:
  312. print(f"ERROR: {s3_pdf_path}, {e}", file=sys.stderr)
  313. logger.exception(e)
  314. if __name__ == '__main__':
  315. main()
  316. # "D:\project/20231108code-clean\pdf_cost_time\竖排例子\净空法师-大乘无量寿.pdf"
  317. # "D:\project/20231108code-clean\pdf_cost_time\竖排例子\三国演义_繁体竖排版.pdf"
  318. # "D:\project/20231108code-clean\pdf_cost_time\scihub\scihub_86800000\libgen.scimag86880000-86880999.zip_10.1021/acsami.1c03109.s002.pdf"
  319. # "D:/project/20231108code-clean/pdf_cost_time/scihub/scihub_18600000/libgen.scimag18645000-18645999.zip_10.1021/om3006239.pdf"
  320. # file_content = read_file("D:/project/20231108code-clean/pdf_cost_time/scihub/scihub_31000000/libgen.scimag31098000-31098999.zip_10.1109/isit.2006.261791.pdf","")
  321. # file_content = read_file("D:\project/20231108code-clean\pdf_cost_time\竖排例子\净空法师_大乘无量寿.pdf","")
  322. # doc = fitz.open("pdf", file_content)
  323. # text_layout_lst = get_pdf_text_layout_per_page(doc)
  324. # print(text_layout_lst)