|
|
@@ -53,32 +53,64 @@ IMG_EXTENSIONS: set[str] = {".jpg", ".jpeg", ".png", ".tiff", ".tif"}
|
|
53
|
53
|
|
|
54
|
54
|
def _get_ocr_engine(name: str, main_lang: str):
|
|
55
|
55
|
if name == "surya":
|
|
56
|
|
- from src.ocr.surya_engine import surya_ocr_image
|
|
57
|
|
- return surya_ocr_image, "Surya 2 VLM"
|
|
|
56
|
+ from src.ocr.surya_engine import surya_ocr_batch
|
|
|
57
|
+ return surya_ocr_batch, "Surya 2 VLM", True
|
|
58
|
58
|
if name == "external-qwen3":
|
|
59
|
59
|
from src.ocr.external_vlm_engine import external_vlm_ocr_image
|
|
60
|
|
- return external_vlm_ocr_image, "Qwen3.8 Max (external VLM)"
|
|
61
|
|
- return functools.partial(ocr_image, main_lang=main_lang), "PP-StructureV3"
|
|
|
60
|
+ return external_vlm_ocr_image, "Qwen3.8 Max (external VLM)", False
|
|
|
61
|
+ return functools.partial(ocr_image, main_lang=main_lang), "PP-StructureV3", False
|
|
62
|
62
|
|
|
63
|
63
|
|
|
64
|
|
-def _merge_markdown_files(output_dir: Path, md_files: list[Path]) -> None:
|
|
|
64
|
+def _merge_markdown_files(output_dir: Path, input_name: str, md_files: list[Path]) -> None:
|
|
65
|
65
|
merged_parts: list[str] = []
|
|
66
|
66
|
for md_path in sorted(md_files):
|
|
67
|
67
|
content = md_path.read_text(encoding="utf-8")
|
|
68
|
68
|
merged_parts.append(f"<!-- page: {md_path.stem} -->\n\n{content}")
|
|
69
|
69
|
if merged_parts:
|
|
70
|
70
|
merged = "# Merged document\n\n" + "\n\n---\n\n".join(merged_parts)
|
|
71
|
|
- merged_path = output_dir / (output_dir.name + MD_EXTENSION)
|
|
|
71
|
+ merged_path = output_dir / (input_name + MD_EXTENSION)
|
|
72
|
72
|
merged_path.write_text(merged, encoding="utf-8")
|
|
73
|
73
|
logger.info("Объединённый документ: %s", merged_path)
|
|
74
|
74
|
|
|
75
|
75
|
|
|
|
76
|
+def _log_low_confidence(raw_json: dict) -> None:
|
|
|
77
|
+ layout = raw_json.get("layout_det_res", {}).get("boxes", [])
|
|
|
78
|
+ if not layout:
|
|
|
79
|
+ return
|
|
|
80
|
+ low = [(b.get("label", "?"), round(b.get("score", 0), 3)) for b in layout if b.get("score", 1) < 0.5]
|
|
|
81
|
+ if low:
|
|
|
82
|
+ labels = ", ".join(f"{lbl}({sc})" for lbl, sc in low)
|
|
|
83
|
+ logger.warning("Низкая уверенность: %d блоков — %s", len(low), labels)
|
|
|
84
|
+
|
|
|
85
|
+
|
|
|
86
|
+def _apply_vlm(blocks, input_path: Path) -> None:
|
|
|
87
|
+ from src.postprocess.corrector import apply_corrector
|
|
|
88
|
+ from src.postprocess.vlm_corrector import vlm_correct_block
|
|
|
89
|
+ corrector = functools.partial(vlm_correct_block, image_path=str(input_path.resolve()))
|
|
|
90
|
+ apply_corrector(blocks, corrector, "VLM")
|
|
|
91
|
+
|
|
|
92
|
+
|
|
|
93
|
+def _save_result(page, input_path: Path, output_dir: Path) -> None:
|
|
|
94
|
+ json_path = output_dir / (input_path.stem + ".json")
|
|
|
95
|
+ if not json_path.exists():
|
|
|
96
|
+ json_payload = page.raw_json or {
|
|
|
97
|
+ "blocks": [{"label": b.label, "content": b.content, "bbox": list(b.bbox)}
|
|
|
98
|
+ for b in page.blocks]
|
|
|
99
|
+ }
|
|
|
100
|
+ json_path.write_text(json.dumps(json_payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
|
101
|
+ logger.debug("JSON сохранён: %s", json_path)
|
|
|
102
|
+
|
|
|
103
|
+ md_path = output_dir / (input_path.stem + MD_EXTENSION)
|
|
|
104
|
+ generate_markdown(page.blocks, md_path)
|
|
|
105
|
+ logger.debug("Markdown сохранён: %s", md_path)
|
|
|
106
|
+ logger.info("%s → %s, %s", input_path.name, md_path.name, json_path.name)
|
|
|
107
|
+
|
|
|
108
|
+
|
|
76
|
109
|
def _process_single(
|
|
77
|
110
|
ocr_fn, input_path: Path,
|
|
78
|
111
|
output_dir_resolved: Path, use_vlm: bool,
|
|
79
|
112
|
) -> None:
|
|
80
|
113
|
"""Обрабатывает одно изображение: OCR → fixups → корректоры → JSON + Markdown."""
|
|
81
|
|
- # Этап OCR — самое долгое и хрупкое место
|
|
82
|
114
|
logger.debug("OCR: начало %s", input_path.name)
|
|
83
|
115
|
t0 = time.time()
|
|
84
|
116
|
page = ocr_fn(str(input_path))
|
|
|
@@ -86,34 +118,15 @@ def _process_single(
|
|
86
|
118
|
logger.info("Распознано %d блоков за %.1fs", len(page.blocks), dt)
|
|
87
|
119
|
logger.debug("OCR: завершено %s (блоков: %d)", input_path.name, len(page.blocks))
|
|
88
|
120
|
|
|
89
|
|
- # Исправление типичных OCR-ошибок (VaR, V^2)
|
|
90
|
121
|
fix_ocr_errors(page.blocks)
|
|
91
|
122
|
logger.debug("Fixups: завершено")
|
|
92
|
123
|
|
|
93
|
|
- # Опциональная коррекция через VLM
|
|
|
124
|
+ _log_low_confidence(page.raw_json)
|
|
|
125
|
+
|
|
94
|
126
|
if use_vlm:
|
|
95
|
|
- from src.postprocess.corrector import apply_corrector
|
|
96
|
|
- logger.debug("Корректоры: инициализация")
|
|
97
|
|
- from src.postprocess.vlm_corrector import vlm_correct_block
|
|
98
|
|
- corrector = functools.partial(vlm_correct_block, image_path=str(input_path.resolve()))
|
|
99
|
|
- apply_corrector(page.blocks, corrector, "VLM")
|
|
100
|
|
- logger.debug("VLM: завершено")
|
|
101
|
|
-
|
|
102
|
|
- # Сохранение JSON (всегда) + Markdown (всегда)
|
|
103
|
|
- json_path = output_dir_resolved / (input_path.stem + ".json")
|
|
104
|
|
- if not json_path.exists():
|
|
105
|
|
- # Если raw_json пуст (Surya), строим из ParsedBlock
|
|
106
|
|
- json_payload = page.raw_json or {
|
|
107
|
|
- "blocks": [{"label": b.label, "content": b.content, "bbox": list(b.bbox)}
|
|
108
|
|
- for b in page.blocks]
|
|
109
|
|
- }
|
|
110
|
|
- json_path.write_text(json.dumps(json_payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
111
|
|
- logger.debug("JSON сохранён: %s", json_path)
|
|
|
127
|
+ _apply_vlm(page.blocks, input_path)
|
|
112
|
128
|
|
|
113
|
|
- md_path = output_dir_resolved / (input_path.stem + MD_EXTENSION)
|
|
114
|
|
- generate_markdown(page.blocks, md_path)
|
|
115
|
|
- logger.debug("Markdown сохранён: %s", md_path)
|
|
116
|
|
- logger.info("%s → %s, %s", input_path.name, md_path.name, json_path.name)
|
|
|
129
|
+ _save_result(page, input_path, output_dir_resolved)
|
|
117
|
130
|
|
|
118
|
131
|
|
|
119
|
132
|
@app.command()
|
|
|
@@ -180,7 +193,7 @@ def image_to_latex(
|
|
180
|
193
|
] = 5,
|
|
181
|
194
|
) -> None:
|
|
182
|
195
|
"""OCR → PostProcess → Markdown. Принимает файл или каталог изображений."""
|
|
183
|
|
- ocr_fn, engine_name = _get_ocr_engine(ocr_engine, main_lang)
|
|
|
196
|
+ ocr_fn, engine_name, is_batch = _get_ocr_engine(ocr_engine, main_lang)
|
|
184
|
197
|
|
|
185
|
198
|
if input.is_dir():
|
|
186
|
199
|
image_paths = sorted(
|
|
|
@@ -214,18 +227,30 @@ def image_to_latex(
|
|
214
|
227
|
|
|
215
|
228
|
total_start = time.time()
|
|
216
|
229
|
|
|
217
|
|
- for i, img_path in enumerate(image_paths):
|
|
218
|
|
- logger.info("[%d/%d] %s...", i + 1 + skipped, len(image_paths) + skipped, img_path.name)
|
|
219
|
|
- _process_single(ocr_fn, img_path, out, use_vlm)
|
|
220
|
|
- if i < len(image_paths) - 1 and pause > 0:
|
|
221
|
|
- time.sleep(pause)
|
|
|
230
|
+ if is_batch:
|
|
|
231
|
+ paths = [str(p) for p in image_paths]
|
|
|
232
|
+ logger.info("Batch-обработка %d изображений...", len(paths))
|
|
|
233
|
+ pages = ocr_fn(paths)
|
|
|
234
|
+ for img_path, page in zip(image_paths, pages):
|
|
|
235
|
+ logger.info("%s → %d блоков", img_path.name, len(page.blocks))
|
|
|
236
|
+ fix_ocr_errors(page.blocks)
|
|
|
237
|
+ _log_low_confidence(page.raw_json)
|
|
|
238
|
+ if use_vlm:
|
|
|
239
|
+ _apply_vlm(page.blocks, img_path)
|
|
|
240
|
+ _save_result(page, img_path, out)
|
|
|
241
|
+ else:
|
|
|
242
|
+ for i, img_path in enumerate(image_paths):
|
|
|
243
|
+ logger.info("[%d/%d] %s...", i + 1 + skipped, len(image_paths) + skipped, img_path.name)
|
|
|
244
|
+ _process_single(ocr_fn, img_path, out, use_vlm)
|
|
|
245
|
+ if i < len(image_paths) - 1 and pause > 0:
|
|
|
246
|
+ time.sleep(pause)
|
|
222
|
247
|
|
|
223
|
248
|
total_dt = time.time() - total_start
|
|
224
|
249
|
processed = len(image_paths)
|
|
225
|
250
|
avg_dt = total_dt / processed if processed else 0
|
|
226
|
251
|
|
|
227
|
252
|
if result_one_document:
|
|
228
|
|
- _merge_markdown_files(out, sorted(out.glob("*.md")))
|
|
|
253
|
+ _merge_markdown_files(out, input.name, sorted(out.glob("*.md")))
|
|
229
|
254
|
|
|
230
|
255
|
logger.info(
|
|
231
|
256
|
"Готово: %d изображений, всего %.1fs, среднее %.1fs/изобр",
|