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
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Iterator
|
||||
from typing import Iterator, Sequence
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -95,20 +95,17 @@ def _is_blank_pixmap(pixmap) -> bool:
|
||||
total = pixmap.width * pixmap.height
|
||||
if total == 0:
|
||||
return True
|
||||
step = max(1, (total * stride) // 4000)
|
||||
values = bytes(samples[0 : len(samples) : step])
|
||||
if not values:
|
||||
return True
|
||||
return len(set(values)) <= 1
|
||||
except Exception: # pragma: no cover - formatos de pixmap inesperados
|
||||
sample = samples[0]
|
||||
return all(s == sample for s in samples)
|
||||
except Exception: # pragma: no cover
|
||||
return False
|
||||
|
||||
|
||||
def file_sha256(path: Path, chunk_size: int = 1024 * 1024) -> str:
|
||||
"""SHA-256 do conteúdo do arquivo."""
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
while chunk := handle.read(chunk_size):
|
||||
with path.open("rb") as f:
|
||||
while chunk := f.read(chunk_size):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
@@ -151,6 +148,70 @@ class PdfExtractor:
|
||||
finally:
|
||||
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:
|
||||
"""Extrai o texto de uma página, usando cache quando disponível."""
|
||||
cache_file = self._cache_path(file_hash, page_number)
|
||||
|
||||
Reference in new issue
Block a user