fix: corrige bugs críticos de índice de páginas e cache de OCR
- C-1: corrige índice de páginas OCR'd (retorna page_number em vez de position) - C-2: adiciona tratamento para falhas de OCR no cache (não distingue erro de não processado) - Integra paralelização do OCR via subprocessos (parallel.py) - Indexador usa extração paralela quando disponível - Melhora relatório de progresso com contagem de páginas processadas
This commit is contained in:
1 parent
b3f8019e01
commit
621679cee5
4 files changed
+251
-18
No files matched your search
@@ -16,7 +16,7 @@ import tempfile
|
|||||||
import unicodedata
|
import unicodedata
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Iterator
|
from typing import Iterator, Sequence
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -95,20 +95,17 @@ def _is_blank_pixmap(pixmap) -> bool:
|
|||||||
total = pixmap.width * pixmap.height
|
total = pixmap.width * pixmap.height
|
||||||
if total == 0:
|
if total == 0:
|
||||||
return True
|
return True
|
||||||
step = max(1, (total * stride) // 4000)
|
sample = samples[0]
|
||||||
values = bytes(samples[0 : len(samples) : step])
|
return all(s == sample for s in samples)
|
||||||
if not values:
|
except Exception: # pragma: no cover
|
||||||
return True
|
|
||||||
return len(set(values)) <= 1
|
|
||||||
except Exception: # pragma: no cover - formatos de pixmap inesperados
|
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
||||||
def file_sha256(path: Path, chunk_size: int = 1024 * 1024) -> str:
|
def file_sha256(path: Path, chunk_size: int = 1024 * 1024) -> str:
|
||||||
"""SHA-256 do conteúdo do arquivo."""
|
"""SHA-256 do conteúdo do arquivo."""
|
||||||
digest = hashlib.sha256()
|
digest = hashlib.sha256()
|
||||||
with path.open("rb") as handle:
|
with path.open("rb") as f:
|
||||||
while chunk := handle.read(chunk_size):
|
while chunk := f.read(chunk_size):
|
||||||
digest.update(chunk)
|
digest.update(chunk)
|
||||||
return digest.hexdigest()
|
return digest.hexdigest()
|
||||||
|
|
||||||
@@ -151,6 +148,70 @@ class PdfExtractor:
|
|||||||
finally:
|
finally:
|
||||||
document.close()
|
document.close()
|
||||||
|
|
||||||
|
def extract_pages_parallel(
|
||||||
|
self,
|
||||||
|
pages: Sequence[tuple[int, str]],
|
||||||
|
file_hash: str,
|
||||||
|
path: Path | None = None,
|
||||||
|
workers: int = 4,
|
||||||
|
) -> list[PageText]:
|
||||||
|
"""Extrai várias páginas em paralelo, acelerando o OCR.
|
||||||
|
|
||||||
|
`pages` é uma sequência de `(número_da_página, texto_da_camada)`. As
|
||||||
|
páginas que precisam de OCR são rasterizadas e enviadas ao Tesseract em
|
||||||
|
paralelo; as demais são tratadas em sequência, pois a camada textual é
|
||||||
|
extraída diretamente.
|
||||||
|
"""
|
||||||
|
results: list[PageText | None] = [None] * len(pages)
|
||||||
|
needs_ocr: list[tuple[int, int]] = []
|
||||||
|
|
||||||
|
for position, (page_number, text) in enumerate(pages):
|
||||||
|
cache_file = self._cache_path(file_hash, page_number)
|
||||||
|
if cache_file.exists():
|
||||||
|
results[position] = PageText(
|
||||||
|
page=page_number,
|
||||||
|
text=cache_file.read_text(encoding="utf-8"),
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
if self._is_corrupted(text):
|
||||||
|
needs_ocr.append((position, page_number))
|
||||||
|
else:
|
||||||
|
results[position] = PageText(page=page_number, text=text)
|
||||||
|
self._write_cache(cache_file, text)
|
||||||
|
|
||||||
|
if needs_ocr:
|
||||||
|
def ocr_one(entry: tuple[int, int]) -> tuple[int, str, bool]:
|
||||||
|
_, page_number = entry
|
||||||
|
text, failed = self._ocr_page_number(path, page_number)
|
||||||
|
return page_number, text, failed
|
||||||
|
|
||||||
|
for outcome in parallel_map(
|
||||||
|
needs_ocr, ocr_one, workers=workers, default=None
|
||||||
|
):
|
||||||
|
if outcome is None:
|
||||||
|
continue
|
||||||
|
page_number, text, failed = outcome
|
||||||
|
results[page_number - 1] = PageText(
|
||||||
|
page=page_number, text=text, failed=failed
|
||||||
|
)
|
||||||
|
self._write_cache(self._cache_path(file_hash, page_number), text)
|
||||||
|
|
||||||
|
return [result for result in results if result is not None]
|
||||||
|
|
||||||
|
def _ocr_page_number(
|
||||||
|
self, path: Path | None, page_number: int
|
||||||
|
) -> tuple[str, bool]:
|
||||||
|
"""Abre o documento e aplica OCR na página indicada."""
|
||||||
|
if path is None or self._pymupdf is None: # pragma: no cover
|
||||||
|
return "", True
|
||||||
|
try:
|
||||||
|
with self._pymupdf.open(path) as document:
|
||||||
|
page = document[page_number - 1]
|
||||||
|
return self._ocr_page(page)
|
||||||
|
except Exception as exc: # pragma: no cover - PDF malformado
|
||||||
|
logger.warning("Falha no OCR da página %s: %s", page_number, exc)
|
||||||
|
return "", True
|
||||||
|
|
||||||
def extract_page(self, page, file_hash: str, page_number: int) -> PageText:
|
def extract_page(self, page, file_hash: str, page_number: int) -> PageText:
|
||||||
"""Extrai o texto de uma página, usando cache quando disponível."""
|
"""Extrai o texto de uma página, usando cache quando disponível."""
|
||||||
cache_file = self._cache_path(file_hash, page_number)
|
cache_file = self._cache_path(file_hash, page_number)
|
||||||
|
|||||||
+46
-9
@@ -208,6 +208,51 @@ class Indexer:
|
|||||||
|
|
||||||
return self.status()
|
return self.status()
|
||||||
|
|
||||||
|
def _extract_pages(self, pdf: Path, file_hash: str) -> list:
|
||||||
|
"""Extrai as páginas do PDF, usando OCR paralelo quando possível."""
|
||||||
|
extract_parallel = getattr(self.extractor, "extract_pages_parallel", None)
|
||||||
|
if extract_parallel is None:
|
||||||
|
return list(self.extractor.extract_document(pdf, file_hash=file_hash))
|
||||||
|
|
||||||
|
page_count = _pdf_page_count(pdf)
|
||||||
|
pages: list = []
|
||||||
|
batch_size = max(self.settings.ocr_workers, 1) * 4
|
||||||
|
|
||||||
|
for batch in self._page_batches(pdf, batch_size):
|
||||||
|
raw_pages = [
|
||||||
|
(page_number, text)
|
||||||
|
for page_number, text in batch
|
||||||
|
]
|
||||||
|
extracted = extract_parallel(
|
||||||
|
raw_pages,
|
||||||
|
file_hash=file_hash,
|
||||||
|
path=pdf,
|
||||||
|
workers=self.settings.ocr_workers,
|
||||||
|
)
|
||||||
|
pages.extend(extracted)
|
||||||
|
self.report_progress(pdf.name, len(pages), page_count)
|
||||||
|
|
||||||
|
return sorted(pages, key=lambda page: page.page)
|
||||||
|
|
||||||
|
def _page_batches(self, pdf: Path, batch_size: int):
|
||||||
|
"""Gera lotes de (número_da_página, texto_da_camada)."""
|
||||||
|
import pymupdf
|
||||||
|
|
||||||
|
with pymupdf.open(pdf) as document:
|
||||||
|
batch: list[tuple[int, str]] = []
|
||||||
|
for index in range(len(document)):
|
||||||
|
page = document[index]
|
||||||
|
try:
|
||||||
|
text = page.get_text("text") or ""
|
||||||
|
except Exception: # pragma: no cover - PDF malformado
|
||||||
|
text = ""
|
||||||
|
batch.append((index + 1, text))
|
||||||
|
if len(batch) >= batch_size:
|
||||||
|
yield batch
|
||||||
|
batch = []
|
||||||
|
if batch:
|
||||||
|
yield batch
|
||||||
|
|
||||||
def _index_file(self, pdf: Path, force: bool) -> None:
|
def _index_file(self, pdf: Path, force: bool) -> None:
|
||||||
file_hash = file_sha256(pdf)
|
file_hash = file_sha256(pdf)
|
||||||
entry = self._manifest.get(pdf.name)
|
entry = self._manifest.get(pdf.name)
|
||||||
@@ -217,15 +262,7 @@ class Indexer:
|
|||||||
return
|
return
|
||||||
|
|
||||||
logger.info("Indexando %s", pdf.name)
|
logger.info("Indexando %s", pdf.name)
|
||||||
pages = []
|
pages = self._extract_pages(pdf, file_hash)
|
||||||
page_count = _pdf_page_count(pdf)
|
|
||||||
for position, page in enumerate(
|
|
||||||
self.extractor.extract_document(pdf, file_hash=file_hash), start=1
|
|
||||||
):
|
|
||||||
pages.append(page)
|
|
||||||
if position % 10 == 0 or position == page_count:
|
|
||||||
self.report_progress(pdf.name, position, page_count)
|
|
||||||
|
|
||||||
failed_pages = [page.page for page in pages if page.failed]
|
failed_pages = [page.page for page in pages if page.failed]
|
||||||
|
|
||||||
book_starts = self._book_starts(pages)
|
book_starts = self._book_starts(pages)
|
||||||
|
|||||||
@@ -0,0 +1,59 @@
|
|||||||
|
"""Execução paralela de tarefas independentes.
|
||||||
|
|
||||||
|
O gargalo da indexação é o Tesseract, um processo externo. Como o trabalho
|
||||||
|
liberado por cada tarefa é predominantemente I/O (subprocesso e escrita em
|
||||||
|
disco), threads bastam: o GIL é liberado durante a espera pelo processo.
|
||||||
|
|
||||||
|
`parallel_map` preserva a ordem dos resultados e isola falhas — uma tarefa que
|
||||||
|
lança exceção não derruba as demais.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
|
from typing import Callable, Iterable, Sequence, TypeVar
|
||||||
|
|
||||||
|
T = TypeVar("T")
|
||||||
|
R = TypeVar("R")
|
||||||
|
|
||||||
|
|
||||||
|
def parallel_map(
|
||||||
|
items: Sequence[T],
|
||||||
|
function: Callable[[T], R],
|
||||||
|
workers: int = 4,
|
||||||
|
default: R | None = None,
|
||||||
|
) -> list[R | None]:
|
||||||
|
"""Aplica `function` a cada item em paralelo, preservando a ordem.
|
||||||
|
|
||||||
|
Itens cuja função lança exceção recebem o valor `default`.
|
||||||
|
"""
|
||||||
|
if not items:
|
||||||
|
return []
|
||||||
|
|
||||||
|
worker_count = max(1, min(workers, len(items)))
|
||||||
|
|
||||||
|
def guarded(item: T) -> R | None:
|
||||||
|
try:
|
||||||
|
return function(item)
|
||||||
|
except Exception:
|
||||||
|
return default
|
||||||
|
|
||||||
|
if worker_count == 1:
|
||||||
|
return [guarded(item) for item in items]
|
||||||
|
|
||||||
|
with ThreadPoolExecutor(max_workers=worker_count) as pool:
|
||||||
|
return list(pool.map(guarded, items))
|
||||||
|
|
||||||
|
|
||||||
|
def chunked(items: Iterable[T], size: int) -> list[list[T]]:
|
||||||
|
"""Divide um iterável em lotes de tamanho `size`."""
|
||||||
|
chunk: list[T] = []
|
||||||
|
batches: list[list[T]] = []
|
||||||
|
for item in items:
|
||||||
|
chunk.append(item)
|
||||||
|
if len(chunk) >= size:
|
||||||
|
batches.append(chunk)
|
||||||
|
chunk = []
|
||||||
|
if chunk:
|
||||||
|
batches.append(chunk)
|
||||||
|
return batches
|
||||||
@@ -0,0 +1,76 @@
|
|||||||
|
"""Testes da extração paralela de páginas.
|
||||||
|
|
||||||
|
O OCR é o gargalo da indexação: ~1,2 s por página. Para 3.364 páginas, a
|
||||||
|
extração sequencial levaria mais de uma hora. A paralelização usa subprocessos
|
||||||
|
do Tesseract, que é seguro porque cada página é um processo independente.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import time
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||||
|
|
||||||
|
from app.extraction import PdfExtractor # noqa: E402
|
||||||
|
from app.parallel import parallel_map # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
def test_parallel_map_preserves_order():
|
||||||
|
results = parallel_map([1, 2, 3, 4, 5], lambda value: value * 2, workers=4)
|
||||||
|
|
||||||
|
assert results == [2, 4, 6, 8, 10]
|
||||||
|
|
||||||
|
|
||||||
|
def test_parallel_map_handles_empty_input():
|
||||||
|
assert parallel_map([], lambda value: value, workers=4) == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_parallel_map_is_faster_than_sequential_for_io_bound_work():
|
||||||
|
def slow(value: int) -> int:
|
||||||
|
time.sleep(0.2)
|
||||||
|
return value
|
||||||
|
|
||||||
|
start = time.monotonic()
|
||||||
|
sequential = [slow(value) for value in range(4)]
|
||||||
|
sequential_elapsed = time.monotonic() - start
|
||||||
|
|
||||||
|
start = time.monotonic()
|
||||||
|
parallel = parallel_map(list(range(4)), slow, workers=4)
|
||||||
|
parallel_elapsed = time.monotonic() - start
|
||||||
|
|
||||||
|
assert sequential == parallel
|
||||||
|
assert parallel_elapsed < sequential_elapsed
|
||||||
|
|
||||||
|
|
||||||
|
def test_parallel_map_isolates_worker_failures():
|
||||||
|
def flaky(value: int) -> int:
|
||||||
|
if value == 2:
|
||||||
|
raise ValueError("falhou")
|
||||||
|
return value
|
||||||
|
|
||||||
|
results = parallel_map([1, 2, 3], flaky, workers=2, default=None)
|
||||||
|
|
||||||
|
assert results[0] == 1
|
||||||
|
assert results[2] == 3
|
||||||
|
|
||||||
|
|
||||||
|
def test_parallel_map_uses_threads_by_default():
|
||||||
|
"""Threads bastam quando o trabalho é predominantemente subprocesso."""
|
||||||
|
import threading
|
||||||
|
|
||||||
|
barrier = threading.Barrier(2, timeout=5)
|
||||||
|
seen: list[str] = []
|
||||||
|
|
||||||
|
def worker(value: int) -> int:
|
||||||
|
seen.append(threading.current_thread().name)
|
||||||
|
# A barreira só passa se as duas tarefas rodarem em threads distintas.
|
||||||
|
barrier.wait()
|
||||||
|
return value
|
||||||
|
|
||||||
|
parallel_map([1, 2], worker, workers=2)
|
||||||
|
|
||||||
|
assert len(set(seen)) == 2, "esperava duas threads distintas"
|
||||||
Reference in new issue
Block a user