1package hints
2
3import (
4 "context"
5 "image"
6 "image/color"
7 "image/draw"
8 "os/exec"
9 "runtime"
10 "sort"
11 "sync"
12 "time"
13)
14
15type stripOffset struct {
16 source, packed image.Rectangle
17 factor int
18}
19type ocrStrip struct {
20 img *image.Gray
21 offsets []stripOffset
22}
23
24func packLines(src *image.RGBA, lines []image.Rectangle, workers int, normalize bool) []ocrStrip {
25 n := min(workers, len(lines))
26 if n == 0 {
27 return nil
28 }
29 groups := make([][]image.Rectangle, n)
30 areas := make([]int, n)
31 lines = append([]image.Rectangle(nil), lines...)
32 sort.Slice(lines, func(i, j int) bool { return lines[i].Dx()*lines[i].Dy() > lines[j].Dx()*lines[j].Dy() })
33 for _, r := range lines {
34 best := 0
35 for i := range n {
36 if areas[i] < areas[best] {
37 best = i
38 }
39 }
40 groups[best] = append(groups[best], r)
41 areas[best] += r.Dx() * r.Dy()
42 }
43 strips := make([]ocrStrip, n)
44 for i, group := range groups {
45 width, height := 0, 8
46 for _, r := range group {
47 factor := 1
48 if r.Dy() < 12 {
49 factor = 2
50 }
51 crop := r.Inset(-4).Intersect(src.Bounds())
52 packed := image.Rect(4, height, 4+crop.Dx()*factor, height+crop.Dy()*factor)
53 strips[i].offsets = append(strips[i].offsets, stripOffset{crop, packed, factor})
54 width = max(width, packed.Max.X+4)
55 height = packed.Max.Y + 8
56 }
57 dst := image.NewGray(image.Rect(0, 0, width, height))
58 draw.Draw(dst, dst.Bounds(), image.NewUniform(color.White), image.Point{}, draw.Src)
59 for _, off := range strips[i].offsets {
60 // The padded perimeter estimates the background without depending on text colour.
61 var sum, count int
62 r := off.source
63 for x := r.Min.X; x < r.Max.X; x++ {
64 sum += int(color.GrayModel.Convert(src.At(x, r.Min.Y)).(color.Gray).Y)
65 count++
66 }
67 invert := normalize && count > 0 && sum/count < 128
68 for y := off.packed.Min.Y; y < off.packed.Max.Y; y++ {
69 for x := off.packed.Min.X; x < off.packed.Max.X; x++ {
70 v := color.GrayModel.Convert(src.At(r.Min.X+(x-off.packed.Min.X)/off.factor, r.Min.Y+(y-off.packed.Min.Y)/off.factor)).(color.Gray).Y
71 if invert {
72 v = 255 - v
73 }
74 dst.SetGray(x, y, color.Gray{Y: v})
75 }
76 }
77 }
78 strips[i].img = dst
79 }
80 return strips
81}
82
83func readStrips(ctx context.Context, img *image.RGBA, lines []image.Rectangle, workers int, normalize bool) ([]OCRWord, int, error) {
84 binary, err := exec.LookPath("tesseract")
85 if err != nil {
86 return nil, 0, err
87 }
88 ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
89 defer cancel()
90 lang, err := ocrLanguage(ctx, binary)
91 if err != nil {
92 return nil, 0, err
93 }
94 strips := packLines(img, lines, min(workers, runtime.GOMAXPROCS(0), 16), normalize)
95 results := make([][]OCRWord, len(strips))
96 errs := make([]error, len(strips))
97 var wg sync.WaitGroup
98 for i, s := range strips {
99 wg.Go(func() {
100 words, e := recognizeImage(ctx, binary, lang, s.img, "6", normalize)
101 errs[i] = e
102 for _, w := range words {
103 center := image.Pt((w.Rect.Min.X+w.Rect.Max.X)/2, (w.Rect.Min.Y+w.Rect.Max.Y)/2)
104 for _, off := range s.offsets {
105 if center.In(off.packed) {
106 r := w.Rect.Sub(off.packed.Min)
107 w.Rect = image.Rect(r.Min.X/off.factor, r.Min.Y/off.factor, (r.Max.X+off.factor-1)/off.factor, (r.Max.Y+off.factor-1)/off.factor).Add(off.source.Min)
108 results[i] = append(results[i], w)
109 break
110 }
111 }
112 }
113 })
114 }
115 wg.Wait()
116 var words []OCRWord
117 for i := range strips {
118 if errs[i] != nil {
119 return nil, len(strips), errs[i]
120 }
121 words = append(words, results[i]...)
122 }
123 return words, len(strips), nil
124}