mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-04 06:40:29 +08:00
Refactor: migrate pdf_parser.py to golang (#16323)
### What problem does this PR solve? Http API based on onnx model. pdf_parser.py to golang ### Type of change - [x] Refactoring
This commit is contained in:
Jack
165
internal/deepdoc/parser/pdf/pdfium/pdfium.go
Normal file
165
internal/deepdoc/parser/pdf/pdfium/pdfium.go
Normal file
@@ -0,0 +1,165 @@
|
||||
// Package pdfium renders PDF pages using the system's libpdfium.so
|
||||
// (bundled with pypdfium2). It exists solely to replace pdf_oxide's
|
||||
// RenderPageRaw for use cases where image quality matters for downstream
|
||||
// OCR/DLA — pdf_oxide still handles all text/char/table extraction.
|
||||
package pdfium
|
||||
|
||||
/*
|
||||
#cgo LDFLAGS: -L/home/shenyushi/cc-workspace/ragflow/.venv/lib/python3.13/site-packages/pypdfium2_raw -lpdfium -lm -lpthread -ldl
|
||||
#cgo linux LDFLAGS: -Wl,-rpath,/home/shenyushi/cc-workspace/ragflow/.venv/lib/python3.13/site-packages/pypdfium2_raw
|
||||
|
||||
#include <stdint.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
typedef struct FPDF_DOCUMENT__ { int unused; } *FPDF_DOCUMENT;
|
||||
typedef struct FPDF_PAGE__ { int unused; } *FPDF_PAGE;
|
||||
typedef struct FPDF_BITMAP__ { int unused; } *FPDF_BITMAP;
|
||||
|
||||
extern void FPDF_InitLibrary(void);
|
||||
extern FPDF_DOCUMENT FPDF_LoadMemDocument(const void* data_buf, int size, const char* password);
|
||||
extern void FPDF_CloseDocument(FPDF_DOCUMENT document);
|
||||
extern int FPDF_GetPageCount(FPDF_DOCUMENT document);
|
||||
extern FPDF_PAGE FPDF_LoadPage(FPDF_DOCUMENT document, int page_index);
|
||||
extern void FPDF_ClosePage(FPDF_PAGE page);
|
||||
extern double FPDF_GetPageWidth(FPDF_PAGE page);
|
||||
extern double FPDF_GetPageHeight(FPDF_PAGE page);
|
||||
extern FPDF_BITMAP FPDFBitmap_Create(int width, int height, int alpha);
|
||||
extern void FPDFBitmap_Destroy(FPDF_BITMAP bitmap);
|
||||
extern void FPDF_RenderPageBitmap(FPDF_BITMAP bitmap, FPDF_PAGE page,
|
||||
int start_x, int start_y, int size_x, int size_y,
|
||||
int rotate, int flags);
|
||||
extern void* FPDFBitmap_GetBuffer(FPDF_BITMAP bitmap);
|
||||
extern int FPDFBitmap_GetWidth(FPDF_BITMAP bitmap);
|
||||
extern int FPDFBitmap_GetHeight(FPDF_BITMAP bitmap);
|
||||
extern int FPDFBitmap_GetStride(FPDF_BITMAP bitmap);
|
||||
*/
|
||||
import "C"
|
||||
import (
|
||||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
"math"
|
||||
"sync"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
var initOnce sync.Once
|
||||
|
||||
// pdfiumMu serializes all pdfium C API access. pdfium is NOT thread-safe —
|
||||
// concurrent calls to FPDF_LoadPage / FPDF_RenderPageBitmap corrupt the
|
||||
// global heap, causing SIGSEGV. See TestPdfiumConcurrentSafety.
|
||||
var pdfiumMu sync.Mutex
|
||||
|
||||
// Init initializes the PDFium library. Safe to call multiple times.
|
||||
func Init() { initOnce.Do(func() { C.FPDF_InitLibrary() }) }
|
||||
|
||||
// PageSize returns the page dimensions in PDF points (1/72 inch) as seen
|
||||
// after rotation. For a page with /Rotate 90 on A4, this returns ~842×595
|
||||
// (swapped from the MediaBox 595×842). The call is cheap — it opens the
|
||||
// document and page, reads dimensions, then closes.
|
||||
func PageSize(pdfData []byte, pageIdx int) (width, height float64, err error) {
|
||||
Init()
|
||||
pdfiumMu.Lock()
|
||||
defer pdfiumMu.Unlock()
|
||||
_, _, pw, ph, closeAll, err := openPage(pdfData, pageIdx)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
closeAll()
|
||||
return pw, ph, nil
|
||||
}
|
||||
|
||||
// RenderPage renders a single page of a PDF to an *image.RGBA at the given DPI.
|
||||
// pdfData is the raw PDF bytes, pageIdx is 0-based.
|
||||
func RenderPage(pdfData []byte, pageIdx int, dpi float64) (*image.RGBA, error) {
|
||||
Init()
|
||||
pdfiumMu.Lock()
|
||||
defer pdfiumMu.Unlock()
|
||||
_, page, pw, ph, closeAll, err := openPage(pdfData, pageIdx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer closeAll()
|
||||
|
||||
scale := dpi / 72.0
|
||||
pxW := int(math.Round(pw * scale))
|
||||
pxH := int(math.Round(ph * scale))
|
||||
|
||||
bitmap := C.FPDFBitmap_Create(C.int(pxW), C.int(pxH), 1) // 1 = RGBA
|
||||
if bitmap == nil {
|
||||
return nil, fmt.Errorf("pdfium: FPDFBitmap_Create(%d,%d) returned nil", pxW, pxH)
|
||||
}
|
||||
defer C.FPDFBitmap_Destroy(bitmap)
|
||||
|
||||
// Fill with opaque white before rendering, so transparent areas
|
||||
// (e.g. outside crop box) are white rather than undefined.
|
||||
stride := int(C.FPDFBitmap_GetStride(bitmap))
|
||||
buf := C.FPDFBitmap_GetBuffer(bitmap)
|
||||
pixels := (*[1 << 30]byte)(unsafe.Pointer(buf))[: pxH*stride : pxH*stride]
|
||||
for i := range pixels {
|
||||
pixels[i] = 255
|
||||
}
|
||||
|
||||
// FPDF_ANNOT (0x01) — render annotations.
|
||||
// LCD text AA (0x02) is left off; default text smoothing is sufficient.
|
||||
C.FPDF_RenderPageBitmap(bitmap, page, 0, 0, C.int(pxW), C.int(pxH), 0, 0x01)
|
||||
|
||||
// pdfium outputs BGRA; convert to RGBA.
|
||||
img := image.NewRGBA(image.Rect(0, 0, pxW, pxH))
|
||||
for y := 0; y < pxH; y++ {
|
||||
for x := 0; x < pxW; x++ {
|
||||
off := y*stride + x*4
|
||||
img.SetRGBA(x, y, color.RGBA{
|
||||
R: pixels[off+2], // B
|
||||
G: pixels[off+1], // G
|
||||
B: pixels[off], // R
|
||||
A: 255,
|
||||
})
|
||||
}
|
||||
}
|
||||
return img, nil
|
||||
}
|
||||
|
||||
// openPage opens a document and page, returning post-rotation dimensions
|
||||
// and a cleanup function. Callers must call closeAll() to free resources.
|
||||
func openPage(pdfData []byte, pageIdx int) (
|
||||
doc C.FPDF_DOCUMENT,
|
||||
page C.FPDF_PAGE,
|
||||
pw, ph float64,
|
||||
closeAll func(),
|
||||
err error,
|
||||
) {
|
||||
cData := C.CBytes(pdfData)
|
||||
|
||||
doc = C.FPDF_LoadMemDocument(unsafe.Pointer(cData), C.int(len(pdfData)), nil)
|
||||
if doc == nil {
|
||||
C.free(cData)
|
||||
err = fmt.Errorf("pdfium: FPDF_LoadMemDocument returned nil")
|
||||
return
|
||||
}
|
||||
|
||||
page = C.FPDF_LoadPage(doc, C.int(pageIdx))
|
||||
if page == nil {
|
||||
C.FPDF_CloseDocument(doc)
|
||||
C.free(cData)
|
||||
err = fmt.Errorf("pdfium: FPDF_LoadPage(%d) returned nil", pageIdx)
|
||||
return
|
||||
}
|
||||
|
||||
pw = float64(C.FPDF_GetPageWidth(page))
|
||||
ph = float64(C.FPDF_GetPageHeight(page))
|
||||
if pw <= 0 || ph <= 0 {
|
||||
C.FPDF_ClosePage(page)
|
||||
C.FPDF_CloseDocument(doc)
|
||||
C.free(cData)
|
||||
err = fmt.Errorf("pdfium: invalid page dimensions %.1fx%.1f", pw, ph)
|
||||
return
|
||||
}
|
||||
|
||||
closeAll = func() {
|
||||
C.FPDF_ClosePage(page)
|
||||
C.FPDF_CloseDocument(doc)
|
||||
C.free(cData)
|
||||
}
|
||||
return
|
||||
}
|
||||
241
internal/deepdoc/parser/pdf/pdfium/pdfium_test.go
Normal file
241
internal/deepdoc/parser/pdf/pdfium/pdfium_test.go
Normal file
@@ -0,0 +1,241 @@
|
||||
package pdfium
|
||||
|
||||
import (
|
||||
"image"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// testdataDir points at the shared test-pdf directory.
|
||||
var testdataDir = filepath.Join("..", "parser", "testdata", "pdfs")
|
||||
|
||||
func readPDF(t *testing.T, name string) []byte {
|
||||
t.Helper()
|
||||
data, err := os.ReadFile(filepath.Join(testdataDir, name))
|
||||
if err != nil {
|
||||
t.Fatalf("read %s: %v", name, err)
|
||||
}
|
||||
return data
|
||||
}
|
||||
|
||||
func TestRenderPage_EnglishSimple(t *testing.T) {
|
||||
data := readPDF(t, "01_english_simple.pdf")
|
||||
img, err := RenderPage(data, 0, 72)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
b := img.Bounds()
|
||||
t.Logf("01_english_simple.pdf @ 72 DPI: %dx%d", b.Dx(), b.Dy())
|
||||
if b.Dx() <= 0 || b.Dy() <= 0 {
|
||||
t.Errorf("expected non-zero dimensions, got %dx%d", b.Dx(), b.Dy())
|
||||
}
|
||||
// Must not be pure white (text should be present).
|
||||
if isPureWhite(img) {
|
||||
t.Error("rendered page is pure white — expected text content")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderPage_ChineseSimple(t *testing.T) {
|
||||
data := readPDF(t, "02_chinese_simple.pdf")
|
||||
img, err := RenderPage(data, 0, 72)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
b := img.Bounds()
|
||||
t.Logf("02_chinese_simple.pdf @ 72 DPI: %dx%d", b.Dx(), b.Dy())
|
||||
if b.Dx() <= 0 || b.Dy() <= 0 {
|
||||
t.Errorf("expected non-zero dimensions, got %dx%d", b.Dx(), b.Dy())
|
||||
}
|
||||
if isPureWhite(img) {
|
||||
t.Error("rendered page is pure white — expected text content")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderPage_MultiPage(t *testing.T) {
|
||||
data := readPDF(t, "03_multipage.pdf")
|
||||
// Render both pages.
|
||||
for pg := 0; pg < 2; pg++ {
|
||||
img, err := RenderPage(data, pg, 72)
|
||||
if err != nil {
|
||||
t.Fatalf("page %d: %v", pg, err)
|
||||
}
|
||||
b := img.Bounds()
|
||||
t.Logf("03_multipage.pdf page %d @ 72 DPI: %dx%d", pg, b.Dx(), b.Dy())
|
||||
if b.Dx() <= 0 || b.Dy() <= 0 {
|
||||
t.Errorf("page %d: expected non-zero dimensions", pg)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderPage_OutOfRange(t *testing.T) {
|
||||
data := readPDF(t, "01_english_simple.pdf")
|
||||
_, err := RenderPage(data, 99, 72)
|
||||
if err == nil {
|
||||
t.Error("expected error for out-of-range page index")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderPage_InvalidPDF(t *testing.T) {
|
||||
_, err := RenderPage([]byte("not a pdf"), 0, 72)
|
||||
if err == nil {
|
||||
t.Error("expected error for invalid PDF data")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderPage_EmptyData(t *testing.T) {
|
||||
_, err := RenderPage(nil, 0, 72)
|
||||
if err == nil {
|
||||
t.Error("expected error for nil data")
|
||||
}
|
||||
_, err = RenderPage([]byte{}, 0, 72)
|
||||
if err == nil {
|
||||
t.Error("expected error for empty data")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderPage_DPI(t *testing.T) {
|
||||
data := readPDF(t, "01_english_simple.pdf")
|
||||
|
||||
// Higher DPI → larger image.
|
||||
low, err := RenderPage(data, 0, 72)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
high, err := RenderPage(data, 0, 144)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
lw, lh := low.Bounds().Dx(), low.Bounds().Dy()
|
||||
hw, hh := high.Bounds().Dx(), high.Bounds().Dy()
|
||||
t.Logf("72 DPI: %dx%d 144 DPI: %dx%d", lw, lh, hw, hh)
|
||||
|
||||
if hw < lw*2-2 || hw > lw*2+2 {
|
||||
t.Errorf("144 DPI width %d not ≈ 2× 72 DPI width %d", hw, lw)
|
||||
}
|
||||
if hh < lh*2-2 || hh > lh*2+2 {
|
||||
t.Errorf("144 DPI height %d not ≈ 2× 72 DPI height %d", hh, lh)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRenderPage_AllTestPDFs(t *testing.T) {
|
||||
entries, err := os.ReadDir(testdataDir)
|
||||
if err != nil {
|
||||
t.Skipf("testdata dir not found: %v", err)
|
||||
}
|
||||
for _, e := range entries {
|
||||
if e.IsDir() || filepath.Ext(e.Name()) != ".pdf" {
|
||||
continue
|
||||
}
|
||||
data, err := os.ReadFile(filepath.Join(testdataDir, e.Name()))
|
||||
if err != nil {
|
||||
t.Errorf("%s: read: %v", e.Name(), err)
|
||||
continue
|
||||
}
|
||||
img, err := RenderPage(data, 0, 72)
|
||||
if err != nil {
|
||||
t.Errorf("%s: RenderPage: %v", e.Name(), err)
|
||||
continue
|
||||
}
|
||||
b := img.Bounds()
|
||||
if b.Dx() <= 0 || b.Dy() <= 0 {
|
||||
t.Errorf("%s: zero dimensions %dx%d", e.Name(), b.Dx(), b.Dy())
|
||||
}
|
||||
t.Logf("%s: %dx%d", e.Name(), b.Dx(), b.Dy())
|
||||
}
|
||||
}
|
||||
|
||||
func isPureWhite(img image.Image) bool {
|
||||
b := img.Bounds()
|
||||
for y := b.Min.Y; y < b.Max.Y; y++ {
|
||||
for x := b.Min.X; x < b.Max.X; x++ {
|
||||
r, g, b, _ := img.At(x, y).RGBA()
|
||||
// RGBA() returns premultiplied values in [0, 65535].
|
||||
if r>>8 < 250 || g>>8 < 250 || b>>8 < 250 {
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func TestPageSize(t *testing.T) {
|
||||
// Non-rotated A4: expect ~595×842
|
||||
data := readPDF(t, "rotate_0.pdf")
|
||||
w, h, err := PageSize(data, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if w < 500 || w > 700 || h < 700 || h > 900 {
|
||||
t.Errorf("rotate_0.pdf: got %.1f×%.1f, want ~595×842", w, h)
|
||||
}
|
||||
t.Logf("rotate_0.pdf: %.1f×%.1f pts", w, h)
|
||||
|
||||
// Rotate=90 A4: expect swapped ~842×595
|
||||
data90 := readPDF(t, "rotate_90.pdf")
|
||||
w90, h90, err := PageSize(data90, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if w90 < 700 || w90 > 950 || h90 < 500 || h90 > 700 {
|
||||
t.Errorf("rotate_90.pdf: got %.1f×%.1f, want ~842×595 (swapped)", w90, h90)
|
||||
}
|
||||
t.Logf("rotate_90.pdf: %.1f×%.1f pts (post-rotation)", w90, h90)
|
||||
|
||||
// Verify dimensions ARE swapped relative to Rotate=0
|
||||
if math.Abs(w-w90) < 50 {
|
||||
t.Errorf("Rotate=90 width %.1f not significantly different from Rotate=0 width %.1f — rotation not reflected?", w90, w)
|
||||
}
|
||||
if math.Abs(w-h90) > 2 || math.Abs(h-w90) > 2 {
|
||||
t.Errorf("Rotate=90 dimensions (%.1f×%.1f) are not swapped from Rotate=0 (%.1f×%.1f)", w90, h90, w, h)
|
||||
}
|
||||
|
||||
// Invalid page index
|
||||
_, _, err = PageSize(data, 999)
|
||||
if err == nil {
|
||||
t.Error("expected error for out-of-range page")
|
||||
}
|
||||
|
||||
// Empty data
|
||||
_, _, err = PageSize([]byte{}, 0)
|
||||
if err == nil {
|
||||
t.Error("expected error for empty PDF data")
|
||||
}
|
||||
}
|
||||
|
||||
// TestPdfiumConcurrentSafety verifies that the pdfiumMu mutex prevents
|
||||
// SIGSEGV from concurrent pdfium access. Without the mutex, 10 goroutines
|
||||
// calling PageSize/RenderPage simultaneously causes heap corruption within
|
||||
// milliseconds (empirically proven). If this test completes without
|
||||
// crashing, the mutex is working.
|
||||
func TestPdfiumConcurrentSafety(t *testing.T) {
|
||||
data := readPDF(t, "01_english_simple.pdf")
|
||||
|
||||
const goroutines = 10
|
||||
const iterations = 3
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < goroutines; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < iterations; j++ {
|
||||
if _, _, err := PageSize(data, 0); err != nil {
|
||||
t.Errorf("PageSize: %v", err)
|
||||
return
|
||||
}
|
||||
if img, err := RenderPage(data, 0, 72); err != nil {
|
||||
t.Errorf("RenderPage: %v", err)
|
||||
return
|
||||
} else if img.Bounds().Dx() <= 0 {
|
||||
t.Error("RenderPage returned zero-width image")
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
// Reaching here without SIGSEGV = mutex is effective.
|
||||
}
|
||||
Reference in New Issue
Block a user