File size: 5,373 Bytes
d4d4e7b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """Датасет для детектора строк из born-digital PDF.
Текстовый слой даёт точные рамки строк бесплатно и в нужном домене — это
разметка, которую иначе пришлось бы рисовать руками. Язык не важен: детектор
ищет строку, а не буквы, поэтому берём все строки подряд, включая русские.
Выпуски, попавшие в оценочную выборку распознавателя, исключаются по имени:
детектор не должен видеть страницы, на которых потом меряется сквозное
качество.
"""
import argparse, json, os, re
from concurrent.futures import ProcessPoolExecutor
import fitz
from PIL import Image
# выпуски, из которых нарезана real_eval (найдены сопоставлением текста),
# плюс весь 2011 год — его выпуск в архиве не опознался однозначно
HELD_OUT = re.compile(r"№ 47 \(2009\)|№ 1 \(2019\)|№ 1 \(2021\)|"
r"1 № \(2013\)|№ 8 \(2017\)|№ 16 \(2015\)|\(2011\)")
def page_boxes(page, scale):
"""Рамки строк текстового слоя в пикселях отрендеренной страницы."""
out = []
for blk in page.get_text("dict")["blocks"]:
for ln in blk.get("lines", []):
if not any(s["text"].strip() for s in ln["spans"]):
continue
x0, y0, x1, y1 = ln["bbox"]
w, h = (x1 - x0) * scale, (y1 - y0) * scale
# вырожденные и гигантские рамки — мусор вёрстки, не строки
if w < 6 or h < 4 or h > 400 or w / h > 200:
continue
out.append([round(x0 * scale, 1), round(y0 * scale, 1),
round(x1 * scale, 1), round(y1 * scale, 1)])
return out
def one_pdf(args):
path, out_dir, dpi, max_pages, quality = args
tag = re.sub(r"[^0-9A-Za-z]+", "_", os.path.splitext(os.path.basename(path))[0])[:48]
rows, scale = [], dpi / 72.0
try:
doc = fitz.open(path)
except Exception as e:
return tag, 0, f"не открылся: {e}"
n = len(doc) if not max_pages else min(len(doc), max_pages)
for pno in range(n):
page = doc[pno]
boxes = page_boxes(page, scale)
if len(boxes) < 20: # обложки и полосы без текстового слоя
continue
pm = page.get_pixmap(dpi=dpi)
im = Image.frombytes("RGB", (pm.width, pm.height), pm.samples).convert("L")
name = f"{tag}_p{pno + 1:03d}.jpg"
im.save(os.path.join(out_dir, name), quality=quality, optimize=True)
rows.append({"f": name, "w": im.width, "h": im.height, "b": boxes})
doc.close()
if rows:
with open(os.path.join(out_dir, f"_part_{tag}.jsonl"), "w", encoding="utf-8") as fh:
for r in rows:
fh.write(json.dumps(r, ensure_ascii=False) + "\n")
return tag, len(rows), ""
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--src", required=True, help="каталог с PDF")
ap.add_argument("--out", required=True)
ap.add_argument("--dpi", type=int, default=150)
ap.add_argument("--pages", type=int, default=0, help="0 = все страницы выпуска")
ap.add_argument("--quality", type=int, default=85)
ap.add_argument("--workers", type=int, default=8)
ap.add_argument("--keep-held-out", action="store_true",
help="собрать, наоборот, только отложенные выпуски")
a = ap.parse_args()
os.makedirs(a.out, exist_ok=True)
files = sorted(f for f in os.listdir(a.src) if f.lower().endswith(".pdf"))
picked = [os.path.join(a.src, f) for f in files
if bool(HELD_OUT.search(f)) == a.keep_held_out]
print(f"PDF всего {len(files)}, берём {len(picked)}", flush=True)
tasks = [(p, a.out, a.dpi, a.pages, a.quality) for p in picked]
done = pages = 0
with ProcessPoolExecutor(max_workers=a.workers) as ex:
for tag, k, err in ex.map(one_pdf, tasks):
done += 1
pages += k
if err:
print(f" [{done}/{len(tasks)}] {tag}: {err}", flush=True)
elif done % 10 == 0:
print(f" [{done}/{len(tasks)}] страниц собрано {pages}", flush=True)
# склеиваем части: каждый воркер пишет свой файл, общий индекс собираем в конце
with open(os.path.join(a.out, "index.jsonl"), "w", encoding="utf-8") as fh:
for name in sorted(os.listdir(a.out)):
if not name.startswith("_part_"):
continue
p = os.path.join(a.out, name)
with open(p, encoding="utf-8") as src:
fh.write(src.read())
os.remove(p)
with open(os.path.join(a.out, "index.jsonl"), encoding="utf-8") as fh:
rows = sum(1 for _ in fh)
print(f"готово: страниц {rows} в {a.out}")
if __name__ == "__main__":
main()
|