1package hints
2
3import (
4 "bufio"
5 "bytes"
6 "context"
7 "fmt"
8 "image"
9 "image/png"
10 "io"
11 "math"
12 "net"
13 "os"
14 "os/exec"
15 "path/filepath"
16 "regexp"
17 "runtime"
18 "sort"
19 "strconv"
20 "strings"
21 "sync"
22 "time"
23
24 "github.com/bnema/neferwl/internal/adapters/logging"
25 "github.com/bnema/neferwl/internal/ports"
26)
27
28// OCRWord is a recognized word in physical output coordinates.
29type OCRWord struct {
30 Text string `json:"text"`
31 Rect image.Rectangle `json:"rect"`
32 Confidence float64 `json:"confidence"`
33 // Underlined is set when a thin line runs under the word. UnderL and
34 // UnderR are the x extent of that line (crop-local), which can reach the
35 // neighbouring words of one underlined link.
36 Underlined bool `json:"underlined,omitempty"`
37 UnderL int `json:"under_l,omitempty"`
38 UnderR int `json:"under_r,omitempty"`
39}
40
41var languageWarning sync.Once
42var failureWarning sync.Once
43
44// languageRunner bounds the external language-discovery process independently
45// of recognition, which has not started yet when discovery runs.
46type languageRunner interface {
47 Output(context.Context, string, ...string) ([]byte, error)
48}
49
50type execLanguageRunner struct{}
51
52func (execLanguageRunner) Output(ctx context.Context, binary string, args ...string) ([]byte, error) {
53 return exec.CommandContext(ctx, binary, args...).Output()
54}
55
56const languageTimeout = 5 * time.Second
57
58func ocrLanguage(ctx context.Context, binary string) (string, error) {
59 return discoverLanguage(ctx, binary, execLanguageRunner{}, languageTimeout)
60}
61
62func discoverLanguage(ctx context.Context, binary string, runner languageRunner, timeout time.Duration) (string, error) {
63 dir := os.Getenv("TESSDATA_PREFIX")
64 if dir == "" {
65 dir = "/usr/share/tessdata"
66 }
67 if _, err := os.Stat(filepath.Join(dir, "eng.traineddata")); err == nil {
68 return "eng", nil
69 }
70 ctx, cancel := context.WithTimeout(ctx, timeout)
71 defer cancel()
72 out, err := runner.Output(ctx, binary, "--list-langs")
73 if err != nil {
74 return "", err
75 }
76 for _, line := range strings.Split(string(out), "\n")[1:] {
77 lang := strings.TrimSpace(line)
78 if lang != "" && lang != "osd" {
79 languageWarning.Do(func() {
80 log := logging.For(ctx, "hints")
81 log.Warn().Str("language", lang).Msg("English OCR data missing; install tesseract-data-eng")
82 })
83 return lang, nil
84 }
85 }
86 return "", fmt.Errorf("no OCR language available")
87}
88
89func parseTSV(data []byte, offset image.Point) ([]OCRWord, error) {
90 var words []OCRWord
91 scanner := bufio.NewScanner(bytes.NewReader(data))
92 scanner.Buffer(make([]byte, 4096), 1024*1024)
93 for scanner.Scan() {
94 fields := strings.SplitN(scanner.Text(), "\t", 12)
95 if len(fields) != 12 || fields[0] != "5" {
96 continue
97 }
98 var values [4]int
99 for i := range values {
100 v, err := strconv.Atoi(fields[6+i])
101 if err != nil {
102 return nil, fmt.Errorf("OCR TSV coordinate: %w", err)
103 }
104 values[i] = v
105 }
106 conf, err := strconv.ParseFloat(fields[10], 64)
107 if err != nil {
108 return nil, fmt.Errorf("OCR TSV confidence: %w", err)
109 }
110 if values[2] <= 0 || values[3] <= 0 || strings.TrimSpace(fields[11]) == "" {
111 continue
112 }
113 r := image.Rect(values[0], values[1], values[0]+values[2], values[1]+values[3]).Add(offset)
114 words = append(words, OCRWord{Text: fields[11], Rect: r, Confidence: conf})
115 }
116 return words, scanner.Err()
117}
118
119// ocrWorkers is the maximum number of bands recognized in parallel. 4 keeps
120// the tesseract processes at about 60% of the memory of 8 for the same time.
121var ocrWorkers = 4
122
123// pngLevel is the compression of non-grayscale images piped to tesseract.
124var pngLevel = png.BestSpeed
125
126// rawPGM feeds grayscale images to tesseract as binary PGM instead of PNG.
127var rawPGM = true
128
129// ocrState is the tesseract installation, resolved once.
130type ocrState struct {
131 done bool
132 binary string
133 lang string
134 err error
135}
136
137func (s *ocrState) resolve(ctx context.Context) error {
138 if !s.done {
139 s.done = true
140 if s.binary, s.err = exec.LookPath("tesseract"); s.err == nil {
141 s.lang, s.err = ocrLanguage(ctx, s.binary)
142 }
143 }
144 return s.err
145}
146
147// ocrReady resolves tesseract on the first call. When it is missing, OCR is
148// disabled for the life of the process and the warning is logged once.
149func (w *worker) ocrReady(ctx context.Context) bool {
150 if err := w.ocr.resolve(ctx); err != nil {
151 failureWarning.Do(func() {
152 log := logging.For(ctx, "hints")
153 log.Warn().Err(err).Msg("OCR unavailable; install tesseract and tesseract-data-eng")
154 })
155 return false
156 }
157 return true
158}
159
160// band is a horizontal slice of a window crop recognized by one tesseract
161// process. It reads rows [top,bottom) and owns the words centred in
162// [ownTop,ownBottom); the margin only supplies context.
163type band struct {
164 idx, n int
165 top, bottom int
166 ownTop, ownBottom int
167}
168
169func bandsOf(height int) []band {
170 n := min(runtime.GOMAXPROCS(0), ocrWorkers, height)
171 bands := make([]band, 0, max(n, 0))
172 for i := range n {
173 ownTop, ownBottom := i*height/n, (i+1)*height/n
174 bands = append(bands, band{idx: i, n: n, top: max(0, ownTop-40), bottom: min(height, ownBottom+40), ownTop: ownTop, ownBottom: ownBottom})
175 }
176 return bands
177}
178
179// recognize reads the bands of jobs that missed the cache, in parallel.
180func (w *worker) recognize(ctx context.Context, jobs []*bandJob) {
181 ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
182 defer cancel()
183 // One scan-wide pool: at most ocrWorkers tesseract processes run at once,
184 // whatever the number of windows.
185 sem := make(chan struct{}, max(ocrWorkers, 1))
186 var wg sync.WaitGroup
187 for _, j := range jobs {
188 if j.hit {
189 continue
190 }
191 select {
192 case sem <- struct{}{}:
193 case <-ctx.Done():
194 j.err = ctx.Err()
195 continue
196 }
197 wg.Go(func() {
198 defer func() { <-sem }()
199 j.words, j.err = readBand(ctx, &w.ocr, j)
200 })
201 }
202 wg.Wait()
203}
204
205// readBand recognizes one band; the words are crop-local and owned by it.
206func readBand(ctx context.Context, st *ocrState, j *bandJob) ([]OCRWord, error) {
207 words, err := recognizeBand(ctx, st.binary, st.lang, j.gray)
208 if err != nil {
209 return nil, err
210 }
211 heights := make([]int, 0, len(words))
212 for _, word := range words {
213 if len(word.Text) >= 3 {
214 heights = append(heights, word.Rect.Dy())
215 }
216 }
217 sort.Ints(heights)
218 if len(heights) > 0 && heights[len(heights)/2] < 12 {
219 words, err = recognizeBand(ctx, st.binary, st.lang, upscaleBand(j.gray))
220 if err != nil {
221 return nil, err
222 }
223 for i := range words {
224 r := words[i].Rect
225 words[i].Rect = image.Rect(r.Min.X/2, r.Min.Y/2, (r.Max.X+1)/2, (r.Max.Y+1)/2)
226 }
227 }
228 markUnderlined(j.gray, words)
229 owned := words[:0]
230 for _, word := range words {
231 // Overlap supplies context, but each word belongs to exactly one band.
232 word.Rect = word.Rect.Add(image.Pt(0, j.ext.Min.Y))
233 if center := (word.Rect.Min.Y + word.Rect.Max.Y) / 2; center >= j.own[0] && center < j.own[1] {
234 owned = append(owned, word)
235 }
236 }
237 return owned, nil
238}
239
240// readWords recognizes the words of img, a window crop, in parallel bands.
241func readWords(ctx context.Context, img *image.Gray) ([]OCRWord, int, error) {
242 w := &worker{}
243 if err := w.ocr.resolve(ctx); err != nil {
244 return nil, 0, err
245 }
246 var jobs []*bandJob
247 for _, b := range bandsOf(img.Bounds().Dy()) {
248 ext := image.Rect(0, b.top, img.Bounds().Dx(), b.bottom)
249 jobs = append(jobs, &bandJob{ext: ext, own: [2]int{b.ownTop, b.ownBottom}, gray: img.SubImage(ext).(*image.Gray)})
250 }
251 w.recognize(ctx, jobs)
252 var words []OCRWord
253 for _, j := range jobs {
254 if j.err != nil {
255 return nil, len(jobs), j.err
256 }
257 for _, word := range j.words {
258 if !duplicateWord(words, word) {
259 words = append(words, word)
260 }
261 }
262 }
263 return words, len(jobs), nil
264}
265
266// The link patterns compile on the first OCR link check, not at process
267// start: with hint mode off nothing here costs memory or CPU.
268var domainPattern = sync.OnceValue(func() *regexp.Regexp {
269 return regexp.MustCompile(`(?i)^(?:[a-z0-9](?:[a-z0-9-]*[a-z0-9])?\.)+([a-z]{2,24})(?::[0-9]+)?(?:[/?#]\S*)?$`)
270})
271var schemePattern = sync.OnceValue(func() *regexp.Regexp { return regexp.MustCompile(`(?i)^(?:https?|htps|file|ftp):/{1,2}\S+$`) })
272var emailPattern = sync.OnceValue(func() *regexp.Regexp {
273 return regexp.MustCompile(`(?i)^[a-z0-9.!#$%&'*+/=?^_` + "`" + `{|}~-]+@(?:[a-z0-9-]+\.)+[a-z]{2,24}$`)
274})
275var fileExtensions = sync.OnceValue(func() []string {
276 return strings.Fields("go md txt toml json yaml yml rs py js ts sh c h png jpg log conf lock mod sum html css")
277})
278
279func isTextLink(text string) bool {
280 text = strings.Trim(text, "()[]<>.,;:'\"")
281 if schemePattern().MatchString(text) || (strings.HasPrefix(strings.ToLower(text), "www.") && len(text) > 4) || emailPattern().MatchString(text) {
282 return true
283 }
284 host := text
285 if h, _, err := net.SplitHostPort(text); err == nil {
286 host = h
287 }
288 if ip := net.ParseIP(host); ip != nil && ip.To4() != nil {
289 return true
290 }
291 m := domainPattern().FindStringSubmatch(text)
292 if m == nil {
293 return false
294 }
295 for _, ext := range fileExtensions() {
296 if strings.EqualFold(m[1], ext) {
297 return false
298 }
299 }
300 return true
301}
302
303func textLinks(words []OCRWord, scale float64) []ports.HintTarget {
304 if scale <= 0 {
305 scale = 1
306 }
307 words = append([]OCRWord(nil), words...)
308 sort.Slice(words, func(i, j int) bool {
309 if words[i].Rect.Min.Y != words[j].Rect.Min.Y {
310 return words[i].Rect.Min.Y < words[j].Rect.Min.Y
311 }
312 return words[i].Rect.Min.X < words[j].Rect.Min.X
313 })
314 var out []ports.HintTarget
315 var urls []image.Rectangle // boxes of the URL targets, in pixels
316 used := make([]bool, len(words))
317 for i, word := range words {
318 if used[i] {
319 continue
320 }
321 text, rect := word.Text, word.Rect
322 // Grow only touching words on the same baseline; ordinary prose is not joined.
323 for {
324 next := -1
325 for j, candidate := range words {
326 if j == i || used[j] || candidate.Rect.Min.X < rect.Max.X-2 {
327 continue
328 }
329 gap := candidate.Rect.Min.X - rect.Max.X
330 if gap > max(3, min(rect.Dy(), candidate.Rect.Dy())/5) {
331 continue
332 }
333 if abs(candidate.Rect.Max.Y-rect.Max.Y) > max(3, rect.Dy()/3) {
334 continue
335 }
336 if next < 0 || candidate.Rect.Min.X < words[next].Rect.Min.X {
337 next = j
338 }
339 }
340 if next < 0 {
341 break
342 }
343 text += words[next].Text
344 rect = rect.Union(words[next].Rect)
345 used[next] = true
346 }
347 if !isTextLink(text) {
348 continue
349 }
350 x, y := int(math.Floor(float64(rect.Min.X)/scale)), int(math.Floor(float64(rect.Min.Y)/scale))
351 right, bottom := int(math.Ceil(float64(rect.Max.X)/scale)), int(math.Ceil(float64(rect.Max.Y)/scale))
352 out = append(out, ports.HintTarget{Kind: ports.HintLink, Text: true, Rect: ports.Rect{X: x, Y: y, W: right - x, H: bottom - y}})
353 urls = append(urls, rect)
354 }
355 return append(out, underlinedLinks(words, urls, scale)...)
356}
357
358func abs(x int) int {
359 if x < 0 {
360 return -x
361 }
362 return x
363}
364func overlap(a, b image.Rectangle) float64 {
365 intersection := a.Intersect(b)
366 area := min(a.Dx()*a.Dy(), b.Dx()*b.Dy())
367 if intersection.Empty() || area <= 0 {
368 return 0
369 }
370 return float64(intersection.Dx()*intersection.Dy()) / float64(area)
371}
372func targetRect(t ports.HintTarget) image.Rectangle {
373 return image.Rect(t.Rect.X, t.Rect.Y, t.Rect.X+t.Rect.W, t.Rect.Y+t.Rect.H)
374}
375func recognizeBand(ctx context.Context, binary, lang string, img image.Image) ([]OCRWord, error) {
376 return recognizeImage(ctx, binary, lang, img, "11", false)
377}
378
379// pgmReader streams a grayscale image as binary PGM (P5), row by row.
380type pgmReader struct {
381 head []byte
382 img *image.Gray
383 y int
384 row []byte
385}
386
387func newPGMReader(img *image.Gray) *pgmReader {
388 b := img.Bounds()
389 return &pgmReader{head: fmt.Appendf(nil, "P5\n%d %d\n255\n", b.Dx(), b.Dy()), img: img}
390}
391
392func (r *pgmReader) Read(p []byte) (int, error) {
393 n := 0
394 for n < len(p) {
395 if len(r.head) > 0 {
396 c := copy(p[n:], r.head)
397 r.head, n = r.head[c:], n+c
398 continue
399 }
400 if len(r.row) == 0 {
401 b := r.img.Bounds()
402 if r.y >= b.Dy() {
403 break
404 }
405 r.row = r.img.Pix[(r.y+b.Min.Y-r.img.Rect.Min.Y)*r.img.Stride+(b.Min.X-r.img.Rect.Min.X):][:b.Dx()]
406 r.y++
407 }
408 c := copy(p[n:], r.row)
409 r.row, n = r.row[c:], n+c
410 }
411 if n == 0 {
412 return 0, io.EOF
413 }
414 return n, nil
415}
416
417func recognizeImage(ctx context.Context, binary, lang string, img image.Image, psm string, normalized bool) ([]OCRWord, error) {
418 args := []string{"stdin", "stdout", "--psm", psm, "-l", lang}
419 if normalized {
420 args = append(args, "-c", "tessedit_do_invert=0")
421 }
422 args = append(args, "tsv")
423 cmd := exec.CommandContext(ctx, binary, args...)
424 cmd.Env = append(os.Environ(), "OMP_THREAD_LIMIT=1")
425 if g, ok := img.(*image.Gray); ok && rawPGM {
426 // Raw grayscale: no deflate, no encoder buffers.
427 cmd.Stdin = newPGMReader(g)
428 } else {
429 // The PNG streams into tesseract; no encoded copy is held.
430 pr, pw := io.Pipe()
431 defer func() { _ = pr.Close() }()
432 go func() {
433 enc := png.Encoder{CompressionLevel: pngLevel}
434 _ = pw.CloseWithError(enc.Encode(pw, img))
435 }()
436 cmd.Stdin = pr
437 }
438 var stderr bytes.Buffer
439 cmd.Stderr = &stderr
440 output, err := cmd.Output()
441 if err != nil {
442 return nil, fmt.Errorf("tesseract: %w: %s", err, stderr.String())
443 }
444 return parseTSV(output, image.Point{})
445}
446
447// upscaleBand doubles a grayscale band with nearest-neighbour sampling.
448func upscaleBand(src *image.Gray) *image.Gray {
449 b := src.Bounds()
450 dst := image.NewGray(image.Rect(0, 0, b.Dx()*2, b.Dy()*2))
451 for y := 0; y < b.Dy(); y++ {
452 row := src.Pix[(y+b.Min.Y-src.Rect.Min.Y)*src.Stride+(b.Min.X-src.Rect.Min.X):][:b.Dx()]
453 top := dst.Pix[2*y*dst.Stride:][:dst.Stride]
454 for x, v := range row {
455 top[2*x], top[2*x+1] = v, v
456 }
457 copy(dst.Pix[(2*y+1)*dst.Stride:], top)
458 }
459 return dst
460}