ppocr-go.go 1.2 KB
Newer Older
L
LKKlein 已提交
1 2 3 4 5
package main

import (
	"flag"
	"log"
L
LKKlein 已提交
6
	"github.com/PaddlePaddle/PaddleOCR/thirdparty/paddleocr-go/ocr"
L
LKKlein 已提交
7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30
)

var (
	confFile     string
	image        string
	imageDir     string
	useServering bool
	port         string
)

func init() {
	flag.StringVar(&confFile, "config", "config/conf.yaml", "config from ocr system. If not given, will use default config.")
	flag.StringVar(&image, "image", "", "image to predict. if not given, will use image_dir")
	flag.StringVar(&imageDir, "image_dir", "", "imgs in dir to be predicted. if not given, will check servering")
	flag.BoolVar(&useServering, "use_servering", false, "whether to use ocr server. [default: false]")
	flag.StringVar(&port, "port", "18600", "which port to serve ocr server. [default: 18600].")
}

func main() {
	flag.Parse()
	sys := ocr.NewOCRSystem(confFile, nil)

	if image != "" {
		img := ocr.ReadImage(image)
L
LKKlein 已提交
31
		defer img.Close()
L
LKKlein 已提交
32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52
		results := sys.PredictOneImage(img)
		for _, res := range results {
			log.Println(res)
		}
		return
	}

	if imageDir != "" {
		results := sys.PredictDirImages(imageDir)
		for k, vs := range results {
			log.Printf("======== image: %v =======\n", k)
			for _, res := range vs {
				log.Println(res)
			}
		}
	}

	if useServering {
		sys.StartServer(port)
	}
}