nexarch/dms/internal/ocr/engine.go

package ocr

import (
	"bytes"
	"context"
	"fmt"
	"image"
	_ "image/jpeg"
	"image/png"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	_ "golang.org/x/image/tiff"

	"gitea.perlbach24.de/scripte/nexarch/dms/internal/ocr/imgproc"
)

// Engine kapselt die externen Werkzeuge (Tesseract, Poppler). Jeder Aufruf
// laeuft mit Zeitlimit und eigenem Temp-Verzeichnis.
type Engine struct {
	Tesseract   string
	Pdftoppm    string
	Pdftotext   string
	Pdfinfo     string
	Lang        string
	DPI         int
	PageTimeout time.Duration
}

func DefaultEngine() Engine {
	return Engine{Tesseract: "tesseract", Pdftoppm: "pdftoppm", Pdftotext: "pdftotext", Pdfinfo: "pdfinfo", Lang: "deu+eng", DPI: 300, PageTimeout: 2 * time.Minute}
}

func (e Engine) timeout() time.Duration {
	if e.PageTimeout > 0 {
		return e.PageTimeout
	}
	return 2 * time.Minute
}

func (e Engine) run(ctx context.Context, env []string, name string, args ...string) ([]byte, error) {
	ctx, cancel := context.WithTimeout(ctx, e.timeout())
	defer cancel()
	cmd := exec.CommandContext(ctx, name, args...)
	cmd.Env = append(os.Environ(), env...)
	cmd.WaitDelay = 5 * time.Second
	var stdout, stderr bytes.Buffer
	cmd.Stdout, cmd.Stderr = &stdout, &stderr
	if err := cmd.Run(); err != nil {
		if ctx.Err() != nil {
			return nil, fmt.Errorf("ocr: %s: zeitlimit %s ueberschritten: %w", filepath.Base(name), e.timeout(), ctx.Err())
		}
		return nil, fmt.Errorf("ocr: %s: %w (%s)", filepath.Base(name), err, strings.TrimSpace(stderr.String()))
	}
	return stdout.Bytes(), nil
}

// PDFPageCount liefert die Seitenzahl.
func (e Engine) PDFPageCount(ctx context.Context, pdfPath string) (int, error) {
	out, err := e.run(ctx, nil, e.Pdfinfo, pdfPath)
	if err != nil {
		return 0, err
	}
	m := regexp.MustCompile(`(?m)^Pages:\s+(\d+)`).FindSubmatch(out)
	if m == nil {
		return 0, fmt.Errorf("ocr: seitenzahl nicht ermittelbar")
	}
	return strconv.Atoi(string(m[1]))
}

// PDFText liefert den eingebetteten Text einer Seite (leer bei reinen Bild-Seiten).
func (e Engine) PDFText(ctx context.Context, pdfPath string, page int) (string, error) {
	out, err := e.run(ctx, nil, e.Pdftotext, "-f", strconv.Itoa(page), "-l", strconv.Itoa(page), "-layout", pdfPath, "-")
	if err != nil {
		return "", err
	}
	return normalizeText(string(out)), nil
}

// RenderPDFPage rendert eine Seite als Graustufenbild.
func (e Engine) RenderPDFPage(ctx context.Context, pdfPath string, page int) (image.Image, error) {
	dir, err := os.MkdirTemp("", "dms-ocr-render-*")
	if err != nil {
		return nil, err
	}
	defer func() { _ = os.RemoveAll(dir) }()
	prefix := filepath.Join(dir, "p")
	dpi := e.DPI
	if dpi <= 0 {
		dpi = 300
	}
	if _, err := e.run(ctx, nil, e.Pdftoppm, "-f", strconv.Itoa(page), "-l", strconv.Itoa(page), "-r", strconv.Itoa(dpi), "-gray", "-png", "-singlefile", pdfPath, prefix); err != nil {
		return nil, err
	}
	f, err := os.Open(prefix + ".png")
	if err != nil {
		return nil, fmt.Errorf("ocr: gerenderte seite lesen: %w", err)
	}
	defer func() { _ = f.Close() }()
	img, err := png.Decode(f)
	if err != nil {
		return nil, fmt.Errorf("ocr: gerenderte seite dekodieren: %w", err)
	}
	return img, nil
}

// PageResult ist das Ergebnis der Erkennung einer Seite.
type PageResult struct {
	Text       string
	Confidence float64
	AngleDeg   float64
}

