geo/server/internal/shared/digitocr/cmd/train/main.go

// 训练工具：读取 samples/0.png ~ samples/9.png，生成 templates.go。
//
// 用法：
//
//	cd server/internal/shared/digitocr
//	# 准备 samples/0.png ... 9.png，每张是同字体的单个数字
//	# 也可以是含多位的整图，命名 NNNN.png（如 4930.png），脚本会按位切分
//	go run ./cmd/train
package main

import (
	"fmt"
	"image"
	_ "image/png"
	"os"
	"path/filepath"
	"strings"

	"github.com/geo-platform/tenant-api/internal/shared/digitocr"
)

func main() {
	groups, err := collectSamples("samples")
	if err != nil {
		exit(err)
	}
	missing := missingDigits(groups)
	if len(missing) > 0 {
		exit(fmt.Errorf("missing samples for digits: %v", missing))
	}

	templates := [10]digitocr.Bitmap{}
	for d := 0; d < 10; d++ {
		templates[d] = voteTemplate(groups[d])
	}

	var sb strings.Builder
	sb.WriteString("package digitocr\n\n")
	sb.WriteString("// Code generated by cmd/train. DO NOT EDIT.\n\n")
	sb.WriteString("var Templates = [10]Bitmap{\n")
	for d := 0; d < 10; d++ {
		sb.WriteString("\t{")
		for i, v := range templates[d] {
			if i > 0 {
				sb.WriteByte(',')
			}
			if v == 0 {
				sb.WriteByte('0')
			} else {
				sb.WriteByte('1')
			}
		}
		fmt.Fprintf(&sb, "}, // d=%d, n=%d\n", d, len(groups[d]))
	}
	sb.WriteString("}\n")

	if err := os.WriteFile("templates.go", []byte(sb.String()), 0o644); err != nil {
		exit(err)
	}
	fmt.Printf("wrote templates.go (samples per digit: ")
	for d := 0; d < 10; d++ {
		fmt.Printf("%d=%d ", d, len(groups[d]))
	}
	fmt.Println(")")
}

// voteTemplate 对一组同数字样本做按位多数投票。
func voteTemplate(bms []digitocr.Bitmap) digitocr.Bitmap {
	var out digitocr.Bitmap
	if len(bms) == 0 {
		return out
	}
	counts := make([]int, digitocr.GridW*digitocr.GridH)
	for _, bm := range bms {
		for i, v := range bm {
			if v == 1 {
				counts[i]++
			}
		}
	}
	threshold := len(bms) / 2
	for i, c := range counts {
		if c > threshold {
			out[i] = 1
		}
	}
	return out
}

// collectSamples 从 dir 下读取样本。支持两种命名：
//   - 0.png ~ 9.png：单字符样本，整张图就是一个数字
//   - NNNN.png（如 4930.png）：多位样本，按位切分，按文件名映射到对应数字
//
// 返回 map[digit][]Bitmap，每个数字累积所有样本以便投票。
func collectSamples(dir string) (map[int][]digitocr.Bitmap, error) {
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, err
	}
	out := map[int][]digitocr.Bitmap{}
	for _, e := range entries {
		if e.IsDir() {
			continue
		}
		name := strings.TrimSuffix(e.Name(), filepath.Ext(e.Name()))
		if !allDigits(name) {
			continue
		}
		bms, err := extractBitmaps(filepath.Join(dir, e.Name()), len(name))
		if err != nil {
			return nil, fmt.Errorf("%s: %w", e.Name(), err)
		}
		for i, ch := range name {
			d := int(ch - '0')
			out[d] = append(out[d], bms[i])
		}
	}
	return out, nil
}

func extractBitmaps(path string, digits int) ([]digitocr.Bitmap, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	img, _, err := image.Decode(f)
	if err != nil {
		return nil, err
	}
	bin, bbox := digitocr.Binarize(img, digitocr.IsOrange)
	if bbox.Empty() {
		return nil, fmt.Errorf("no foreground pixels")
	}
	cells := digitocr.SegmentDigits(bin, bbox, digits)
	if len(cells) != digits {
		return nil, fmt.Errorf("segmented %d cells, expected %d", len(cells), digits)
	}
	bms := make([]digitocr.Bitmap, digits)
	for i, cell := range cells {
		bms[i] = digitocr.Normalize(bin, cell)
	}
	return bms, nil
}

func allDigits(s string) bool {
	if s == "" {
		return false
	}
	for _, r := range s {
		if r < '0' || r > '9' {
			return false
		}
	}
	return true
}

func missingDigits(m map[int][]digitocr.Bitmap) []int {
	var miss []int
	for d := 0; d < 10; d++ {
		if len(m[d]) == 0 {
			miss = append(miss, d)
		}
	}
	return miss
}

func exit(err error) {
	fmt.Fprintln(os.Stderr, "train:", err)
	os.Exit(1)
}