visual-parser 2.2.0__py3-none-any.whl → 2.2.2__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- visual_parser/__init__.py +1 -1
- visual_parser/cli.py +51 -9
- visual_parser/cli_main.py +50 -9
- visual_parser/config.py +37 -2
- visual_parser/figure_describer.py +14 -1
- visual_parser/image_describer.py +329 -0
- visual_parser/image_tracker.py +71 -0
- visual_parser/jsonl_writer.py +64 -3
- visual_parser/kb_redo.py +148 -0
- visual_parser/metadata_extractor.py +87 -13
- visual_parser/pdf_tracker.py +25 -0
- visual_parser/pipeline.py +139 -24
- visual_parser/prompts.py +78 -0
- visual_parser/text_extractor.py +20 -29
- visual_parser/vision_llm.py +46 -22
- {visual_parser-2.2.0.dist-info → visual_parser-2.2.2.dist-info}/METADATA +9 -5
- visual_parser-2.2.2.dist-info/RECORD +26 -0
- visual_parser-2.2.0.dist-info/RECORD +0 -23
- {visual_parser-2.2.0.dist-info → visual_parser-2.2.2.dist-info}/WHEEL +0 -0
- {visual_parser-2.2.0.dist-info → visual_parser-2.2.2.dist-info}/entry_points.txt +0 -0
- {visual_parser-2.2.0.dist-info → visual_parser-2.2.2.dist-info}/licenses/LICENSE +0 -0
- {visual_parser-2.2.0.dist-info → visual_parser-2.2.2.dist-info}/top_level.txt +0 -0
visual_parser/__init__.py
CHANGED
visual_parser/cli.py
CHANGED
|
@@ -44,11 +44,13 @@ def _build_arg_parser() -> argparse.ArgumentParser:
|
|
|
44
44
|
p = argparse.ArgumentParser(
|
|
45
45
|
prog="visual-parser",
|
|
46
46
|
description=(
|
|
47
|
-
"Visual-RAG PDF Parser — detects new PDFs
|
|
48
|
-
"figure descriptions, and writes
|
|
49
|
-
"
|
|
50
|
-
"
|
|
51
|
-
"
|
|
47
|
+
"Visual-RAG PDF Parser — detects new PDFs and standalone images, "
|
|
48
|
+
"extracts text and figure/image descriptions, and writes JSONL "
|
|
49
|
+
"knowledge bases:\n"
|
|
50
|
+
" 01_chunks_kb.jsonl text chunks\n"
|
|
51
|
+
" 02_visuals_kb.jsonl PDF figure descriptions\n"
|
|
52
|
+
" 03_metadata_kb.jsonl document metadata\n"
|
|
53
|
+
" image_descriptions.jsonl standalone image descriptions"
|
|
52
54
|
),
|
|
53
55
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
54
56
|
epilog=USAGE_EXAMPLES,
|
|
@@ -67,8 +69,9 @@ def _build_arg_parser() -> argparse.ArgumentParser:
|
|
|
67
69
|
"--input-dir", "-i",
|
|
68
70
|
required=False,
|
|
69
71
|
metavar="DIR",
|
|
70
|
-
help="Directory to scan for PDF files
|
|
71
|
-
"
|
|
72
|
+
help="Directory to scan for PDF files and standalone images "
|
|
73
|
+
"(searched recursively, auto-detected by extension -- a single "
|
|
74
|
+
"directory may mix both). Required unless --list-models is given.",
|
|
72
75
|
)
|
|
73
76
|
io_group.add_argument(
|
|
74
77
|
"--output-dir", "-o",
|
|
@@ -108,6 +111,15 @@ def _build_arg_parser() -> argparse.ArgumentParser:
|
|
|
108
111
|
metavar="N",
|
|
109
112
|
help="Overlap characters between adjacent chunks (default: 100).",
|
|
110
113
|
)
|
|
114
|
+
text_group.add_argument(
|
|
115
|
+
"--nougat-crop-margin",
|
|
116
|
+
action="store_true",
|
|
117
|
+
help=(
|
|
118
|
+
"Let Nougat autocrop whitespace margins (its own default). "
|
|
119
|
+
"Off by default -- autocropping can silently crop into real "
|
|
120
|
+
"content on unusual page layouts."
|
|
121
|
+
),
|
|
122
|
+
)
|
|
111
123
|
|
|
112
124
|
# ---- Vision LLM ---------------------------------------------------------
|
|
113
125
|
vision_group = p.add_argument_group("Vision LLM (figure descriptions & metadata)")
|
|
@@ -204,8 +216,22 @@ def _build_arg_parser() -> argparse.ArgumentParser:
|
|
|
204
216
|
"--rebuild",
|
|
205
217
|
action="store_true",
|
|
206
218
|
help=(
|
|
207
|
-
"
|
|
208
|
-
"Use after changing prompts, chunking strategy, or switching models."
|
|
219
|
+
"Clear and reprocess every PDF/image currently in --input-dir. "
|
|
220
|
+
"Use after changing prompts, chunking strategy, or switching models. "
|
|
221
|
+
"No duplicate rows -- old entries are removed before reprocessing."
|
|
222
|
+
),
|
|
223
|
+
)
|
|
224
|
+
misc_group.add_argument(
|
|
225
|
+
"--redo",
|
|
226
|
+
action="append",
|
|
227
|
+
default=None,
|
|
228
|
+
metavar="NAME",
|
|
229
|
+
help=(
|
|
230
|
+
"Clear ALL existing KB rows and registry entries for this PDF or "
|
|
231
|
+
"image basename (case-sensitive, path-stripped), then let normal "
|
|
232
|
+
"discovery reprocess it in this same run -- no duplicates. "
|
|
233
|
+
"Repeatable: --redo doc1.pdf --redo diagram2.png. Combine with "
|
|
234
|
+
"--rebuild to clear+reprocess everything in --input-dir instead."
|
|
209
235
|
),
|
|
210
236
|
)
|
|
211
237
|
misc_group.add_argument(
|
|
@@ -226,6 +252,19 @@ def _build_arg_parser() -> argparse.ArgumentParser:
|
|
|
226
252
|
help="Verbosity level written to 05_pipeline.log (default: ERROR).",
|
|
227
253
|
)
|
|
228
254
|
|
|
255
|
+
# ---- Image processing ----------------------------------------------------
|
|
256
|
+
image_group = p.add_argument_group("Image processing")
|
|
257
|
+
image_group.add_argument(
|
|
258
|
+
"--skip-images",
|
|
259
|
+
action="store_true",
|
|
260
|
+
help=(
|
|
261
|
+
"Skip standalone-image processing entirely, even if image files "
|
|
262
|
+
"(.png/.jpg/.jpeg/.tif/.tiff/.bmp) are found in --input-dir. "
|
|
263
|
+
"Reuses --vision-provider/--vision-model/--vision-detail/"
|
|
264
|
+
"--reasoning-effort/--max-workers -- no separate image vision flags."
|
|
265
|
+
),
|
|
266
|
+
)
|
|
267
|
+
|
|
229
268
|
return p
|
|
230
269
|
|
|
231
270
|
|
|
@@ -336,6 +375,7 @@ def main(argv=None) -> int:
|
|
|
336
375
|
nougat_model = args.nougat_model,
|
|
337
376
|
chunk_size = args.chunk_size,
|
|
338
377
|
chunk_overlap = args.chunk_overlap,
|
|
378
|
+
nougat_crop_margin = args.nougat_crop_margin,
|
|
339
379
|
vision_provider = args.vision_provider,
|
|
340
380
|
gpt_vision_model = args.vision_model if args.vision_provider == "gpt" else "gpt-5.4",
|
|
341
381
|
gemini_vision_model = args.vision_model if args.vision_provider == "gemini" else "gemini-3.8-flash",
|
|
@@ -346,7 +386,9 @@ def main(argv=None) -> int:
|
|
|
346
386
|
vision_context_pages = args.vision_context_pages,
|
|
347
387
|
max_workers = args.max_workers,
|
|
348
388
|
rebuild = args.rebuild,
|
|
389
|
+
redo_names = args.redo or [],
|
|
349
390
|
skip_text = args.skip_text,
|
|
391
|
+
skip_images = args.skip_images,
|
|
350
392
|
log_level = args.log_level,
|
|
351
393
|
)
|
|
352
394
|
|
visual_parser/cli_main.py
CHANGED
|
@@ -42,11 +42,13 @@ def _build_arg_parser() -> argparse.ArgumentParser:
|
|
|
42
42
|
parser = argparse.ArgumentParser(
|
|
43
43
|
prog="visual-parser",
|
|
44
44
|
description=(
|
|
45
|
-
"Visual-RAG PDF Parser - detects new PDFs
|
|
46
|
-
"figure descriptions, and writes
|
|
47
|
-
"
|
|
48
|
-
"
|
|
49
|
-
"
|
|
45
|
+
"Visual-RAG PDF Parser - detects new PDFs and standalone images, "
|
|
46
|
+
"extracts text and figure/image descriptions, and writes JSONL "
|
|
47
|
+
"knowledge bases:\n"
|
|
48
|
+
" 01_chunks_kb.jsonl text chunks\n"
|
|
49
|
+
" 02_visuals_kb.jsonl PDF figure descriptions\n"
|
|
50
|
+
" 03_metadata_kb.jsonl document metadata\n"
|
|
51
|
+
" image_descriptions.jsonl standalone image descriptions"
|
|
50
52
|
),
|
|
51
53
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
52
54
|
epilog=USAGE_EXAMPLES,
|
|
@@ -65,8 +67,9 @@ def _build_arg_parser() -> argparse.ArgumentParser:
|
|
|
65
67
|
"-i",
|
|
66
68
|
required=False,
|
|
67
69
|
metavar="DIR",
|
|
68
|
-
help="Directory to scan for PDF files
|
|
69
|
-
"
|
|
70
|
+
help="Directory to scan for PDF files and standalone images "
|
|
71
|
+
"(searched recursively, auto-detected by extension -- a single "
|
|
72
|
+
"directory may mix both). Required unless --list-models is given.",
|
|
70
73
|
)
|
|
71
74
|
io_group.add_argument(
|
|
72
75
|
"--output-dir",
|
|
@@ -106,6 +109,15 @@ def _build_arg_parser() -> argparse.ArgumentParser:
|
|
|
106
109
|
metavar="N",
|
|
107
110
|
help="Overlap characters between adjacent chunks (default: 100).",
|
|
108
111
|
)
|
|
112
|
+
text_group.add_argument(
|
|
113
|
+
"--nougat-crop-margin",
|
|
114
|
+
action="store_true",
|
|
115
|
+
help=(
|
|
116
|
+
"Let Nougat autocrop whitespace margins (its own default). "
|
|
117
|
+
"Off by default -- autocropping can silently crop into real "
|
|
118
|
+
"content on unusual page layouts."
|
|
119
|
+
),
|
|
120
|
+
)
|
|
109
121
|
|
|
110
122
|
vision_group = parser.add_argument_group("Vision LLM (figure descriptions & metadata)")
|
|
111
123
|
vision_group.add_argument(
|
|
@@ -174,6 +186,18 @@ def _build_arg_parser() -> argparse.ArgumentParser:
|
|
|
174
186
|
help="Number of front pages sent to the vision LLM for metadata extraction (default: 2).",
|
|
175
187
|
)
|
|
176
188
|
|
|
189
|
+
image_group = parser.add_argument_group("Image processing")
|
|
190
|
+
image_group.add_argument(
|
|
191
|
+
"--skip-images",
|
|
192
|
+
action="store_true",
|
|
193
|
+
help=(
|
|
194
|
+
"Skip standalone-image processing entirely, even if image files "
|
|
195
|
+
"(.png/.jpg/.jpeg/.tif/.tiff/.bmp) are found in --input-dir. "
|
|
196
|
+
"Reuses --vision-provider/--vision-model/--vision-detail/"
|
|
197
|
+
"--reasoning-effort/--max-workers -- no separate image vision flags."
|
|
198
|
+
),
|
|
199
|
+
)
|
|
200
|
+
|
|
177
201
|
perf_group = parser.add_argument_group("Performance")
|
|
178
202
|
perf_group.add_argument(
|
|
179
203
|
"--max-workers",
|
|
@@ -199,8 +223,22 @@ def _build_arg_parser() -> argparse.ArgumentParser:
|
|
|
199
223
|
"--rebuild",
|
|
200
224
|
action="store_true",
|
|
201
225
|
help=(
|
|
202
|
-
"
|
|
203
|
-
"Use after changing prompts, chunking strategy, or switching models."
|
|
226
|
+
"Clear and reprocess every PDF/image currently in --input-dir. "
|
|
227
|
+
"Use after changing prompts, chunking strategy, or switching models. "
|
|
228
|
+
"No duplicate rows -- old entries are removed before reprocessing."
|
|
229
|
+
),
|
|
230
|
+
)
|
|
231
|
+
misc_group.add_argument(
|
|
232
|
+
"--redo",
|
|
233
|
+
action="append",
|
|
234
|
+
default=None,
|
|
235
|
+
metavar="NAME",
|
|
236
|
+
help=(
|
|
237
|
+
"Clear ALL existing KB rows and registry entries for this PDF or "
|
|
238
|
+
"image basename (case-sensitive, path-stripped), then let normal "
|
|
239
|
+
"discovery reprocess it in this same run -- no duplicates. "
|
|
240
|
+
"Repeatable: --redo doc1.pdf --redo diagram2.png. Combine with "
|
|
241
|
+
"--rebuild to clear+reprocess everything in --input-dir instead."
|
|
204
242
|
),
|
|
205
243
|
)
|
|
206
244
|
misc_group.add_argument(
|
|
@@ -253,6 +291,7 @@ def main(argv=None) -> int:
|
|
|
253
291
|
nougat_model=args.nougat_model,
|
|
254
292
|
chunk_size=args.chunk_size,
|
|
255
293
|
chunk_overlap=args.chunk_overlap,
|
|
294
|
+
nougat_crop_margin=args.nougat_crop_margin,
|
|
256
295
|
vision_provider=args.vision_provider,
|
|
257
296
|
gpt_vision_model=args.vision_model if args.vision_provider == "gpt" else "gpt-5.4",
|
|
258
297
|
gemini_vision_model=(
|
|
@@ -265,7 +304,9 @@ def main(argv=None) -> int:
|
|
|
265
304
|
vision_context_pages=args.vision_context_pages,
|
|
266
305
|
max_workers=args.max_workers,
|
|
267
306
|
rebuild=args.rebuild,
|
|
307
|
+
redo_names=args.redo or [],
|
|
268
308
|
skip_text=args.skip_text,
|
|
309
|
+
skip_images=args.skip_images,
|
|
269
310
|
log_level=args.log_level,
|
|
270
311
|
)
|
|
271
312
|
|
visual_parser/config.py
CHANGED
|
@@ -12,7 +12,7 @@ from __future__ import annotations
|
|
|
12
12
|
import os
|
|
13
13
|
from dataclasses import dataclass, field
|
|
14
14
|
from pathlib import Path
|
|
15
|
-
from typing import Literal, Optional
|
|
15
|
+
from typing import List, Literal, Optional
|
|
16
16
|
|
|
17
17
|
from dotenv import load_dotenv
|
|
18
18
|
|
|
@@ -79,6 +79,14 @@ class ParserConfig:
|
|
|
79
79
|
nougat_model: str = "facebook/nougat-small"
|
|
80
80
|
"""HuggingFace model identifier for Nougat."""
|
|
81
81
|
|
|
82
|
+
nougat_crop_margin: bool = False
|
|
83
|
+
"""
|
|
84
|
+
Passed through to Nougat's image processor as do_crop_margin. False
|
|
85
|
+
(default) matches this package's long-standing behavior -- autocropping
|
|
86
|
+
whitespace margins can silently crop into real content on unusual page
|
|
87
|
+
layouts. Set True to opt into Nougat's own autocrop behavior instead.
|
|
88
|
+
"""
|
|
89
|
+
|
|
82
90
|
chunk_size: int = 500
|
|
83
91
|
"""Target character count per text chunk."""
|
|
84
92
|
|
|
@@ -139,7 +147,23 @@ class ParserConfig:
|
|
|
139
147
|
|
|
140
148
|
# --- Misc ----------------------------------------------------------------
|
|
141
149
|
rebuild: bool = False
|
|
142
|
-
"""
|
|
150
|
+
"""
|
|
151
|
+
If True, clear and reprocess every PDF/image currently in input_dir,
|
|
152
|
+
regardless of the registries -- implemented via kb_redo.redo_entries()
|
|
153
|
+
(clear KB rows + registry entries first, then let normal discovery pick
|
|
154
|
+
everything back up as new), so this produces no duplicate rows.
|
|
155
|
+
"""
|
|
156
|
+
|
|
157
|
+
redo_names: List[str] = field(default_factory=list)
|
|
158
|
+
"""
|
|
159
|
+
PDF or image basenames (case-sensitive, path-stripped) to clear from
|
|
160
|
+
every KB JSONL file and both registries before this run's discovery
|
|
161
|
+
step, via kb_redo.redo_entries() -- so the normal pipeline naturally
|
|
162
|
+
reprocesses them cleanly in this same run, with no duplicate rows,
|
|
163
|
+
unlike blindly reprocessing over an already-populated KB. A name that's
|
|
164
|
+
no longer present in input_dir degrades gracefully to a clean removal
|
|
165
|
+
(nothing left to reprocess).
|
|
166
|
+
"""
|
|
143
167
|
|
|
144
168
|
skip_text: bool = False
|
|
145
169
|
"""
|
|
@@ -150,6 +174,14 @@ class ParserConfig:
|
|
|
150
174
|
03_metadata_kb.jsonl are skipped automatically.
|
|
151
175
|
"""
|
|
152
176
|
|
|
177
|
+
skip_images: bool = False
|
|
178
|
+
"""
|
|
179
|
+
If True, skip standalone-image processing entirely, even if image files
|
|
180
|
+
are found in input_dir. Mirrors skip_text's opt-out pattern -- an escape
|
|
181
|
+
hatch for existing PDF-only workflows that don't want stray image files
|
|
182
|
+
sitting in input_dir to suddenly start being processed.
|
|
183
|
+
"""
|
|
184
|
+
|
|
153
185
|
log_level: str = "ERROR"
|
|
154
186
|
|
|
155
187
|
# -------------------------------------------------------------------------
|
|
@@ -162,6 +194,7 @@ class ParserConfig:
|
|
|
162
194
|
output_dir = os.getenv("VISUAL_PARSER_OUTPUT_DIR", ""),
|
|
163
195
|
text_mode = os.getenv("VISUAL_PARSER_TEXT_MODE", "nougat"), # type: ignore[arg-type]
|
|
164
196
|
nougat_model = os.getenv("VISUAL_PARSER_NOUGAT_MODEL", "facebook/nougat-small"),
|
|
197
|
+
nougat_crop_margin = os.getenv("VISUAL_PARSER_NOUGAT_CROP_MARGIN", "false").lower() == "true",
|
|
165
198
|
chunk_size = int(os.getenv("VISUAL_PARSER_CHUNK_SIZE", "500")),
|
|
166
199
|
chunk_overlap = int(os.getenv("VISUAL_PARSER_CHUNK_OVERLAP", "100")),
|
|
167
200
|
vision_provider = os.getenv("VISUAL_PARSER_VISION_PROVIDER", "gpt"), # type: ignore[arg-type]
|
|
@@ -176,7 +209,9 @@ class ParserConfig:
|
|
|
176
209
|
vision_context_pages = int(os.getenv("VISUAL_PARSER_VISION_CONTEXT_PAGES", "0")),
|
|
177
210
|
max_workers = int(os.getenv("VISUAL_PARSER_MAX_WORKERS", "4")),
|
|
178
211
|
rebuild = os.getenv("VISUAL_PARSER_REBUILD", "false").lower() == "true",
|
|
212
|
+
redo_names = [n.strip() for n in os.getenv("VISUAL_PARSER_REDO_NAMES", "").split(",") if n.strip()],
|
|
179
213
|
skip_text = os.getenv("VISUAL_PARSER_SKIP_TEXT", "false").lower() == "true",
|
|
214
|
+
skip_images = os.getenv("VISUAL_PARSER_SKIP_IMAGES", "false").lower() == "true",
|
|
180
215
|
log_level = os.getenv("VISUAL_PARSER_LOG_LEVEL", "ERROR"),
|
|
181
216
|
)
|
|
182
217
|
|
|
@@ -117,6 +117,13 @@ def describe_figures_for_new_pdfs(
|
|
|
117
117
|
0 (default) reproduces the exact single-page
|
|
118
118
|
behavior this function had before this parameter
|
|
119
119
|
existed.
|
|
120
|
+
|
|
121
|
+
Always resumes from 02_visuals_kb.jsonl's existing content (page-level,
|
|
122
|
+
via the done_pages set below) -- callers that want a page redescribed
|
|
123
|
+
must first remove its existing row(s) (see kb_redo.redo_entries()) and
|
|
124
|
+
its basename from the processed-PDFs registry, then pass it in here as
|
|
125
|
+
if new; this function has no bypass of its own, so there is no path to
|
|
126
|
+
producing a duplicate row.
|
|
120
127
|
"""
|
|
121
128
|
# -----------------------------------------------------------------------
|
|
122
129
|
# Step 1 – Rasterise every page of every new PDF
|
|
@@ -249,7 +256,13 @@ def describe_figures_for_new_pdfs(
|
|
|
249
256
|
if not isinstance(caption, dict):
|
|
250
257
|
continue
|
|
251
258
|
description = caption.get("description")
|
|
252
|
-
if description
|
|
259
|
+
if not isinstance(description, str):
|
|
260
|
+
logger.warning(
|
|
261
|
+
"Figure %d on %s page %d has a non-string 'description' "
|
|
262
|
+
"(got %s) -- model likely returned a malformed/nested "
|
|
263
|
+
"response; skipping.",
|
|
264
|
+
fig_idx, pdf_name, page_number, type(description).__name__,
|
|
265
|
+
)
|
|
253
266
|
continue
|
|
254
267
|
page_rows.append({
|
|
255
268
|
"source": pdf_name,
|
|
@@ -0,0 +1,329 @@
|
|
|
1
|
+
"""
|
|
2
|
+
image_describer.py — Send each standalone image to a Vision LLM for a
|
|
3
|
+
holistic description, one call per image (no page
|
|
4
|
+
rasterisation, no context window).
|
|
5
|
+
|
|
6
|
+
Extracted-pattern sibling of figure_describer.py: same ThreadPoolExecutor +
|
|
7
|
+
as_completed concurrency shape (worker only computes/returns; only the single-
|
|
8
|
+
threaded completion loop calls append_to_jsonl), but for arbitrary user-
|
|
9
|
+
supplied image files instead of PyMuPDF-rasterized PDF pages.
|
|
10
|
+
|
|
11
|
+
Output
|
|
12
|
+
------
|
|
13
|
+
One record per image is appended to ``image_descriptions.jsonl`` in
|
|
14
|
+
*output_dir*:
|
|
15
|
+
|
|
16
|
+
{
|
|
17
|
+
"source": "diagram1.png",
|
|
18
|
+
"image_id": "a1b2c3d4e5f6g7h8",
|
|
19
|
+
"title": "Vertical Parabolic Gate Schematic",
|
|
20
|
+
"description": "**Subject:** ..."
|
|
21
|
+
}
|
|
22
|
+
"""
|
|
23
|
+
|
|
24
|
+
from __future__ import annotations
|
|
25
|
+
|
|
26
|
+
import json
|
|
27
|
+
import logging
|
|
28
|
+
import os
|
|
29
|
+
import re
|
|
30
|
+
from collections import defaultdict
|
|
31
|
+
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
32
|
+
from io import BytesIO
|
|
33
|
+
from typing import Any, Dict, List, Optional
|
|
34
|
+
|
|
35
|
+
from PIL import Image
|
|
36
|
+
|
|
37
|
+
from visual_parser.jsonl_writer import append_to_jsonl, make_document_id
|
|
38
|
+
from visual_parser.prompts import IMAGE_DESCRIPTION_PROMPT
|
|
39
|
+
from visual_parser.vision_llm import call_vision_llm
|
|
40
|
+
|
|
41
|
+
logger = logging.getLogger(__name__)
|
|
42
|
+
|
|
43
|
+
IMAGE_DESCRIPTIONS_FILE = "image_descriptions.jsonl"
|
|
44
|
+
|
|
45
|
+
# Decoupled from figure_describer.py's context-window-scaled formula (10 min
|
|
46
|
+
# for 1 image there): IMAGE_DESCRIPTION_PROMPT's "Verbatim Text Content"
|
|
47
|
+
# requirement can demand much more generated output than FIGURE_PROMPT ever
|
|
48
|
+
# needs, since it asks for full verbatim transcription rather than a bounded
|
|
49
|
+
# structured description. Measured directly against this model/hardware: a
|
|
50
|
+
# dense, text-heavy page took 594s (9.9 min) for a single, uncontended call --
|
|
51
|
+
# real concurrent runs (multiple images via --max-workers) push slower still.
|
|
52
|
+
# 900s (15 min) calibrated from that real baseline, not guessed.
|
|
53
|
+
_OLLAMA_IMAGE_TIMEOUT_SECONDS = 900.0
|
|
54
|
+
|
|
55
|
+
_SUBJECT_RE = re.compile(r"\*\*Subject:\*\*\s*(.*?)(?:\n|$)")
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
# ---------------------------------------------------------------------------
|
|
59
|
+
# JSON response parser
|
|
60
|
+
# ---------------------------------------------------------------------------
|
|
61
|
+
|
|
62
|
+
def _parse_single_image_response(raw: str, image_name: str) -> Optional[Dict]:
|
|
63
|
+
"""
|
|
64
|
+
Parse the Vision LLM's single-JSON-object response for one image.
|
|
65
|
+
|
|
66
|
+
A sibling of figure_describer.py's _parse_llm_response: that function's
|
|
67
|
+
prose-recovery regex is list-shaped (looks for "[...]") and cannot
|
|
68
|
+
recover a bare "{...}" object wrapped in prose, which is the expected
|
|
69
|
+
shape here (one image always yields exactly one description, never a
|
|
70
|
+
list). Also defensively unwraps the inverse case -- a model that ignores
|
|
71
|
+
"NOT a list" and returns a single-element list anyway.
|
|
72
|
+
"""
|
|
73
|
+
body = raw.strip()
|
|
74
|
+
for fence in ("```json", "```"):
|
|
75
|
+
body = body.replace(fence, "")
|
|
76
|
+
|
|
77
|
+
parsed: Any = None
|
|
78
|
+
try:
|
|
79
|
+
parsed = json.loads(body)
|
|
80
|
+
except json.JSONDecodeError:
|
|
81
|
+
# A regex (greedy or lazy) can't reliably find "the" {...} block: the
|
|
82
|
+
# prompt's own OUTPUT FORMAT section embeds a literal, syntactically
|
|
83
|
+
# valid example object, so a model that echoes it alongside its real
|
|
84
|
+
# answer produces two {...} blocks -- greedy spans both into one
|
|
85
|
+
# invalid blob, while lazy stops at the first literal '}' even if
|
|
86
|
+
# that's just a brace character sitting inside a string value (e.g.
|
|
87
|
+
# a description that quotes set notation like "{A, B, C}"). Properly
|
|
88
|
+
# JSON-aware: try decoding starting at each '{' in turn and take the
|
|
89
|
+
# first one that parses as valid JSON, which correctly treats braces
|
|
90
|
+
# inside quoted strings as plain characters, not structure.
|
|
91
|
+
decoder = json.JSONDecoder()
|
|
92
|
+
start = body.find("{")
|
|
93
|
+
while start != -1:
|
|
94
|
+
try:
|
|
95
|
+
parsed, _ = decoder.raw_decode(body, start)
|
|
96
|
+
break
|
|
97
|
+
except json.JSONDecodeError:
|
|
98
|
+
start = body.find("{", start + 1)
|
|
99
|
+
|
|
100
|
+
if isinstance(parsed, list) and len(parsed) == 1:
|
|
101
|
+
parsed = parsed[0]
|
|
102
|
+
|
|
103
|
+
if isinstance(parsed, dict):
|
|
104
|
+
return parsed
|
|
105
|
+
|
|
106
|
+
logger.warning("Could not parse JSON for image %s: %r", image_name, body[:200])
|
|
107
|
+
return None
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
def _flatten_to_rgb(pil_image: Image.Image) -> Image.Image:
|
|
111
|
+
"""
|
|
112
|
+
Convert *pil_image* to RGB, compositing any transparency onto a white
|
|
113
|
+
background first.
|
|
114
|
+
|
|
115
|
+
Image.convert("RGB") alone does NOT do this -- it just drops the alpha
|
|
116
|
+
channel and exposes whatever RGB values are stored underneath, which are
|
|
117
|
+
commonly zeroed by PNG encoders for fully-transparent pixels. A
|
|
118
|
+
transparent-background PNG (a very common export format for standalone
|
|
119
|
+
diagrams/logos/screenshots -- exactly this feature's target content)
|
|
120
|
+
would silently turn mostly/fully black with no exception or log, giving
|
|
121
|
+
the vision LLM a black image instead of a diagram on a light background.
|
|
122
|
+
"""
|
|
123
|
+
has_alpha = pil_image.mode in ("RGBA", "LA") or (
|
|
124
|
+
pil_image.mode == "P" and "transparency" in pil_image.info
|
|
125
|
+
)
|
|
126
|
+
if not has_alpha:
|
|
127
|
+
return pil_image.convert("RGB")
|
|
128
|
+
|
|
129
|
+
rgba = pil_image.convert("RGBA")
|
|
130
|
+
background = Image.new("RGB", rgba.size, (255, 255, 255))
|
|
131
|
+
background.paste(rgba, mask=rgba.split()[-1])
|
|
132
|
+
return background
|
|
133
|
+
|
|
134
|
+
|
|
135
|
+
def _derive_title(description: str) -> str:
|
|
136
|
+
"""
|
|
137
|
+
Derive a concise title from *description*.
|
|
138
|
+
|
|
139
|
+
The prompt's heading 1 is always "**Subject:** [Title]" -- extract it
|
|
140
|
+
directly rather than truncating raw characters. Falls back to the first
|
|
141
|
+
sentence boundary, then a hard cut, if the model didn't follow the format.
|
|
142
|
+
"""
|
|
143
|
+
match = _SUBJECT_RE.search(description)
|
|
144
|
+
if match:
|
|
145
|
+
subject = match.group(1).strip()
|
|
146
|
+
if subject:
|
|
147
|
+
return subject
|
|
148
|
+
|
|
149
|
+
sentence_end = description.find(". ")
|
|
150
|
+
if 0 < sentence_end <= 150:
|
|
151
|
+
return description[:sentence_end + 1].strip()
|
|
152
|
+
|
|
153
|
+
return description[:100].strip()
|
|
154
|
+
|
|
155
|
+
|
|
156
|
+
# ---------------------------------------------------------------------------
|
|
157
|
+
# Main function
|
|
158
|
+
# ---------------------------------------------------------------------------
|
|
159
|
+
|
|
160
|
+
def describe_images(
|
|
161
|
+
new_image_paths: List[str],
|
|
162
|
+
output_dir: str,
|
|
163
|
+
vision_provider: str,
|
|
164
|
+
vision_api_key: str,
|
|
165
|
+
vision_model: str,
|
|
166
|
+
vision_detail: str = "low",
|
|
167
|
+
image_prompt: str = IMAGE_DESCRIPTION_PROMPT,
|
|
168
|
+
reasoning_effort: Optional[str] = "medium",
|
|
169
|
+
max_workers: int = 4,
|
|
170
|
+
) -> List[str]:
|
|
171
|
+
"""
|
|
172
|
+
For each image in *new_image_paths*, call the Vision LLM once (no
|
|
173
|
+
context window -- a single standalone image per call), parse the
|
|
174
|
+
description, and append the result to ``image_descriptions.jsonl`` in
|
|
175
|
+
*output_dir*.
|
|
176
|
+
|
|
177
|
+
Returns the basenames actually written a row for -- NOT the full input
|
|
178
|
+
list. Unlike PyMuPDF-rasterized PDF pages (always well-formed), arbitrary
|
|
179
|
+
user-supplied image files can be corrupt, truncated, or unreadable, and a
|
|
180
|
+
failed vision call must not be marked as processed by the caller, or the
|
|
181
|
+
image would silently disappear from future find_new_images() output with
|
|
182
|
+
no row ever written.
|
|
183
|
+
|
|
184
|
+
Args:
|
|
185
|
+
new_image_paths: Full paths of images to describe.
|
|
186
|
+
output_dir: Directory where ``image_descriptions.jsonl`` is written.
|
|
187
|
+
vision_provider: ``'gpt'``, ``'gemini'``, or ``'ollama'``.
|
|
188
|
+
vision_api_key: API key for the chosen provider (unused for ollama).
|
|
189
|
+
vision_model: Vision model name string.
|
|
190
|
+
vision_detail: Image detail level (GPT/Ollama only).
|
|
191
|
+
image_prompt: The instruction prompt sent with each image.
|
|
192
|
+
max_workers: Thread-pool size for concurrent vision-LLM calls.
|
|
193
|
+
|
|
194
|
+
Always resumes from image_descriptions.jsonl's existing content (via the
|
|
195
|
+
done_images set below) -- callers that want an image redescribed must
|
|
196
|
+
first remove its existing row and its basename from the processed-
|
|
197
|
+
images registry (see kb_redo.redo_entries()), then pass it in here as if
|
|
198
|
+
new; this function has no bypass of its own, so there is no path to
|
|
199
|
+
producing a duplicate row.
|
|
200
|
+
"""
|
|
201
|
+
descriptions_path = os.path.join(output_dir, IMAGE_DESCRIPTIONS_FILE)
|
|
202
|
+
|
|
203
|
+
# -------------------------------------------------------------------
|
|
204
|
+
# Build set of images already on disk (resume safety).
|
|
205
|
+
# -------------------------------------------------------------------
|
|
206
|
+
done_images: set = set()
|
|
207
|
+
if os.path.exists(descriptions_path):
|
|
208
|
+
with open(descriptions_path, encoding="utf-8") as _fh:
|
|
209
|
+
for _line in _fh:
|
|
210
|
+
_line = _line.strip()
|
|
211
|
+
if not _line:
|
|
212
|
+
continue
|
|
213
|
+
try:
|
|
214
|
+
_rec = json.loads(_line)
|
|
215
|
+
_src = _rec.get("source", "")
|
|
216
|
+
if _src:
|
|
217
|
+
done_images.add(_src)
|
|
218
|
+
except Exception:
|
|
219
|
+
pass
|
|
220
|
+
if done_images:
|
|
221
|
+
logger.info(
|
|
222
|
+
"Resuming: %d image(s) already in %s — will skip.",
|
|
223
|
+
len(done_images), IMAGE_DESCRIPTIONS_FILE,
|
|
224
|
+
)
|
|
225
|
+
|
|
226
|
+
tasks = [
|
|
227
|
+
{"path": path, "name": os.path.basename(path)}
|
|
228
|
+
for path in new_image_paths
|
|
229
|
+
if os.path.basename(path) not in done_images
|
|
230
|
+
]
|
|
231
|
+
|
|
232
|
+
if not tasks:
|
|
233
|
+
logger.info("No new images to describe — all already processed.")
|
|
234
|
+
return []
|
|
235
|
+
|
|
236
|
+
# -------------------------------------------------------------------
|
|
237
|
+
# Call Vision LLM per task, concurrently. Each worker only computes
|
|
238
|
+
# and returns a result; only the single-threaded completion loop below
|
|
239
|
+
# calls append_to_jsonl, so the file is never written to from more
|
|
240
|
+
# than one thread at a time (writing it directly from workers would
|
|
241
|
+
# not be safe -- append_to_jsonl opens the file fresh per call).
|
|
242
|
+
# -------------------------------------------------------------------
|
|
243
|
+
def _describe_one(task: Dict[str, Any]) -> Dict[str, Any]:
|
|
244
|
+
image_name = task["name"]
|
|
245
|
+
try:
|
|
246
|
+
with open(task["path"], "rb") as fh:
|
|
247
|
+
raw_bytes = fh.read()
|
|
248
|
+
# Re-encode to PNG in memory, mirroring figure_describer.py's own
|
|
249
|
+
# pix.tobytes("png") step. Not cosmetic: OpenAI's vision API does
|
|
250
|
+
# not accept BMP or TIFF at all, regardless of what the data: URL
|
|
251
|
+
# claims, so a raw passthrough would fail outright for those.
|
|
252
|
+
pil_image = _flatten_to_rgb(Image.open(BytesIO(raw_bytes)))
|
|
253
|
+
buf = BytesIO()
|
|
254
|
+
pil_image.save(buf, format="PNG")
|
|
255
|
+
png_bytes = buf.getvalue()
|
|
256
|
+
|
|
257
|
+
raw_response = call_vision_llm(
|
|
258
|
+
images=[png_bytes],
|
|
259
|
+
prompt=image_prompt,
|
|
260
|
+
provider=vision_provider,
|
|
261
|
+
api_key=vision_api_key,
|
|
262
|
+
model=vision_model,
|
|
263
|
+
detail=vision_detail,
|
|
264
|
+
reasoning_effort=reasoning_effort,
|
|
265
|
+
ollama_timeout_seconds=_OLLAMA_IMAGE_TIMEOUT_SECONDS,
|
|
266
|
+
)
|
|
267
|
+
|
|
268
|
+
parsed = _parse_single_image_response(raw_response, image_name)
|
|
269
|
+
if parsed is None:
|
|
270
|
+
return {"image_name": image_name, "row": None}
|
|
271
|
+
|
|
272
|
+
description = parsed.get("description")
|
|
273
|
+
if not isinstance(description, str):
|
|
274
|
+
logger.warning(
|
|
275
|
+
"'description' field for image %s is not a string (got %s) -- "
|
|
276
|
+
"model likely returned a malformed/nested response; skipping.",
|
|
277
|
+
image_name, type(description).__name__,
|
|
278
|
+
)
|
|
279
|
+
return {"image_name": image_name, "row": None}
|
|
280
|
+
|
|
281
|
+
row = {
|
|
282
|
+
"source": image_name,
|
|
283
|
+
"image_id": make_document_id(image_name),
|
|
284
|
+
"title": _derive_title(description),
|
|
285
|
+
"description": description,
|
|
286
|
+
}
|
|
287
|
+
return {"image_name": image_name, "row": row}
|
|
288
|
+
|
|
289
|
+
except Exception as exc:
|
|
290
|
+
logger.error("Vision LLM failed for image %s: %s", image_name, exc)
|
|
291
|
+
return {"image_name": image_name, "row": None}
|
|
292
|
+
|
|
293
|
+
written_basenames: List[str] = []
|
|
294
|
+
|
|
295
|
+
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
296
|
+
futures = {executor.submit(_describe_one, task): task for task in tasks}
|
|
297
|
+
for future in as_completed(futures):
|
|
298
|
+
task = futures[future]
|
|
299
|
+
try:
|
|
300
|
+
result = future.result()
|
|
301
|
+
except Exception as exc:
|
|
302
|
+
logger.error("Error collecting result for image %s: %s", task["name"], exc)
|
|
303
|
+
continue
|
|
304
|
+
|
|
305
|
+
row = result["row"]
|
|
306
|
+
if row is not None:
|
|
307
|
+
# Only mark this basename "written" if the row actually
|
|
308
|
+
# landed on disk -- the caller uses this to persist
|
|
309
|
+
# parsed_images.txt, and unlike the PDF/figure side (which
|
|
310
|
+
# resumes by reading 02_visuals_kb.jsonl's own content back,
|
|
311
|
+
# self-healing after a transient write failure), the image
|
|
312
|
+
# registry is a separate file with no such self-correction.
|
|
313
|
+
# A silently-swallowed disk/permission error here previously
|
|
314
|
+
# still counted as "done" -- permanent data loss with no
|
|
315
|
+
# retry path short of --redo/--rebuild.
|
|
316
|
+
if append_to_jsonl(descriptions_path, [row]):
|
|
317
|
+
written_basenames.append(result["image_name"])
|
|
318
|
+
else:
|
|
319
|
+
logger.error(
|
|
320
|
+
"Failed to write description for %s -- not marking processed, will retry next run.",
|
|
321
|
+
result["image_name"],
|
|
322
|
+
)
|
|
323
|
+
|
|
324
|
+
if written_basenames:
|
|
325
|
+
print(f"[IMAGES] Wrote {len(written_basenames)} image description(s) to {IMAGE_DESCRIPTIONS_FILE}.")
|
|
326
|
+
else:
|
|
327
|
+
logger.info("No new image descriptions written.")
|
|
328
|
+
|
|
329
|
+
return written_basenames
|