package main
import (
"bytes"
"context"
"errors"
"flag"
"fmt"
"image"
_ "image/jpeg"
_ "image/png"
"io"
"log"
"os"
"path/filepath"
"chirami/internal/ocr"
"chirami/internal/ocr/fetch"
"chirami/internal/ocr/onnx"
"chirami/internal/paths"
)
// defaultModelDir is where `chirami setup -ocr` puts the models.
func defaultModelDir() string {
if env := os.Getenv("CHIRAMI_OCR_MODELS"); env != "" {
return env
}
return filepath.Join(paths.DataDir(), "ocr")
}
func main() {
log.SetFlags(0)
log.SetPrefix("chirami-ocr: ")
modelsFlag := flag.String("models", defaultModelDir(), "path to directory containing OCR models")
ortLibFlag := flag.String("ort-lib", "", "explicit path to libonnxruntime shared library")
downloadFlag := flag.Bool("download", false, "ensure all models are downloaded to -models dir and exit")
prepareFlag := flag.Bool("prepare", false, "load every model once (writing the ORT-format copies) and exit")
flag.Parse()
modelsDir := *modelsFlag
if *downloadFlag {
log.Printf("checking and downloading models into %s...", modelsDir)
if err := fetch.EnsureModels(context.Background(), modelsDir); err != nil {
log.Fatalf("download models: %v", err)
}
log.Printf("all models verified in %s", modelsDir)
return
}
// setup verified the hashes; here only a cheap check that nothing is
// missing, since hashing 480 MB would slow every start
if missing := fetch.Absent(modelsDir); len(missing) > 0 {
log.Fatalf("%d model file(s) missing in %s; run `chirami setup -ocr`", len(missing), modelsDir)
}
if err := onnx.InitRuntime(*ortLibFlag); err != nil {
log.Fatalf("init onnxruntime: %v", err)
}
pipe, err := onnx.NewPipeline(modelsDir)
if err != nil {
log.Fatalf("initialize pipeline: %v", err)
}
defer pipe.Close()
if *prepareFlag {
return
}
// Main request loop: read 4-byte big-endian length + PNG on stdin, reply JSON on stdout
for {
req, err := ocr.ReadRequest(os.Stdin)
if err != nil {
if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
return
}
log.Fatalf("read request: %v", err)
}
img, _, err := image.Decode(bytes.NewReader(req))
if err != nil {
_ = ocr.WriteResult(os.Stdout, ocr.Result{
Engine: "onnx",
Error: fmt.Sprintf("decode image: %v", err),
})
continue
}
res, err := pipe.Recognize(img)
if err != nil {
res = ocr.Result{
Engine: "onnx",
Error: err.Error(),
}
}
if err := ocr.WriteResult(os.Stdout, res); err != nil {
log.Fatalf("write result: %v", err)
}
}
}