106 lines
4.1 KiB
Python
106 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
为已合并的 OCR 书籍 md 补插图：从 _chunks/*.jsonl 提取每页 markdown.images，
|
||
下载到 <md_dir>/images/，并在每页 '<!-- page N -->' 标记后插入图片引用。
|
||
|
||
用法: python ocr_book_insert_images.py <book.md>
|
||
（jsonl 缓存需在 md 同目录 _chunks/ 下，即 ocr_book_to_md.py 的产出布局）
|
||
"""
|
||
import json
|
||
import glob
|
||
import os
|
||
import re
|
||
import sys
|
||
import requests
|
||
|
||
PAGE_RE = re.compile(r"<!-- page (\d+) -->")
|
||
|
||
|
||
def main():
|
||
if len(sys.argv) < 2:
|
||
print(__doc__)
|
||
sys.exit(1)
|
||
md_path = os.path.abspath(sys.argv[1])
|
||
md_dir = os.path.dirname(md_path)
|
||
chunks_dir = os.path.join(md_dir, "_chunks")
|
||
img_root = os.path.join(md_dir, "images")
|
||
|
||
# 1) 按页序遍历所有 lpr，收集每页 images（全局游标推进）
|
||
page_imgs = [] # list of dict(key->url)，与页序对齐
|
||
for jf in sorted(glob.glob(os.path.join(chunks_dir, "chunk_*.jsonl"))):
|
||
for line in open(jf, encoding="utf-8"):
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
o = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
for res in (o.get("result") or {}).get("layoutParsingResults") or []:
|
||
imgs = (res.get("markdown") or {}).get("images") or {}
|
||
page_imgs.append(imgs)
|
||
print(f"pages from jsonl: {len(page_imgs)}")
|
||
|
||
# 2) 下载（同 URL 去重），保存到 images/<key>
|
||
os.makedirs(img_root, exist_ok=True)
|
||
url_to_key = {}
|
||
n_dl = n_skip = n_fail = 0
|
||
for pidx, imgs in enumerate(page_imgs, 1):
|
||
for key, url in imgs.items():
|
||
dst = os.path.join(img_root, key)
|
||
if url in url_to_key:
|
||
n_skip += 1
|
||
continue
|
||
if os.path.exists(dst) and os.path.getsize(dst) > 0:
|
||
url_to_key[url] = key # 已下载过，仅登记去重
|
||
n_skip += 1
|
||
continue
|
||
try:
|
||
os.makedirs(os.path.dirname(dst), exist_ok=True)
|
||
r = requests.get(url, timeout=60)
|
||
if r.status_code != 200:
|
||
raise RuntimeError(f"HTTP {r.status_code}")
|
||
with open(dst, "wb") as f:
|
||
f.write(r.content)
|
||
url_to_key[url] = key
|
||
n_dl += 1
|
||
print(f" [page {pidx}] dl {key} ({len(r.content)} B)")
|
||
except Exception as e:
|
||
n_fail += 1
|
||
print(f" [page {pidx}] FAIL {key}: {e}")
|
||
print(f"downloaded: {n_dl}, dedup-skip: {n_skip}, failed: {n_fail}")
|
||
|
||
# 3) 读 md，按页补齐图片引用（幂等）：
|
||
# VL 的 markdown.text 自带 HTML <img src="imgs/xxx.jpg">，路径需修为 images/imgs/；
|
||
# 仅对 text 中未出现的图（images 字段有但 text 无引用）插入 markdown 引用，避免双图。
|
||
with open(md_path, encoding="utf-8") as f:
|
||
md_text = f.read()
|
||
# 修 VL 原生 HTML 路径（幂等：已修的 images/imgs/ 不再变）
|
||
md_text = re.sub(r'(<img src=")imgs/', r"\1images/imgs/", md_text)
|
||
# 已有引用集合（HTML 或 markdown 形式，均以 imgs/<key> 结尾）
|
||
have = set(re.findall(r"images/imgs/([^\")\s]+\.jpg)", md_text))
|
||
lines = md_text.split("\n")
|
||
out = []
|
||
page_no = 0
|
||
inserted = 0
|
||
for line in lines:
|
||
out.append(line)
|
||
m = PAGE_RE.search(line)
|
||
if m:
|
||
page_no = int(m.group(1))
|
||
imgs = page_imgs[page_no - 1] if page_no - 1 < len(page_imgs) else {}
|
||
for i, (key, _url) in enumerate(imgs.items(), 1):
|
||
rel = key.replace("\\", "/")
|
||
if os.path.basename(key) in have: # key 形如 imgs/xxx.jpg，have 存 basename
|
||
continue # text 已有 HTML 引用，不重复
|
||
out.append(f"")
|
||
have.add(os.path.basename(key))
|
||
inserted += 1
|
||
with open(md_path, "w", encoding="utf-8") as f:
|
||
f.write("\n".join(out))
|
||
print(f"inserted {inserted} missing image refs into {md_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|