/
main.go
97 lines (81 loc) · 1.92 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
package main
import (
"flag"
"fmt"
"image"
_ "image/jpeg"
_ "image/png"
"io/ioutil"
"log"
"os"
"path"
"runtime/pprof"
"strings"
"github.com/mrfuxi/sudoku"
"github.com/mrfuxi/sudoku/digits"
"github.com/mrfuxi/sudoku/nngrid"
)
const (
exampleDir = "examples"
saveLocation = "examples_out"
)
func getExampleImage(name string) (image.Image, error) {
filePath := path.Join(exampleDir, name)
reader, err := os.Open(filePath)
if err != nil {
log.Fatalln(err)
}
img, _, err := image.Decode(reader)
return img, err
}
func findSudoku(filename string, debug bool) (sudoku.Sudoku, error) {
img, err := getExampleImage(filename)
if err != nil {
log.Fatalln(err)
}
return sudoku.NewSudoku(img)
}
func main() {
os.RemoveAll(saveLocation)
os.MkdirAll(saveLocation, os.ModePerm)
var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
var debug = flag.Bool("debug", false, "prepare debug images")
var file = flag.String("file", "", "file to process")
var nnFile = flag.String("nn", "", "neural network")
var gnnFile = flag.String("gnn", "", "grid neural network")
flag.Parse()
if *cpuprofile != "" {
f, err := os.Create(*cpuprofile)
if err != nil {
log.Fatal(err)
}
pprof.StartCPUProfile(f)
defer pprof.StopCPUProfile()
}
var s sudoku.Sudoku
var err error
if *nnFile == "" {
fmt.Println("Neural network file not provided. Use -nn FileName")
os.Exit(1)
}
digits.LoadNetwork(*nnFile)
nngrid.LoadNetwork(*gnnFile)
if *file != "" {
s, err = findSudoku(*file, *debug)
} else {
fileInfos, err := ioutil.ReadDir(exampleDir)
if err != nil {
log.Fatal(err)
}
for _, fileInfo := range fileInfos {
if strings.HasSuffix(fileInfo.Name(), ".png") || strings.HasSuffix(fileInfo.Name(), ".jpg") {
s, err = findSudoku(fileInfo.Name(), *debug)
}
}
}
if err == nil && s != nil {
fmt.Println("Sudoku has been found!")
} else {
fmt.Println("Could not find sudoku")
}
}