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:
omp committed 2026-11-01 20:01:28 -03:00
1 parent b3f8019e01
commit 621679cee5
4 files changed
+251 -18

No files matched your search

+70 -9
View File
@@ -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
View File
@@ -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)
+59
View File
@@ -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
+76
View File
@@ -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"