100 lines
2.2 KiB
Go
100 lines
2.2 KiB
Go
package ocr
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"image"
|
|
"image/png"
|
|
|
|
"github.com/gen2brain/go-fitz"
|
|
"github.com/otiai10/gosseract/v2"
|
|
)
|
|
|
|
type Box struct {
|
|
Text string
|
|
X int
|
|
Y int
|
|
W int
|
|
H int
|
|
}
|
|
|
|
// RenderPage rasterizes the first page of a PDF.
|
|
func RenderPage(pdfData []byte, dpi int) (image.Image, int, int, error) {
|
|
doc, err := fitz.NewFromMemory(pdfData)
|
|
if err != nil {
|
|
return nil, 0, 0, fmt.Errorf("open pdf: %w", err)
|
|
}
|
|
defer doc.Close()
|
|
|
|
if doc.NumPage() == 0 {
|
|
return nil, 0, 0, fmt.Errorf("pdf has no pages")
|
|
}
|
|
|
|
img, err := doc.ImageDPI(0, float64(dpi))
|
|
if err != nil {
|
|
return nil, 0, 0, fmt.Errorf("render page: %w", err)
|
|
}
|
|
bounds := img.Bounds()
|
|
return img, bounds.Dx(), bounds.Dy(), nil
|
|
}
|
|
|
|
// OCRImage runs Tesseract word-level OCR on a rendered page image.
|
|
func (e *Engine) OCRImage(img image.Image) ([]Box, error) {
|
|
var buf bytes.Buffer
|
|
if err := png.Encode(&buf, img); err != nil {
|
|
return nil, fmt.Errorf("encode png: %w", err)
|
|
}
|
|
|
|
client := e.Acquire()
|
|
defer e.Release(client)
|
|
|
|
if err := client.SetImageFromBytes(buf.Bytes()); err != nil {
|
|
return nil, fmt.Errorf("tesseract image: %w", err)
|
|
}
|
|
|
|
boxes, err := client.GetBoundingBoxes(gosseract.RIL_WORD)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("tesseract ocr: %w", err)
|
|
}
|
|
|
|
out := make([]Box, 0, len(boxes))
|
|
for _, b := range boxes {
|
|
text := trim(b.Word)
|
|
if text == "" || b.Confidence <= 0 {
|
|
continue
|
|
}
|
|
out = append(out, Box{
|
|
Text: text,
|
|
X: b.Box.Min.X,
|
|
Y: b.Box.Min.Y,
|
|
W: b.Box.Dx(),
|
|
H: b.Box.Dy(),
|
|
})
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// ExtractBoxes renders the first PDF page and runs OCR with a pooled Tesseract client.
|
|
func (e *Engine) ExtractBoxes(pdfData []byte, dpi int) ([]Box, int, int, error) {
|
|
img, pageW, pageH, err := RenderPage(pdfData, dpi)
|
|
if err != nil {
|
|
return nil, 0, 0, err
|
|
}
|
|
boxes, err := e.OCRImage(img)
|
|
if err != nil {
|
|
return nil, 0, 0, err
|
|
}
|
|
return boxes, pageW, pageH, nil
|
|
}
|
|
|
|
func trim(s string) string {
|
|
start, end := 0, len(s)
|
|
for start < end && (s[start] == ' ' || s[start] == '\t' || s[start] == '\n') {
|
|
start++
|
|
}
|
|
for end > start && (s[end-1] == ' ' || s[end-1] == '\t' || s[end-1] == '\n') {
|
|
end--
|
|
}
|
|
return s[start:end]
|
|
}
|