-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
121 lines (114 loc) · 2.83 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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
package main
import (
"errors"
"flag"
"fmt"
"image"
"image/png"
"io"
"os"
"path/filepath"
"strings"
"giautm.dev/captcha/binimg"
"giautm.dev/captcha/fisheye"
"giautm.dev/captcha/labeled"
"github.com/manifoldco/promptui"
)
var (
dir = flag.String("dir", "./data/incorrect", "Input directory with labeled images")
doneDir = flag.String("done", "./data/done", "Input directory with labeled images")
outputDir = flag.String("output", "./data/labeled", "Output directory with labeled images")
)
const (
captchaLen = 5
binWidth = 10
testRowIndex = 42
)
func main() {
flag.Parse()
if *dir == "" || *outputDir == "" {
flag.Usage()
return
}
prompt := promptui.Prompt{
Label: "What's the correct captcha?",
Validate: func(input string) error {
if len(input) != captchaLen {
return errors.New("Cần nhập đủ 5 ký tự")
}
return nil
},
}
err := filepath.Walk(*dir, func(path string, info os.FileInfo, err error) error {
if err != nil || info.IsDir() {
return err
}
return relabelFile(path, info, prompt.Run)
})
if err != nil {
fmt.Printf("error walking the path %q: %v\n", *dir, err)
return
}
}
func relabelFile(path string, info os.FileInfo, askCaptcha func() (string, error)) error {
fmt.Printf("-> Process: %s\n", path)
file, err := os.OpenFile(path, os.O_RDONLY, 0)
if err != nil {
return err
}
defer file.Close()
preview, err := os.Create("./data/captcha.png")
if err != nil {
return err
}
defer preview.Close()
if _, err = io.Copy(preview, file); err != nil {
return err
}
result, err := askCaptcha()
if err != nil {
fmt.Printf("Prompt failed %v\n", err)
return err
}
name := fmt.Sprintf("%s-%d.png", result, info.ModTime().Unix())
if err = os.Rename(path, filepath.Join(*doneDir, name)); err != nil {
return err
}
captcha := []rune(labeled.CaptchaFromName(info.Name()))
wrongs := map[int]string{}
wrongChars := []string{}
for idx, c := range result {
if captcha[idx] != c {
wrongs[idx] = (string)(c)
wrongChars = append(wrongChars, wrongs[idx])
}
}
if c := len(wrongs); c > 0 {
fmt.Printf("Found %d wrong character: %s\n", c, strings.Join(wrongChars, ","))
// Reset to head for decode img
file.Seek(0, 0)
img, _, err := image.Decode(file)
if err != nil {
return err
}
img, distance := fisheye.FindDistance(img, testRowIndex)
if distance < 0 {
return errors.New("preprocess: can not detect distance")
}
for pos, label := range wrongs {
if err = genBinFile(img, pos, label, name); err != nil {
return err
}
}
}
return nil
}
func genBinFile(img image.Image, pos int, label, name string) error {
file, err := labeled.CreateLabeledFile(*outputDir, label, name)
if err != nil {
return err
}
defer file.Close()
bImg := binimg.MarkPosition(binimg.ExpandLeft(img, binWidth*captchaLen), binWidth, pos)
return png.Encode(file, bImg)
}