mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-25 01:20:11 +08:00
@@ -1,4 +1,4 @@
|
||||
package parser
|
||||
package pdf
|
||||
|
||||
import (
|
||||
"context"
|
||||
@@ -8,52 +8,36 @@ import (
|
||||
"log/slog"
|
||||
"sync"
|
||||
|
||||
inf "ragflow/internal/deepdoc/parser/pdf/inference"
|
||||
lyt "ragflow/internal/deepdoc/parser/pdf/layout"
|
||||
tbl "ragflow/internal/deepdoc/parser/pdf/table"
|
||||
pdf "ragflow/internal/deepdoc/parser/pdf/type"
|
||||
util "ragflow/internal/deepdoc/parser/pdf/util"
|
||||
)
|
||||
|
||||
// Parser is the main PDF text/layout extraction pipeline.
|
||||
// Parser is the core PDF text/layout extraction pipeline.
|
||||
// It corresponds to RAGFlowPdfParser in pdf_parser.py.
|
||||
// Parser is stateless after construction — safe to reuse across documents.
|
||||
// Stateless after construction — safe to reuse across documents.
|
||||
type Parser struct {
|
||||
Config pdf.ParserConfig
|
||||
|
||||
// DeepDoc is the required document layout / OCR / table recognition
|
||||
// service. Set at construction time by NewParser.
|
||||
DeepDoc pdf.DocAnalyzer
|
||||
|
||||
// SampleChars samples up to n chars from a page for English detection.
|
||||
// Defaults to random sampling (matching Python's random.choices).
|
||||
// Inject a deterministic sampler for reproducible tests.
|
||||
SampleChars pdf.SampleFunc
|
||||
|
||||
// tableBuilder is the TSR model adapter. Set at construction time
|
||||
//
|
||||
// different implementation via Config.TableBuilder.
|
||||
tableBuilder pdf.TableBuilder
|
||||
}
|
||||
|
||||
// NewParser creates a new Parser with the required DeepDoc service.
|
||||
func NewParser(cfg pdf.ParserConfig, doc pdf.DocAnalyzer) *Parser {
|
||||
tb := cfg.TableBuilder
|
||||
if tb == nil {
|
||||
tb = NewTableBuilderFor(doc)
|
||||
}
|
||||
return &Parser{
|
||||
Config: cfg,
|
||||
DeepDoc: doc,
|
||||
tableBuilder: tb,
|
||||
}
|
||||
// pageResult holds per-page output from extractPages.
|
||||
type pageResult struct {
|
||||
pg int
|
||||
ocrBoxes []pdf.TextBox
|
||||
chars []pdf.TextChar
|
||||
ocrUsed bool
|
||||
pageImg image.Image
|
||||
err error
|
||||
}
|
||||
|
||||
// New creates a new Parser with the given config.
|
||||
func NewParser(cfg pdf.ParserConfig) *Parser {
|
||||
return &Parser{Config: cfg}
|
||||
}
|
||||
|
||||
// ── TableBuilder factory ───────────────────────────────────────────────────
|
||||
|
||||
// tableBuilderFactory holds a model-specific TableBuilder factory registered
|
||||
// by EE packages via RegisterTableBuilder. If nil, the default OSS
|
||||
// implementation is used.
|
||||
var tableBuilderFactory func(pdf.DocAnalyzer) pdf.TableBuilder
|
||||
|
||||
// RegisterTableBuilder registers a TableBuilder factory for the PDF parser.
|
||||
@@ -62,30 +46,20 @@ func RegisterTableBuilder(factory func(pdf.DocAnalyzer) pdf.TableBuilder) {
|
||||
tableBuilderFactory = factory
|
||||
}
|
||||
|
||||
// NewTableBuilderFor creates the right TableBuilder, chosen by the registry.
|
||||
// Checks the registry first for EE-registered implementations, falling back
|
||||
// to the default OSS DeepDocTableBuilder. Label taxonomies are injected
|
||||
// before construction.
|
||||
func NewTableBuilderFor(doc pdf.DocAnalyzer) pdf.TableBuilder {
|
||||
if tableBuilderFactory != nil {
|
||||
return tableBuilderFactory(doc)
|
||||
}
|
||||
if c, ok := doc.(*inf.InferenceClient); ok {
|
||||
c.DLALabels = inf.DefaultDLALabels()
|
||||
c.TSRLabels = inf.DefaultTSRLabels()
|
||||
}
|
||||
return tbl.NewDeepDocTableBuilder(doc)
|
||||
}
|
||||
|
||||
// Parse runs the full PDF extraction pipeline: chars → boxes →
|
||||
// column assignment → text merge → vertical merge → sections.
|
||||
//
|
||||
// For documents larger than Config.BatchSize pages, processes in batches
|
||||
// to bound memory usage (matching Python's batch_size=50).
|
||||
//
|
||||
// Returns a pdf.ParseResult containing sections, tables, page images, figures,
|
||||
// and pipeline stage metrics. Parser itself remains stateless.
|
||||
func (p *Parser) Parse(ctx context.Context, engine pdf.PDFEngine) (*pdf.ParseResult, error) {
|
||||
// ── Public API ─────────────────────────────────────────────────────────────
|
||||
|
||||
// ParseRaw is the internal entry point: runs the core pipeline on an
|
||||
// already-opened engine. Exported for tests that inject mock engines.
|
||||
func (p *Parser) ParseRaw(ctx context.Context, engine pdf.PDFEngine, docAnalyzer pdf.DocAnalyzer) (*pdf.ParseResult, error) {
|
||||
tb := NewTableBuilderFor(docAnalyzer)
|
||||
|
||||
// Normalize page range
|
||||
pageCount, err := engine.PageCount()
|
||||
if err != nil {
|
||||
@@ -103,11 +77,10 @@ func (p *Parser) Parse(ctx context.Context, engine pdf.PDFEngine) (*pdf.ParseRes
|
||||
totalPages := toPage - fromPage + 1
|
||||
batchSize := p.Config.BatchSize
|
||||
if batchSize <= 0 {
|
||||
batchSize = 50 // default, matching Python's batch_size
|
||||
batchSize = 50
|
||||
}
|
||||
|
||||
// ── Prescan: lightweight char extraction for language/noise detection ──
|
||||
// No rendering, no OCR — just raw chars for global decisions.
|
||||
// ── Prescan ──
|
||||
prescanChars := make(map[int][]pdf.TextChar)
|
||||
prescanMedianH := make(map[int]float64)
|
||||
prescanMedianW := make(map[int]float64)
|
||||
@@ -115,26 +88,27 @@ func (p *Parser) Parse(ctx context.Context, engine pdf.PDFEngine) (*pdf.ParseRes
|
||||
chars, extractErr := engine.ExtractChars(pg)
|
||||
if extractErr != nil {
|
||||
slog.Warn("prescan: ExtractChars failed", "page", pg, "err", extractErr)
|
||||
chars = nil // skip broken pages (matching old behavior)
|
||||
chars = nil
|
||||
}
|
||||
prescanChars[pg] = chars
|
||||
prescanMedianH[pg] = util.MedianCharHeight(chars)
|
||||
prescanMedianW[pg] = util.MedianCharWidth(chars)
|
||||
}
|
||||
isEnglish := util.DetectEnglish(prescanChars, totalPages, p.SampleChars)
|
||||
isEnglish := util.DetectEnglish(prescanChars, totalPages, nil)
|
||||
scanNoise := util.IsScanNoise(util.FullTextFromChars(prescanChars))
|
||||
|
||||
// ── Extract PDF outlines/bookmarks (best-effort, non-fatal) ──
|
||||
// ── Outlines ──
|
||||
outlines, outlineErr := engine.Outlines()
|
||||
if outlineErr != nil {
|
||||
slog.Warn("Failed to extract PDF outlines; continuing without them", "err", outlineErr)
|
||||
outlines = nil
|
||||
}
|
||||
|
||||
// ── Small document: process all at once (no batching overhead) ──
|
||||
// ── Small document ──
|
||||
if totalPages <= batchSize {
|
||||
result, err := p.processPages(ctx, engine, fromPage, toPage,
|
||||
prescanChars, prescanMedianH, prescanMedianW, isEnglish, scanNoise)
|
||||
prescanChars, prescanMedianH, prescanMedianW, isEnglish, scanNoise,
|
||||
docAnalyzer, tb)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -142,7 +116,7 @@ func (p *Parser) Parse(ctx context.Context, engine pdf.PDFEngine) (*pdf.ParseRes
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// ── Large document: process in batches to bound memory ──
|
||||
// ── Large document: batched ──
|
||||
slog.Info("batched processing", "pages", totalPages, "batchSize", batchSize)
|
||||
result := &pdf.ParseResult{PageImages: make(map[int]image.Image)}
|
||||
for start := fromPage; start <= toPage; start += batchSize {
|
||||
@@ -151,7 +125,6 @@ func (p *Parser) Parse(ctx context.Context, engine pdf.PDFEngine) (*pdf.ParseRes
|
||||
}
|
||||
end := min(start+batchSize-1, toPage)
|
||||
|
||||
// Slice prescan data for this batch.
|
||||
batchChars := make(map[int][]pdf.TextChar, end-start+1)
|
||||
batchMH := make(map[int]float64, end-start+1)
|
||||
batchMW := make(map[int]float64, end-start+1)
|
||||
@@ -162,15 +135,14 @@ func (p *Parser) Parse(ctx context.Context, engine pdf.PDFEngine) (*pdf.ParseRes
|
||||
}
|
||||
|
||||
batch, err := p.processPages(ctx, engine, start, end,
|
||||
batchChars, batchMH, batchMW, isEnglish, scanNoise)
|
||||
batchChars, batchMH, batchMW, isEnglish, scanNoise,
|
||||
docAnalyzer, tb)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Merge batch results.
|
||||
result.Sections = append(result.Sections, batch.Sections...)
|
||||
result.Tables = append(result.Tables, batch.Tables...)
|
||||
// Figures() is computed on demand from Sections.
|
||||
for pg, img := range batch.PageImages {
|
||||
result.PageImages[pg] = img
|
||||
}
|
||||
@@ -184,33 +156,22 @@ func (p *Parser) Parse(ctx context.Context, engine pdf.PDFEngine) (*pdf.ParseRes
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// extractPages runs per-page OCR (detect + recognize) for the given page
|
||||
// range, returning text boxes, char data, whether any page used OCR, and
|
||||
// any errors encountered. Partial results are returned even when some
|
||||
// pages fail — callers should inspect the error for diagnostics but may
|
||||
// still use the returned boxes and chars.
|
||||
// ── Internal pipeline steps ────────────────────────────────────────────────
|
||||
|
||||
func (p *Parser) extractPages(ctx context.Context, engine pdf.PDFEngine,
|
||||
fromPage, toPage int,
|
||||
prescanChars map[int][]pdf.TextChar,
|
||||
medianHeights, medianWidths map[int]float64,
|
||||
pageImages map[int]image.Image,
|
||||
docAnalyzer pdf.DocAnalyzer,
|
||||
) ([]pdf.TextBox, map[int][]pdf.TextChar, bool, error) {
|
||||
var boxes []pdf.TextBox
|
||||
pageChars := make(map[int][]pdf.TextChar)
|
||||
ocrUsedAny := false
|
||||
|
||||
type pr struct {
|
||||
pg int
|
||||
ocrBoxes []pdf.TextBox
|
||||
chars []pdf.TextChar
|
||||
ocrUsed bool
|
||||
pageImg image.Image
|
||||
err error
|
||||
}
|
||||
pageCount := toPage - fromPage + 1
|
||||
results := make([]pr, pageCount)
|
||||
results := make([]pageResult, pageCount)
|
||||
|
||||
// Semaphore cap: 0 → sequential; >0 → bounded parallelism.
|
||||
cap := p.Config.MaxOCRConcurrency
|
||||
if cap <= 0 {
|
||||
cap = 1
|
||||
@@ -222,16 +183,15 @@ func (p *Parser) extractPages(ctx context.Context, engine pdf.PDFEngine,
|
||||
pg := fromPage + i
|
||||
chars := prescanChars[pg]
|
||||
|
||||
// Fast path: pages with embedded chars → sequential inline (no HTTP OCR).
|
||||
if len(chars) > 0 && !util.IsGarbledPage(chars) {
|
||||
pageImg, renderErr := renderPageToImage(engine, pg)
|
||||
pageImg, renderErr := RenderPageToImage(engine, pg)
|
||||
if renderErr == nil && pageImg != nil {
|
||||
pageImages[pg] = pageImg
|
||||
}
|
||||
var ocrBoxes []pdf.TextBox
|
||||
ocrUsed := false
|
||||
if !p.Config.SkipOCR && renderErr == nil && pageImg != nil {
|
||||
ocrBoxes = ocrMergeChars(ctx, pageImg, chars, p.DeepDoc, pg)
|
||||
ocrBoxes = ocrMergeChars(ctx, pageImg, chars, docAnalyzer, pg)
|
||||
if ocrBoxes == nil {
|
||||
ocrBoxes = lyt.CharsToBoxes(chars, pg, p.Config.SortByTop)
|
||||
} else {
|
||||
@@ -241,30 +201,28 @@ func (p *Parser) extractPages(ctx context.Context, engine pdf.PDFEngine,
|
||||
} else {
|
||||
ocrBoxes = lyt.CharsToBoxes(chars, pg, p.Config.SortByTop)
|
||||
}
|
||||
results[i] = pr{pg: pg, ocrBoxes: ocrBoxes, chars: chars, ocrUsed: ocrUsed}
|
||||
results[i] = pageResult{pg: pg, ocrBoxes: ocrBoxes, chars: chars, ocrUsed: ocrUsed}
|
||||
continue
|
||||
}
|
||||
|
||||
// OCR path: render + detect + recognize (potentially parallel).
|
||||
wg.Add(1)
|
||||
go func(i, pg int, chars []pdf.TextChar) {
|
||||
defer wg.Done()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
results[i] = pr{pg: pg, err: ctx.Err()}
|
||||
results[i] = pageResult{pg: pg, err: ctx.Err()}
|
||||
return
|
||||
case sem <- struct{}{}:
|
||||
}
|
||||
defer func() { <-sem }()
|
||||
|
||||
pageImg, err := renderPageToImage(engine, pg)
|
||||
pageImg, err := RenderPageToImage(engine, pg)
|
||||
if err != nil {
|
||||
results[i] = pr{pg: pg, err: err}
|
||||
results[i] = pageResult{pg: pg, err: err}
|
||||
return
|
||||
}
|
||||
// Check if context was cancelled during render.
|
||||
if err := ctx.Err(); err != nil {
|
||||
results[i] = pr{pg: pg, err: err}
|
||||
results[i] = pageResult{pg: pg, err: err}
|
||||
return
|
||||
}
|
||||
|
||||
@@ -275,7 +233,7 @@ func (p *Parser) extractPages(ctx context.Context, engine pdf.PDFEngine,
|
||||
if len(chars) > 0 {
|
||||
label = "garbled page"
|
||||
}
|
||||
ocrBoxes = ocrDetectAndRecognize(ctx, pageImg, p.DeepDoc, pg, label)
|
||||
ocrBoxes = ocrDetectAndRecognize(ctx, pageImg, docAnalyzer, pg, label)
|
||||
if ocrBoxes != nil {
|
||||
for j := range ocrBoxes {
|
||||
for _, r := range ocrBoxes[j].Text {
|
||||
@@ -286,9 +244,8 @@ func (p *Parser) extractPages(ctx context.Context, engine pdf.PDFEngine,
|
||||
ocrUsed = true
|
||||
}
|
||||
}
|
||||
// Merged OCR path for pages with both embedded and OCR chars.
|
||||
if !ocrUsed && len(chars) > 0 && !p.Config.SkipOCR {
|
||||
ocrBoxes = ocrMergeChars(ctx, pageImg, chars, p.DeepDoc, pg)
|
||||
ocrBoxes = ocrMergeChars(ctx, pageImg, chars, docAnalyzer, pg)
|
||||
if ocrBoxes != nil {
|
||||
ocrUsed = true
|
||||
}
|
||||
@@ -298,15 +255,252 @@ func (p *Parser) extractPages(ctx context.Context, engine pdf.PDFEngine,
|
||||
ocrBoxes = lyt.CharsToBoxes(chars, pg, p.Config.SortByTop)
|
||||
}
|
||||
}
|
||||
results[i] = pr{pg: pg, ocrBoxes: ocrBoxes, chars: chars, ocrUsed: ocrUsed, pageImg: pageImg}
|
||||
results[i] = pageResult{pg: pg, ocrBoxes: ocrBoxes, chars: chars, ocrUsed: ocrUsed, pageImg: pageImg}
|
||||
}(i, pg, chars)
|
||||
}
|
||||
wg.Wait()
|
||||
return mergePageResults(results, boxes, pageImages, pageChars, ocrUsedAny, medianHeights, medianWidths)
|
||||
}
|
||||
|
||||
// Merge results in page order.
|
||||
func (p *Parser) retryScanNoise(ctx context.Context, engine pdf.PDFEngine,
|
||||
fromPage, toPage int,
|
||||
pageImages map[int]image.Image,
|
||||
pageChars map[int][]pdf.TextChar,
|
||||
medianHeights, medianWidths map[int]float64,
|
||||
ocrUsedAny bool,
|
||||
docAnalyzer pdf.DocAnalyzer,
|
||||
) ([]pdf.TextBox, map[int][]pdf.TextChar, bool) {
|
||||
slog.Warn("scan noise: OCR retry", "from", fromPage, "to", toPage)
|
||||
var boxes []pdf.TextBox
|
||||
for pg := fromPage; pg <= toPage; pg++ {
|
||||
img := pageImages[pg]
|
||||
if img == nil {
|
||||
var err error
|
||||
img, err = RenderPageToImage(engine, pg)
|
||||
if err != nil {
|
||||
slog.Warn("scan noise: page render failed", "page", pg, "err", err)
|
||||
continue
|
||||
}
|
||||
pageImages[pg] = img
|
||||
}
|
||||
ocrBoxes := ocrDetectAndRecognize(ctx, img, docAnalyzer, pg, "scan page")
|
||||
if ocrBoxes == nil {
|
||||
slog.Warn("scan noise: page OCR empty", "page", pg)
|
||||
continue
|
||||
}
|
||||
boxes = append(boxes, ocrBoxes...)
|
||||
var chars []pdf.TextChar
|
||||
for _, b := range ocrBoxes {
|
||||
for _, r := range b.Text {
|
||||
chars = append(chars, pdf.TextChar{Text: string(r), Top: b.Top, Bottom: b.Bottom, PageNumber: pg})
|
||||
break
|
||||
}
|
||||
}
|
||||
pageChars[pg] = chars
|
||||
medianHeights[pg] = util.MedianCharHeight(chars)
|
||||
medianWidths[pg] = util.MedianCharWidth(chars)
|
||||
}
|
||||
slog.Debug("scan noise OCR retry complete", "pages", toPage-fromPage+1, "boxes", len(boxes))
|
||||
return boxes, pageChars, true
|
||||
}
|
||||
|
||||
func (p *Parser) retryZoom(ctx context.Context, engine pdf.PDFEngine,
|
||||
fromPage, toPage int,
|
||||
pageImages map[int]image.Image,
|
||||
boxes []pdf.TextBox, ocrUsedAny bool,
|
||||
docAnalyzer pdf.DocAnalyzer,
|
||||
) ([]pdf.TextBox, bool) {
|
||||
retryZoomVal := p.Config.Zoom * pdf.DlaScale
|
||||
retryDPI := retryZoomVal * 72
|
||||
slog.Info("zoom retry: re-rendering", "oldZoom", p.Config.Zoom, "newZoom", retryZoomVal)
|
||||
for pg := fromPage; pg <= toPage; pg++ {
|
||||
img, err := engine.RenderPageImage(pg, retryDPI)
|
||||
if err != nil {
|
||||
slog.Warn("zoom retry: render failed", "page", pg, "err", err)
|
||||
continue
|
||||
}
|
||||
pageImages[pg] = img
|
||||
if retryDPI != pdf.DlaDPI {
|
||||
if dlaImg, dlaErr := engine.RenderPageImage(pg, pdf.DlaDPI); dlaErr == nil {
|
||||
pageImages[pg] = dlaImg
|
||||
}
|
||||
}
|
||||
ocrBoxes := ocrDetectAndRecognize(ctx, img, docAnalyzer, pg, "zoom retry")
|
||||
if ocrBoxes == nil {
|
||||
continue
|
||||
}
|
||||
scaleFactor := retryZoomVal / p.Config.Zoom
|
||||
for i := range ocrBoxes {
|
||||
ocrBoxes[i].X0 /= scaleFactor
|
||||
ocrBoxes[i].X1 /= scaleFactor
|
||||
ocrBoxes[i].Top /= scaleFactor
|
||||
ocrBoxes[i].Bottom /= scaleFactor
|
||||
}
|
||||
boxes = append(boxes, ocrBoxes...)
|
||||
ocrUsedAny = true
|
||||
}
|
||||
return boxes, ocrUsedAny
|
||||
}
|
||||
|
||||
func (p *Parser) buildLayout(ctx context.Context,
|
||||
result *pdf.ParseResult, engine pdf.PDFEngine,
|
||||
boxes []pdf.TextBox, pageChars map[int][]pdf.TextChar,
|
||||
medianHeights, medianWidths map[int]float64,
|
||||
fromPage, toPage int, ocrUsedAny bool, isEnglish bool,
|
||||
docAnalyzer pdf.DocAnalyzer, tb pdf.TableBuilder,
|
||||
) error {
|
||||
result.Metrics.BoxesInitial = len(boxes)
|
||||
|
||||
result.Tables = p.enrichWithDeepDoc(ctx, result, engine, boxes, result.PageImages, docAnalyzer, tb)
|
||||
result.Metrics.TablesCount = len(result.Tables)
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
boxes = lyt.AssignColumn(boxes, p.Config.Zoom)
|
||||
boxes = lyt.TextMerge(boxes, medianHeights, p.Config.Zoom)
|
||||
result.Metrics.BoxesTextMerge = len(boxes)
|
||||
|
||||
lyt.SortByPageThenY(boxes, p.Config.SortByTop)
|
||||
|
||||
if ocrUsedAny {
|
||||
isEnglish = util.DetectEnglish(pageChars, toPage-fromPage+1, nil)
|
||||
}
|
||||
boxes = lyt.NaiveVerticalMerge(boxes, medianHeights, medianWidths, isEnglish)
|
||||
result.Metrics.BoxesVertMerge = len(boxes)
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
boxes = tbl.ExtractTableAndReplace(boxes, result.Tables)
|
||||
boxes = tbl.ConsolidateFigures(boxes)
|
||||
|
||||
pageHeights := make(map[int]float64, len(result.PageImages))
|
||||
for pg, img := range result.PageImages {
|
||||
pageHeights[pg] = float64(img.Bounds().Dy()) / p.Config.Zoom
|
||||
}
|
||||
result.Sections = lyt.BoxesToSections(boxes, pageHeights)
|
||||
result.Metrics.BoxesFinal = len(result.Sections)
|
||||
result.Sections = tbl.MergeCaptions(result.Sections, result.Figures())
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *Parser) processPages(ctx context.Context, engine pdf.PDFEngine,
|
||||
fromPage, toPage int,
|
||||
prescanChars map[int][]pdf.TextChar,
|
||||
medianHeights, medianWidths map[int]float64,
|
||||
isEnglish, isScanNoiseDoc bool,
|
||||
docAnalyzer pdf.DocAnalyzer, tb pdf.TableBuilder,
|
||||
) (*pdf.ParseResult, error) {
|
||||
result := &pdf.ParseResult{PageImages: make(map[int]image.Image)}
|
||||
|
||||
boxes, pageChars, ocrUsedAny, ocrErr := p.extractPages(ctx, engine,
|
||||
fromPage, toPage, prescanChars,
|
||||
medianHeights, medianWidths, result.PageImages, docAnalyzer)
|
||||
if ocrErr != nil {
|
||||
slog.Warn("extractPages: some pages failed OCR", "err", ocrErr)
|
||||
}
|
||||
|
||||
if isScanNoiseDoc {
|
||||
boxes, pageChars, ocrUsedAny = p.retryScanNoise(ctx, engine,
|
||||
fromPage, toPage, result.PageImages,
|
||||
pageChars, medianHeights, medianWidths, ocrUsedAny, docAnalyzer)
|
||||
}
|
||||
|
||||
if len(boxes) == 0 && p.Config.Zoom < 9 && !p.Config.SkipOCR {
|
||||
boxes, ocrUsedAny = p.retryZoom(ctx, engine, fromPage, toPage,
|
||||
result.PageImages, boxes, ocrUsedAny, docAnalyzer)
|
||||
}
|
||||
|
||||
if len(boxes) == 0 {
|
||||
return result, nil
|
||||
}
|
||||
|
||||
if err := p.buildLayout(ctx, result, engine, boxes, pageChars,
|
||||
medianHeights, medianWidths, fromPage, toPage, ocrUsedAny, isEnglish,
|
||||
docAnalyzer, tb); err != nil {
|
||||
return nil, fmt.Errorf("buildLayout: %w", err)
|
||||
}
|
||||
p.fillSectionImages(result)
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (p *Parser) fillSectionImages(result *pdf.ParseResult) {
|
||||
if len(result.PageImages) == 0 {
|
||||
return
|
||||
}
|
||||
tableImgByRegion := make(map[string]string, len(result.Tables))
|
||||
for _, tbl := range result.Tables {
|
||||
if tbl.ImageB64 == "" {
|
||||
continue
|
||||
}
|
||||
pg := 0
|
||||
if len(tbl.Positions) > 0 && len(tbl.Positions[0].PageNumbers) > 0 {
|
||||
pg = tbl.Positions[0].PageNumbers[0]
|
||||
}
|
||||
key := fmt.Sprintf("%d_%.1f_%.1f_%.1f_%.1f",
|
||||
pg, tbl.RegionLeft, tbl.RegionRight, tbl.RegionTop, tbl.RegionBottom)
|
||||
tableImgByRegion[key] = tbl.ImageB64
|
||||
}
|
||||
for i := range result.Sections {
|
||||
if result.Sections[i].LayoutType == pdf.LayoutTypeTable {
|
||||
if img, ok := matchTableImage(&result.Sections[i], tableImgByRegion); ok {
|
||||
result.Sections[i].Image = img
|
||||
continue
|
||||
}
|
||||
}
|
||||
if result.Sections[i].LayoutType == pdf.LayoutTypeFigure && len(result.Sections[i].Positions) > 0 {
|
||||
if dlaImg := util.CropSectionByDLA(result.Sections[i], result.DLADebug, result.PageImages); dlaImg != "" {
|
||||
result.Sections[i].Image = dlaImg
|
||||
continue
|
||||
}
|
||||
}
|
||||
img := util.CropSectionImage(result.Sections[i].PositionTag, result.PageImages, p.Config.Zoom)
|
||||
result.Sections[i].Image = img
|
||||
if img == "" && result.Sections[i].Text != "" {
|
||||
tag := result.Sections[i].PositionTag
|
||||
slog.Warn("cropSectionImage empty for non-empty section",
|
||||
"section", i, "posTag", tag[:min(80, len(tag))])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// matchTableImage looks up a pre-rendered table image for a section.
|
||||
// Uses Positions if available; falls back to TableItem Region boundaries.
|
||||
func matchTableImage(sec *pdf.Section, tableImgByRegion map[string]string) (string, bool) {
|
||||
pg := 0
|
||||
if len(sec.Positions) > 0 {
|
||||
pos := sec.Positions[0]
|
||||
if len(pos.PageNumbers) > 0 {
|
||||
pg = pos.PageNumbers[0]
|
||||
}
|
||||
key := fmt.Sprintf("%d_%.1f_%.1f_%.1f_%.1f", pg, pos.Left, pos.Right, pos.Top, pos.Bottom)
|
||||
if img, ok := tableImgByRegion[key]; ok {
|
||||
return img, true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
if sec.TableItem != nil {
|
||||
if len(sec.TableItem.Positions) > 0 && len(sec.TableItem.Positions[0].PageNumbers) > 0 {
|
||||
pg = sec.TableItem.Positions[0].PageNumbers[0]
|
||||
}
|
||||
key := fmt.Sprintf("%d_%.1f_%.1f_%.1f_%.1f", pg,
|
||||
sec.TableItem.RegionLeft, sec.TableItem.RegionRight,
|
||||
sec.TableItem.RegionTop, sec.TableItem.RegionBottom)
|
||||
if img, ok := tableImgByRegion[key]; ok {
|
||||
return img, true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// mergePageResults collects per-page OCR results into the final output.
|
||||
func mergePageResults(results []pageResult, boxes []pdf.TextBox, pageImages map[int]image.Image,
|
||||
pageChars map[int][]pdf.TextChar, ocrUsedAny bool,
|
||||
medianHeights, medianWidths map[int]float64,
|
||||
) ([]pdf.TextBox, map[int][]pdf.TextChar, bool, error) {
|
||||
var errs []error
|
||||
for i := 0; i < pageCount; i++ {
|
||||
r := results[i]
|
||||
for _, r := range results {
|
||||
if r.err != nil {
|
||||
slog.Warn("page OCR failed", "page", r.pg, "err", r.err)
|
||||
errs = append(errs, fmt.Errorf("page %d: %w", r.pg, r.err))
|
||||
@@ -329,233 +523,3 @@ func (p *Parser) extractPages(ctx context.Context, engine pdf.PDFEngine,
|
||||
}
|
||||
return boxes, pageChars, ocrUsedAny, errors.Join(errs...)
|
||||
}
|
||||
|
||||
// retryScanNoise re-runs OCR on all pages when prescan detects scan noise,
|
||||
// overwriting page-level state with fresh detect+recognize results.
|
||||
func (p *Parser) retryScanNoise(ctx context.Context, engine pdf.PDFEngine,
|
||||
fromPage, toPage int,
|
||||
pageImages map[int]image.Image,
|
||||
pageChars map[int][]pdf.TextChar,
|
||||
medianHeights, medianWidths map[int]float64,
|
||||
ocrUsedAny bool,
|
||||
) ([]pdf.TextBox, map[int][]pdf.TextChar, bool) {
|
||||
slog.Warn("scan noise: OCR retry", "from", fromPage, "to", toPage)
|
||||
var boxes []pdf.TextBox
|
||||
for pg := fromPage; pg <= toPage; pg++ {
|
||||
img := pageImages[pg]
|
||||
if img == nil {
|
||||
var err error
|
||||
img, err = renderPageToImage(engine, pg)
|
||||
if err != nil {
|
||||
slog.Warn("scan noise: page render failed", "page", pg, "err", err)
|
||||
continue
|
||||
}
|
||||
pageImages[pg] = img
|
||||
}
|
||||
ocrBoxes := ocrDetectAndRecognize(ctx, img, p.DeepDoc, pg, "scan page")
|
||||
if ocrBoxes == nil {
|
||||
slog.Warn("scan noise: page OCR empty", "page", pg)
|
||||
continue
|
||||
}
|
||||
boxes = append(boxes, ocrBoxes...)
|
||||
var chars []pdf.TextChar
|
||||
for _, b := range ocrBoxes {
|
||||
for _, r := range b.Text {
|
||||
chars = append(chars, pdf.TextChar{Text: string(r), Top: b.Top, Bottom: b.Bottom, PageNumber: pg})
|
||||
break
|
||||
}
|
||||
}
|
||||
pageChars[pg] = chars
|
||||
medianHeights[pg] = util.MedianCharHeight(chars)
|
||||
medianWidths[pg] = util.MedianCharWidth(chars)
|
||||
}
|
||||
slog.Debug("scan noise OCR retry complete", "pages", toPage-fromPage+1, "boxes", len(boxes))
|
||||
return boxes, pageChars, true
|
||||
}
|
||||
|
||||
// retryZoom re-renders pages at higher resolution and re-runs OCR when the
|
||||
// initial extraction produced zero boxes. Box coordinates are scaled back
|
||||
// to Config.Zoom space. Matches Python's __images__ retry.
|
||||
func (p *Parser) retryZoom(ctx context.Context, engine pdf.PDFEngine,
|
||||
fromPage, toPage int,
|
||||
pageImages map[int]image.Image,
|
||||
boxes []pdf.TextBox, ocrUsedAny bool,
|
||||
) ([]pdf.TextBox, bool) {
|
||||
retryZoom := p.Config.Zoom * pdf.DlaScale
|
||||
retryDPI := retryZoom * 72
|
||||
slog.Info("zoom retry: re-rendering", "oldZoom", p.Config.Zoom, "newZoom", retryZoom)
|
||||
for pg := fromPage; pg <= toPage; pg++ {
|
||||
img, err := engine.RenderPageImage(pg, retryDPI)
|
||||
if err != nil {
|
||||
slog.Warn("zoom retry: render failed", "page", pg, "err", err)
|
||||
continue
|
||||
}
|
||||
pageImages[pg] = img
|
||||
// Downstream DLA/TSR assumes pdf.DlaDPI. Re-render at standard
|
||||
// resolution so layout coordinates are scaled correctly.
|
||||
if retryDPI != pdf.DlaDPI {
|
||||
if dlaImg, dlaErr := engine.RenderPageImage(pg, pdf.DlaDPI); dlaErr == nil {
|
||||
pageImages[pg] = dlaImg
|
||||
}
|
||||
}
|
||||
ocrBoxes := ocrDetectAndRecognize(ctx, img, p.DeepDoc, pg, "zoom retry")
|
||||
if ocrBoxes == nil {
|
||||
continue
|
||||
}
|
||||
scaleFactor := retryZoom / p.Config.Zoom
|
||||
for i := range ocrBoxes {
|
||||
ocrBoxes[i].X0 /= scaleFactor
|
||||
ocrBoxes[i].X1 /= scaleFactor
|
||||
ocrBoxes[i].Top /= scaleFactor
|
||||
ocrBoxes[i].Bottom /= scaleFactor
|
||||
}
|
||||
boxes = append(boxes, ocrBoxes...)
|
||||
ocrUsedAny = true
|
||||
}
|
||||
return boxes, ocrUsedAny
|
||||
}
|
||||
|
||||
// buildLayout runs the DLA → TSR → Column → TextMerge → VM → pdf.Section
|
||||
// pipeline and populates result.Metrics, result.Tables, result.Sections,
|
||||
// and result.Sections. Matches Python's _parse_loaded_window_into_bboxes
|
||||
// order.
|
||||
func (p *Parser) buildLayout(ctx context.Context,
|
||||
result *pdf.ParseResult, engine pdf.PDFEngine,
|
||||
boxes []pdf.TextBox, pageChars map[int][]pdf.TextChar,
|
||||
medianHeights, medianWidths map[int]float64,
|
||||
fromPage, toPage int, ocrUsedAny bool, isEnglish bool,
|
||||
) error {
|
||||
result.Metrics.BoxesInitial = len(boxes)
|
||||
|
||||
result.Tables = p.enrichWithDeepDoc(ctx, result, engine, boxes, result.PageImages)
|
||||
result.Metrics.TablesCount = len(result.Tables)
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
boxes = lyt.AssignColumn(boxes, p.Config.Zoom)
|
||||
boxes = lyt.TextMerge(boxes, medianHeights, p.Config.Zoom)
|
||||
result.Metrics.BoxesTextMerge = len(boxes)
|
||||
|
||||
lyt.SortByPageThenY(boxes, p.Config.SortByTop)
|
||||
|
||||
if ocrUsedAny {
|
||||
isEnglish = util.DetectEnglish(pageChars, toPage-fromPage+1, p.SampleChars)
|
||||
}
|
||||
boxes = lyt.NaiveVerticalMerge(boxes, medianHeights, medianWidths, isEnglish)
|
||||
result.Metrics.BoxesVertMerge = len(boxes)
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
boxes = tbl.ExtractTableAndReplace(boxes, result.Tables)
|
||||
boxes = tbl.ConsolidateFigures(boxes)
|
||||
|
||||
pageHeights := make(map[int]float64, len(result.PageImages))
|
||||
for pg, img := range result.PageImages {
|
||||
pageHeights[pg] = float64(img.Bounds().Dy()) / p.Config.Zoom
|
||||
}
|
||||
result.Sections = lyt.BoxesToSections(boxes, pageHeights)
|
||||
result.Metrics.BoxesFinal = len(result.Sections)
|
||||
result.Sections = tbl.MergeCaptions(result.Sections, result.Figures())
|
||||
return nil
|
||||
}
|
||||
|
||||
// processPages runs the full pipeline on pages [fromPage, toPage].
|
||||
// prescanChars provides pre-extracted chars (avoids double extraction).
|
||||
func (p *Parser) processPages(ctx context.Context, engine pdf.PDFEngine,
|
||||
fromPage, toPage int,
|
||||
prescanChars map[int][]pdf.TextChar,
|
||||
medianHeights, medianWidths map[int]float64,
|
||||
isEnglish, isScanNoiseDoc bool,
|
||||
) (*pdf.ParseResult, error) {
|
||||
result := &pdf.ParseResult{PageImages: make(map[int]image.Image)}
|
||||
|
||||
// 1. OCR extraction — per-page detect + recognize + char merge.
|
||||
boxes, pageChars, ocrUsedAny, ocrErr := p.extractPages(ctx, engine,
|
||||
fromPage, toPage, prescanChars,
|
||||
medianHeights, medianWidths, result.PageImages)
|
||||
if ocrErr != nil {
|
||||
slog.Warn("extractPages: some pages failed OCR", "err", ocrErr)
|
||||
}
|
||||
// 2. Scan noise retry — re-OCR all pages when prescan detects scan noise.
|
||||
if isScanNoiseDoc {
|
||||
boxes, pageChars, ocrUsedAny = p.retryScanNoise(ctx, engine,
|
||||
fromPage, toPage, result.PageImages,
|
||||
pageChars, medianHeights, medianWidths, ocrUsedAny)
|
||||
}
|
||||
|
||||
// 3. Zoom retry — re-render at higher resolution if OCR produced zero boxes.
|
||||
if len(boxes) == 0 && p.Config.Zoom < 9 && !p.Config.SkipOCR {
|
||||
boxes, ocrUsedAny = p.retryZoom(ctx, engine, fromPage, toPage,
|
||||
result.PageImages, boxes, ocrUsedAny)
|
||||
}
|
||||
|
||||
if len(boxes) == 0 {
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// 4. Layout pipeline — DLA → TSR → Column → TextMerge → VM → Sections.
|
||||
if err := p.buildLayout(ctx, result, engine, boxes, pageChars,
|
||||
medianHeights, medianWidths, fromPage, toPage, ocrUsedAny, isEnglish); err != nil {
|
||||
return nil, fmt.Errorf("buildLayout: %w", err)
|
||||
}
|
||||
// 5. Crop section images from page renders.
|
||||
p.fillSectionImages(result)
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// fillSectionImages populates result.Sections[i].Image with cropped
|
||||
// page images. Table sections are matched to their TableItem image;
|
||||
// figure sections try DLA-aware cropping first, then fall back to
|
||||
// position-tag-based cropping.
|
||||
func (p *Parser) fillSectionImages(result *pdf.ParseResult) {
|
||||
if len(result.PageImages) == 0 {
|
||||
return
|
||||
}
|
||||
// Build lookup: DLA region -> table image (base64).
|
||||
tableImgByRegion := make(map[string]string, len(result.Tables))
|
||||
for _, tbl := range result.Tables {
|
||||
if tbl.ImageB64 == "" {
|
||||
continue
|
||||
}
|
||||
pg := 0
|
||||
if len(tbl.Positions) > 0 && len(tbl.Positions[0].PageNumbers) > 0 {
|
||||
pg = tbl.Positions[0].PageNumbers[0]
|
||||
}
|
||||
key := fmt.Sprintf("%d_%.1f_%.1f_%.1f_%.1f",
|
||||
pg, tbl.RegionLeft, tbl.RegionRight, tbl.RegionTop, tbl.RegionBottom)
|
||||
tableImgByRegion[key] = tbl.ImageB64
|
||||
}
|
||||
for i := range result.Sections {
|
||||
if result.Sections[i].LayoutType == pdf.LayoutTypeTable && len(result.Sections[i].Positions) > 0 {
|
||||
pos := result.Sections[i].Positions[0]
|
||||
pg := 0
|
||||
if len(pos.PageNumbers) > 0 {
|
||||
pg = pos.PageNumbers[0]
|
||||
}
|
||||
key := fmt.Sprintf("%d_%.1f_%.1f_%.1f_%.1f",
|
||||
pg, pos.Left, pos.Right, pos.Top, pos.Bottom)
|
||||
if img, ok := tableImgByRegion[key]; ok {
|
||||
result.Sections[i].Image = img
|
||||
continue
|
||||
}
|
||||
}
|
||||
// Try DLA-aware cropping for figure sections (matching Python's
|
||||
// cropout which uses DLA region boundaries instead of text boxes).
|
||||
if result.Sections[i].LayoutType == pdf.LayoutTypeFigure && len(result.Sections[i].Positions) > 0 {
|
||||
if dlaImg := util.CropSectionByDLA(result.Sections[i], result.DLADebug, result.PageImages); dlaImg != "" {
|
||||
result.Sections[i].Image = dlaImg
|
||||
continue
|
||||
}
|
||||
}
|
||||
img := util.CropSectionImage(result.Sections[i].PositionTag, result.PageImages, p.Config.Zoom)
|
||||
result.Sections[i].Image = img
|
||||
if img == "" && result.Sections[i].Text != "" {
|
||||
tag := result.Sections[i].PositionTag
|
||||
slog.Warn("cropSectionImage empty for non-empty section",
|
||||
"section", i, "posTag", tag[:min(80, len(tag))])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user