// Package ocr reads Japanese text out of screenshots through a helper
// process, chirami-ocr, started on demand: Live Text on macOS, ONNX models on
// Linux. Keeping the engines out of process keeps chirami small while idle.
//
// The protocol on the helper's stdin/stdout: each request is a 4-byte
// big-endian length followed by that many bytes of PNG; each reply is one
// line of JSON, a Result. The helper exits when stdin closes.
package ocr
import (
"bufio"
"context"
"encoding/binary"
"encoding/json"
"errors"
"fmt"
"io"
"os"
"os/exec"
"path/filepath"
"strings"
"sync"
"time"
)
// Line is one line or column of recognised text, in reading order.
type Line struct {
Text string `json:"text"`
Vertical bool `json:"vertical,omitempty"`
}
// Result is what the helper read in one image.
type Result struct {
Lines []Line `json:"lines"`
Engine string `json:"engine,omitempty"` // which engine read it, for logs
Error string `json:"error,omitempty"`
}
// Text joins the lines the way Japanese is written: without spaces.
func (r Result) Text() string {
var b strings.Builder
for _, l := range r.Lines {
b.WriteString(strings.TrimSpace(l.Text))
}
return b.String()
}
// HelperName is the helper binary, looked up next to the chirami executable
// first (as invoked, then with symlinks resolved, since nix links both into
// one bin directory), then on $PATH.
const HelperName = "chirami-ocr"
// FindHelper returns the helper's path.
func FindHelper() (string, error) {
var dirs []string
if strings.ContainsRune(os.Args[0], filepath.Separator) {
dirs = append(dirs, filepath.Dir(os.Args[0]))
}
if exe, err := os.Executable(); err == nil {
dirs = append(dirs, filepath.Dir(exe))
}
for _, d := range dirs {
p := filepath.Join(d, HelperName)
if _, err := os.Stat(p); err == nil {
return p, nil
}
}
return exec.LookPath(HelperName)
}
// Client talks to one helper process, starting it on first use and
// stopping it after Idle without requests.
type Client struct {
Path string // helper binary
Args []string // e.g. the model directory
Idle time.Duration
mu sync.Mutex
cmd *exec.Cmd
in io.WriteCloser
out *bufio.Reader
timer *time.Timer
}
// Recognize reads the text in a PNG image.
func (c *Client) Recognize(ctx context.Context, png []byte) (Result, error) {
c.mu.Lock()
defer c.mu.Unlock()
if c.cmd == nil {
if err := c.start(); err != nil {
return Result{}, err
}
}
if c.timer != nil {
c.timer.Stop()
}
type reply struct {
r Result
err error
}
done := make(chan reply, 1)
go func() {
var hdr [4]byte
binary.BigEndian.PutUint32(hdr[:], uint32(len(png)))
if _, err := c.in.Write(append(hdr[:], png...)); err != nil {
done <- reply{err: err}
return
}
line, err := c.out.ReadBytes('\n')
if err != nil {
done <- reply{err: err}
return
}
var r Result
err = json.Unmarshal(line, &r)
done <- reply{r, err}
}()
select {
case <-ctx.Done():
c.stop() // the helper is mid-reply; a fresh one is simpler than resyncing
return Result{}, ctx.Err()
case rep := <-done:
if rep.err != nil {
c.stop()
return Result{}, fmt.Errorf("ocr helper: %w", rep.err)
}
c.arm()
if rep.r.Error != "" {
return rep.r, errors.New(rep.r.Error)
}
return rep.r, nil
}
}
// Close stops the helper.
func (c *Client) Close() {
c.mu.Lock()
defer c.mu.Unlock()
c.stop()
}
func (c *Client) start() error {
cmd := exec.Command(c.Path, c.Args...)
cmd.Stderr = os.Stderr
in, err := cmd.StdinPipe()
if err != nil {
return err
}
out, err := cmd.StdoutPipe()
if err != nil {
return err
}
if err := cmd.Start(); err != nil {
return fmt.Errorf("ocr helper: %w", err)
}
c.cmd, c.in, c.out = cmd, in, bufio.NewReader(out)
return nil
}
func (c *Client) arm() {
if c.Idle <= 0 {
return
}
cmd := c.cmd
c.timer = time.AfterFunc(c.Idle, func() {
c.mu.Lock()
defer c.mu.Unlock()
if c.cmd == cmd {
c.stop()
}
})
}
// stop closes stdin, which ends the helper, and reaps it.
func (c *Client) stop() {
if c.cmd == nil {
return
}
if c.timer != nil {
c.timer.Stop()
}
c.in.Close()
done := make(chan struct{})
go func() { c.cmd.Wait(); close(done) }()
select {
case <-done:
case <-time.After(2 * time.Second):
c.cmd.Process.Kill()
<-done
}
c.cmd, c.in, c.out, c.timer = nil, nil, nil, nil
}
// ReadRequest reads one request on the helper side; io.EOF means stdin
// closed and the helper should exit.
func ReadRequest(r io.Reader) ([]byte, error) {
var hdr [4]byte
if _, err := io.ReadFull(r, hdr[:]); err != nil {
return nil, err
}
n := binary.BigEndian.Uint32(hdr[:])
if n > 64<<20 {
return nil, fmt.Errorf("request of %d bytes", n)
}
buf := make([]byte, n)
_, err := io.ReadFull(r, buf)
return buf, err
}
// WriteResult writes one reply on the helper side.
func WriteResult(w io.Writer, r Result) error {
b, err := json.Marshal(r)
if err != nil {
return err
}
_, err = w.Write(append(b, '\n'))
return err
}