diff --git a/args.go b/args.go new file mode 100644 index 0000000..887f859 --- /dev/null +++ b/args.go @@ -0,0 +1,64 @@ +package main + +import ( + "github.com/juju/gnuflag" + + "errors" +) + +var ( + ArgInput string + ArgInputIsSet bool = false + ArgOutput string + ArgOutputIsSet bool = false + ArgThreshold string + ArgThresholdIsSet bool = false + + ArgJobs int = 0 + ArgDefaultNoiseLevel int = 0 +) + +func Use(vals ...interface{}) { + for _, val := range vals { + _ = val + } +} + +func init() { + gnuflag.IntVar(&ArgJobs, "jobs", 0, "") + gnuflag.IntVar(&ArgJobs, "j", 0, "") + gnuflag.IntVar(&ArgDefaultNoiseLevel, "noise", -1, "") + gnuflag.IntVar(&ArgDefaultNoiseLevel, "n", -1, "") +} + +func parseArgs() error { + gnuflag.Parse(false) + + for i := 0; i < 3 && i < gnuflag.NArg(); i++ { + switch i { + case 0: + ArgInput = gnuflag.Arg(0) + ArgInputIsSet = true + case 1: + ArgOutput = gnuflag.Arg(1) + ArgOutputIsSet = true + case 2: + ArgThreshold = gnuflag.Arg(2) + ArgThresholdIsSet = true + } + } + + if !ArgInputIsSet { + return errors.New("Input directory not specified") + } + + if !ArgOutputIsSet { + return errors.New("Output file not specified") + } + + if !ArgThresholdIsSet { + return errors.New("Threshold not specified") + } + + return nil +} diff --git a/main.go b/main.go new file mode 100644 index 0000000..34a59ee --- /dev/null +++ b/main.go @@ -0,0 +1,142 @@ +package main + +import ( + "bufio" + "image" + "io/ioutil" + "os" + "sort" + "strconv" + "strings" + + _ "image/png" + + "github.com/dsoprea/go-perceptualhash" + + "range-generator/output" +) + +const ( + EXIT_SUCCESS int = 0 + EXIT_ARG_ERR int = 1 + EXIT_FILE_READ_ERR int = 2 + EXIT_DIR_READ_ERR int = 3 + EXIT_FILE_WRITE_ERR int = 4 +) + +var ( + hashes map[string]string + names []string +) + +func main() { + output.SetEndOfLineNative() + + if err := parseArgs(); err != nil { + output.Errorln(err) + os.Exit(EXIT_ARG_ERR) + } + + Threshold, err := strconv.Atoi(ArgThreshold) + if err != nil { + output.Errorln("Can't convert threshold to integer") + os.Exit(EXIT_ARG_ERR) + } + + files, err := ioutil.ReadDir(ArgInput) + if err != nil { + output.Errorln("Read dir error") + os.Exit(EXIT_DIR_READ_ERR) + } + + hashes = make(map[string]string) + for _, file := range files { + if !file.IsDir() && strings.HasSuffix(file.Name(), ".png") { + names = append(names, file.Name()) + hashes[file.Name()], err = calchash(ArgInput + file.Name()) + + if err != nil { + output.Errorln(err) + } + } + } + sort.Strings(names) + + var ranges strings.Builder + var prevHash string = "" + var dist int = 0 + var startName = "" + var rangeNoise string = "" + + if ArgDefaultNoiseLevel >= 0 { + rangeNoise = "\t" + string(ArgDefaultNoiseLevel) + } + + for i := 0; i < len(names); i++ { + name := names[i] + if startName == "" { + startName = name + } + dist = hammingDistance(prevHash, hashes[name]) + if dist >= Threshold { + ranges.WriteString(startName + "\t" + names[i-1] + rangeNoise + output.EOL()) + startName = name + } + prevHash = hashes[name] + } + ranges.WriteString(startName + "\t" + names[len(names)-1] + rangeNoise + output.EOL()) + + mode := os.FileMode(int(0644)) + targetFile, err := os.OpenFile(ArgOutput, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, mode) + defer targetFile.Close() + if err == nil { + datawriter := bufio.NewWriter(targetFile) + _, err = datawriter.WriteString(ranges.String()) + if err == nil { + err = datawriter.Flush() + } + } + + if err != nil { + output.Errorln(err) + os.Exit(EXIT_FILE_WRITE_ERR) + } + + os.Exit(EXIT_SUCCESS) +} + +func calchash(filepath string) (string, error) { + file, err := os.Open(filepath) + if err != nil { + return "", err + } + defer file.Close() + + image, _, err := image.Decode(file) + if err != nil { + return "", err + } + + return blockhash.NewBlockhash(image, 16).Hexdigest(), nil +} + +func hammingDistance(prev, cur string) int { + var dist int = 0 + var p, c int64 + + if prev == "" || cur == "" { + return 0 + } + + for i := 0; i < len(cur); i++ { + p, _ = strconv.ParseInt(string([]rune(prev)[i]), 16, 64) + c, _ = strconv.ParseInt(string([]rune(cur)[i]), 16, 64) + if p > c { + dist += int(p - c) + } else { + dist += int(c - p) + } + } + + return dist +} diff --git a/output/output.go b/output/output.go new file mode 100644 index 0000000..7b16205 --- /dev/null +++ b/output/output.go @@ -0,0 +1,68 @@ +package output + +import ( + "fmt" + "io/ioutil" + "log" + "os" + "runtime" +) + +var ( + endOfLine string = "\n" + devNull *log.Logger = log.New(ioutil.Discard, "", 0) + stdout *log.Logger = log.New(os.Stdout, "", 0) + stderr *log.Logger = log.New(os.Stderr, "", 0) +) + +func SetQuiet(enabled bool) { + if enabled { + stdout = devNull + stderr = devNull + } +} + +func SetEndOfLineNative() { + switch os := runtime.GOOS; os { + case "windows": + setEndOfLineWindows() + default: + setEndOfLineUnix() + } +} + +func EOL() string { + return endOfLine +} + +func setEndOfLineUnix() { + endOfLine = "\n" +} + +func setEndOfLineWindows() { + endOfLine = "\r\n" +} + +func Print(v ...interface{}) { + stdout.Print(v...) +} + +func Printf(format string, v ...interface{}) { + stdout.Printf(format, v...) +} + +func Println(v ...interface{}) { + stdout.Print(fmt.Sprint(v...) + endOfLine) +} + +func Error(v ...interface{}) { + stderr.Print(v...) +} + +func Errorf(format string, v ...interface{}) { + stderr.Printf(format, v...) +} + +func Errorln(v ...interface{}) { + stderr.Print(fmt.Sprint(v...) + endOfLine) +}