// Recognize bereitet ein Bild auf (Rauschen, Entzerrung) und erkennt den Text.
func (e Engine) Recognize(ctx context.Context, img image.Image, preprocess bool) (PageResult, error) {
	return e.RecognizeLang(ctx, img, preprocess, e.Lang)
}

// RecognizeLang wie Recognize, mit explizitem Sprachmodell (leer: Standard der Engine).
func (e Engine) RecognizeLang(ctx context.Context, img image.Image, preprocess bool, lang string) (PageResult, error) {
	gray := imgproc.ToGray(img)
	var angle float64
	if preprocess {
		gray = imgproc.Median3(gray)
		gray, angle = imgproc.Deskew(gray)
	}
	dir, err := os.MkdirTemp("", "dms-ocr-*")
	if err != nil {
		return PageResult{}, err
	}
	defer func() { _ = os.RemoveAll(dir) }()
	in := filepath.Join(dir, "page.png")
	f, err := os.Create(in)
	if err != nil {
		return PageResult{}, err
	}
	if err := png.Encode(f, gray); err != nil {
		_ = f.Close()
		return PageResult{}, err
	}
	if err := f.Close(); err != nil {
		return PageResult{}, err
	}
	if lang == "" {
		lang = e.Lang
	}
	if lang == "" {
		lang = "deu+eng"
	}
	out, err := e.run(ctx, []string{"OMP_THREAD_LIMIT=1"}, e.Tesseract, in, "stdout", "-l", lang, "--psm", "3", "tsv")
	if err != nil {
		return PageResult{}, err
	}
	text, conf := parseTSV(string(out))
	return PageResult{Text: text, Confidence: conf, AngleDeg: angle}, nil
}

var installedLangs sync.Map

// HasModel meldet, ob Tesseract das Sprachmodell (z. B. "deu") installiert hat.
func (e Engine) HasModel(ctx context.Context, code string) bool {
	if code == "" {
		return false
	}
	if cached, ok := installedLangs.Load(e.Tesseract); ok {
		return cached.(map[string]bool)[code]
	}
	out, err := e.run(ctx, nil, e.Tesseract, "--list-langs")
	if err != nil {
		return false
	}
	set := map[string]bool{}
	for i, line := range strings.Split(string(out), "\n") {
		if line = strings.TrimSpace(line); i > 0 && line != "" {
			set[line] = true
		}
	}
	installedLangs.Store(e.Tesseract, set)
	return set[code]
}

// parseTSV baut aus der Wortliste von Tesseract Zeilen und die mittlere Konfidenz.
func parseTSV(tsv string) (string, float64) {
	type lineKey struct{ block, par, line int }
	lines := map[lineKey][]string{}
	order := []lineKey{}
	var sum float64
	var n int
	for i, row := range strings.Split(tsv, "\n") {
		if i == 0 || strings.TrimSpace(row) == "" {
			continue
		}
		cols := strings.Split(row, "\t")
		if len(cols) < 12 || cols[0] != "5" {
			continue
		}
		word := strings.TrimSpace(cols[11])
		conf, _ := strconv.ParseFloat(cols[10], 64)
		if word == "" || conf < 0 {
			continue
		}
		b, _ := strconv.Atoi(cols[2])
		p, _ := strconv.Atoi(cols[3])
		l, _ := strconv.Atoi(cols[4])
		k := lineKey{b, p, l}
		if _, ok := lines[k]; !ok {
			order = append(order, k)
		}
		lines[k] = append(lines[k], word)
		sum += conf
		n++
	}
	sort.SliceStable(order, func(i, j int) bool {
		if order[i].block != order[j].block {
			return order[i].block < order[j].block
		}
		if order[i].par != order[j].par {
			return order[i].par < order[j].par
		}
		return order[i].line < order[j].line
	})
	var b strings.Builder
	for _, k := range order {
		b.WriteString(strings.Join(lines[k], " "))
		b.WriteString("\n")
	}
	mean := 0.0
	if n > 0 {
		mean = sum / float64(n)
	}
	return strings.TrimSpace(b.String()), mean
}

func normalizeText(s string) string {
	s = strings.ReplaceAll(s, "\f", "")
	return strings.TrimSpace(strings.ReplaceAll(s, "\r\n", "\n"))
}