mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-06-29 15:31:05 +08:00
239 lines
6.6 KiB
Go
239 lines
6.6 KiB
Go
|
package parser
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"image"
|
|||
|
|
"testing"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// mockRotationDoc implements DocAnalyzer with deterministic OCR results per angle.
|
|||
|
|
// The mock tracks the call sequence: evaluateTableOrientation tests angles in
|
|||
|
|
// order 0°, 90°, 180°, 270°. Each call to OCRDetect increments an internal
|
|||
|
|
// counter and returns data for the corresponding angle.
|
|||
|
|
type mockRotationDoc struct {
|
|||
|
|
// angle → {regions count, average confidence, error}
|
|||
|
|
angles map[int]struct {
|
|||
|
|
regions int
|
|||
|
|
avgConf float64
|
|||
|
|
err error
|
|||
|
|
}
|
|||
|
|
callSeq int // incremented per OCRDetect call, selects the angle's data
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
var rotationOrder = []int{0, 90, 180, 270}
|
|||
|
|
|
|||
|
|
func (m *mockRotationDoc) DLA(_ context.Context, _ image.Image) ([]DLARegion, error) { return nil, nil }
|
|||
|
|
func (m *mockRotationDoc) TSR(_ context.Context, _ image.Image) ([]TSRCell, error) { return nil, nil }
|
|||
|
|
func (m *mockRotationDoc) OCR(_ image.Image) (string, error) { return "", nil }
|
|||
|
|
func (m *mockRotationDoc) Health() bool { return true }
|
|||
|
|
func (m *mockRotationDoc) ModelType() ModelType { return ModelSaas }
|
|||
|
|
|
|||
|
|
func (m *mockRotationDoc) currentAngle() int {
|
|||
|
|
idx := m.callSeq % len(rotationOrder)
|
|||
|
|
return rotationOrder[idx]
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (m *mockRotationDoc) OCRDetect(_ context.Context, img image.Image) ([]OCRBox, error) {
|
|||
|
|
defer func() { m.callSeq++ }()
|
|||
|
|
angle := m.currentAngle()
|
|||
|
|
cfg, ok := m.angles[angle]
|
|||
|
|
if !ok {
|
|||
|
|
cfg = m.angles[0] // fallback to 0° config
|
|||
|
|
}
|
|||
|
|
if cfg.err != nil {
|
|||
|
|
return nil, cfg.err
|
|||
|
|
}
|
|||
|
|
if cfg.regions == 0 {
|
|||
|
|
return nil, nil
|
|||
|
|
}
|
|||
|
|
w, h := img.Bounds().Dx(), img.Bounds().Dy()
|
|||
|
|
boxes := make([]OCRBox, cfg.regions)
|
|||
|
|
step := w / (cfg.regions + 1)
|
|||
|
|
for i := 0; i < cfg.regions; i++ {
|
|||
|
|
x := step * (i + 1)
|
|||
|
|
boxes[i] = OCRBox{
|
|||
|
|
X0: float64(x), Y0: float64(h / 4),
|
|||
|
|
X1: float64(x + 20), Y1: float64(h / 4),
|
|||
|
|
X2: float64(x + 20), Y2: float64(h * 3 / 4),
|
|||
|
|
X3: float64(x), Y3: float64(h * 3 / 4),
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return boxes, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (m *mockRotationDoc) OCRRecognizeBatch(_ context.Context, cropped []image.Image) ([][]OCRText, []error) {
|
|||
|
|
results := make([][]OCRText, len(cropped))
|
|||
|
|
errs := make([]error, len(cropped))
|
|||
|
|
for i, img := range cropped {
|
|||
|
|
results[i], errs[i] = m.OCRRecognize(context.Background(), img)
|
|||
|
|
}
|
|||
|
|
return results, errs
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (m *mockRotationDoc) OCRRecognize(_ context.Context, _ image.Image) ([]OCRText, error) {
|
|||
|
|
angle := rotationOrder[(m.callSeq-1)%len(rotationOrder)] // use angle from last Detect call
|
|||
|
|
cfg, ok := m.angles[angle]
|
|||
|
|
if !ok {
|
|||
|
|
cfg = m.angles[0]
|
|||
|
|
}
|
|||
|
|
if cfg.err != nil {
|
|||
|
|
return nil, cfg.err
|
|||
|
|
}
|
|||
|
|
if cfg.regions == 0 {
|
|||
|
|
return nil, nil
|
|||
|
|
}
|
|||
|
|
texts := make([]OCRText, cfg.regions)
|
|||
|
|
for i := 0; i < cfg.regions; i++ {
|
|||
|
|
texts[i] = OCRText{Text: "X", Confidence: cfg.avgConf}
|
|||
|
|
}
|
|||
|
|
return texts, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func makeTestTableImage() image.Image {
|
|||
|
|
return image.NewRGBA(image.Rect(0, 0, 200, 100))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func TestEvaluateTableOrientation(t *testing.T) {
|
|||
|
|
t.Run("normal table 0° wins", func(t *testing.T) {
|
|||
|
|
doc := &mockRotationDoc{
|
|||
|
|
angles: map[int]struct {
|
|||
|
|
regions int
|
|||
|
|
avgConf float64
|
|||
|
|
err error
|
|||
|
|
}{
|
|||
|
|
0: {regions: 10, avgConf: 0.9},
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
angle, _, scores := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
|
|||
|
|
if angle != 0 {
|
|||
|
|
t.Errorf("expected 0°, got %d° (scores: %v)", angle, scores)
|
|||
|
|
}
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
t.Run("90° rotated table wins", func(t *testing.T) {
|
|||
|
|
doc := &mockRotationDoc{
|
|||
|
|
angles: map[int]struct {
|
|||
|
|
regions int
|
|||
|
|
avgConf float64
|
|||
|
|
err error
|
|||
|
|
}{
|
|||
|
|
0: {regions: 2, avgConf: 0.2},
|
|||
|
|
90: {regions: 10, avgConf: 0.9},
|
|||
|
|
180: {regions: 2, avgConf: 0.2},
|
|||
|
|
270: {regions: 2, avgConf: 0.2},
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
angle, _, scores := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
|
|||
|
|
if angle != 90 {
|
|||
|
|
t.Errorf("expected 90°, got %d° (scores: %v)", angle, scores)
|
|||
|
|
}
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
t.Run("180° rotated table wins", func(t *testing.T) {
|
|||
|
|
doc := &mockRotationDoc{
|
|||
|
|
angles: map[int]struct {
|
|||
|
|
regions int
|
|||
|
|
avgConf float64
|
|||
|
|
err error
|
|||
|
|
}{
|
|||
|
|
0: {regions: 1, avgConf: 0.1},
|
|||
|
|
90: {regions: 1, avgConf: 0.1},
|
|||
|
|
180: {regions: 8, avgConf: 0.85},
|
|||
|
|
270: {regions: 1, avgConf: 0.1},
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
angle, _, scores := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
|
|||
|
|
if angle != 180 {
|
|||
|
|
t.Errorf("expected 180°, got %d° (scores: %v)", angle, scores)
|
|||
|
|
}
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
t.Run("270° rotated table wins", func(t *testing.T) {
|
|||
|
|
doc := &mockRotationDoc{
|
|||
|
|
angles: map[int]struct {
|
|||
|
|
regions int
|
|||
|
|
avgConf float64
|
|||
|
|
err error
|
|||
|
|
}{
|
|||
|
|
0: {regions: 1, avgConf: 0.1},
|
|||
|
|
90: {regions: 1, avgConf: 0.1},
|
|||
|
|
180: {regions: 1, avgConf: 0.1},
|
|||
|
|
270: {regions: 9, avgConf: 0.88},
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
angle, _, scores := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
|
|||
|
|
if angle != 270 {
|
|||
|
|
t.Errorf("expected 270°, got %d° (scores: %v)", angle, scores)
|
|||
|
|
}
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
t.Run("threshold protection — 0° keeps when diff too small", func(t *testing.T) {
|
|||
|
|
// Region-count scoring: 8 vs 9 is too close (< 1.4×) → 0° wins.
|
|||
|
|
doc := &mockRotationDoc{
|
|||
|
|
angles: map[int]struct {
|
|||
|
|
regions int
|
|||
|
|
avgConf float64
|
|||
|
|
err error
|
|||
|
|
}{
|
|||
|
|
0: {regions: 8},
|
|||
|
|
90: {regions: 9},
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
angle, _, _ := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
|
|||
|
|
if angle != 0 {
|
|||
|
|
t.Errorf("expected 0° (threshold protection), got %d°", angle)
|
|||
|
|
}
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
t.Run("threshold pass — 90° wins when region count is clearly higher", func(t *testing.T) {
|
|||
|
|
// 0° has few regions AND 90° has ≥1.4× more → 90° wins.
|
|||
|
|
doc := &mockRotationDoc{
|
|||
|
|
angles: map[int]struct {
|
|||
|
|
regions int
|
|||
|
|
avgConf float64
|
|||
|
|
err error
|
|||
|
|
}{
|
|||
|
|
0: {regions: 4},
|
|||
|
|
90: {regions: 10},
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
angle, _, _ := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
|
|||
|
|
if angle != 90 {
|
|||
|
|
t.Errorf("expected 90° (threshold passed), got %d°", angle)
|
|||
|
|
}
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
t.Run("all angles fail OCR → fallback 0°", func(t *testing.T) {
|
|||
|
|
doc := &mockRotationDoc{
|
|||
|
|
angles: map[int]struct {
|
|||
|
|
regions int
|
|||
|
|
avgConf float64
|
|||
|
|
err error
|
|||
|
|
}{
|
|||
|
|
0: {err: errMockOCR},
|
|||
|
|
90: {err: errMockOCR},
|
|||
|
|
180: {err: errMockOCR},
|
|||
|
|
270: {err: errMockOCR},
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
angle, img, scores := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
|
|||
|
|
if angle != 0 {
|
|||
|
|
t.Errorf("expected 0° fallback, got %d°", angle)
|
|||
|
|
}
|
|||
|
|
if img == nil {
|
|||
|
|
t.Error("expected non-nil fallback image")
|
|||
|
|
}
|
|||
|
|
for _, s := range scores {
|
|||
|
|
if s != 0 {
|
|||
|
|
t.Error("all scores should be 0 on OCR failure")
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
})
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
var errMockOCR = &mockError{"mock OCR failure"}
|
|||
|
|
|
|||
|
|
type mockError struct{ msg string }
|
|||
|
|
|
|||
|
|
func (e *mockError) Error() string { return e.msg }
|