change project structure
This commit is contained in:
96
cmd/range-gen/args.go
Normal file
96
cmd/range-gen/args.go
Normal file
@ -0,0 +1,96 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"github.com/juju/gnuflag"
|
||||
|
||||
"range-gen/internal/output"
|
||||
|
||||
"errors"
|
||||
"os"
|
||||
)
|
||||
|
||||
var (
|
||||
ArgInput string
|
||||
ArgInputIsSet bool = false
|
||||
ArgOutput string
|
||||
ArgOutputIsSet bool = false
|
||||
ArgThreshold string
|
||||
ArgThresholdIsSet bool = false
|
||||
|
||||
ArgJobs int = 0
|
||||
ArgDefaultNoiseLevel int = 0
|
||||
|
||||
ArgVersion bool = false
|
||||
ArgHelp bool = false
|
||||
)
|
||||
|
||||
func printHelp() {
|
||||
output.Println("Сreates a list of scene ranges based on a set of frames from the video")
|
||||
output.Println("")
|
||||
output.Println("Usage: range-gen [option]... <input_dir> <output_file> <threshold>")
|
||||
output.Println("input_dir Directory with png images")
|
||||
output.Println("output_file Range list file")
|
||||
output.Println("threshold Image similarity threshold (0-1024)")
|
||||
output.Println("")
|
||||
output.Println("Options:")
|
||||
output.Println(" -j, --jobs N Allow N jobs at once")
|
||||
output.Println(" -n, --noise Default noise level for each range")
|
||||
output.Println(" -h, --help Show this page")
|
||||
output.Println(" -v, --version Show version")
|
||||
}
|
||||
|
||||
func printVersion() {
|
||||
output.Println("range-gen ", Version)
|
||||
}
|
||||
|
||||
func init() {
|
||||
gnuflag.IntVar(&ArgJobs, "jobs", 0, "")
|
||||
gnuflag.IntVar(&ArgJobs, "j", 0, "")
|
||||
gnuflag.IntVar(&ArgDefaultNoiseLevel, "noise", -1, "")
|
||||
gnuflag.IntVar(&ArgDefaultNoiseLevel, "n", -1, "")
|
||||
gnuflag.BoolVar(&ArgVersion, "version", false, "")
|
||||
gnuflag.BoolVar(&ArgVersion, "v", false, "")
|
||||
gnuflag.BoolVar(&ArgHelp, "help", false, "")
|
||||
gnuflag.BoolVar(&ArgHelp, "h", false, "")
|
||||
}
|
||||
|
||||
func parseArgs() error {
|
||||
gnuflag.Parse(false)
|
||||
|
||||
switch {
|
||||
case ArgHelp:
|
||||
printHelp()
|
||||
os.Exit(EXIT_SUCCESS)
|
||||
case ArgVersion:
|
||||
printVersion()
|
||||
os.Exit(EXIT_SUCCESS)
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
197
cmd/range-gen/main.go
Normal file
197
cmd/range-gen/main.go
Normal file
@ -0,0 +1,197 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"image"
|
||||
"io/ioutil"
|
||||
"os"
|
||||
"os/signal"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
|
||||
_ "image/png"
|
||||
|
||||
"github.com/dsoprea/go-perceptualhash"
|
||||
|
||||
"range-gen/internal/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 (
|
||||
Version string = "dev"
|
||||
)
|
||||
|
||||
func main() {
|
||||
closeHandler() // Ctrl+C
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
allFiles, err := ioutil.ReadDir(ArgInput)
|
||||
if err != nil {
|
||||
output.Errorln("Read dir error")
|
||||
os.Exit(EXIT_DIR_READ_ERR)
|
||||
}
|
||||
|
||||
if ArgJobs == 0 {
|
||||
ArgJobs = runtime.NumCPU()
|
||||
}
|
||||
runtime.GOMAXPROCS(ArgJobs)
|
||||
|
||||
pngFiles := pngList(allFiles)
|
||||
hashes := calcHashes(pngFiles)
|
||||
ranges := calcRanges(hashes, Threshold)
|
||||
writeRanges(ranges)
|
||||
|
||||
os.Exit(EXIT_SUCCESS)
|
||||
}
|
||||
|
||||
func closeHandler() {
|
||||
interrupt := make(chan os.Signal, 2)
|
||||
signal.Notify(interrupt, os.Interrupt, syscall.SIGTERM)
|
||||
go func() {
|
||||
<-interrupt
|
||||
output.Println("Closed by SIGINT")
|
||||
os.Exit(EXIT_SUCCESS)
|
||||
}()
|
||||
}
|
||||
|
||||
func pngList(files []os.FileInfo) []os.FileInfo {
|
||||
var pngFiles []os.FileInfo
|
||||
|
||||
for _, file := range files {
|
||||
if !file.IsDir() && strings.HasSuffix(strings.ToLower(file.Name()), ".png") {
|
||||
pngFiles = append(pngFiles, file)
|
||||
}
|
||||
}
|
||||
|
||||
return pngFiles
|
||||
}
|
||||
|
||||
func calcHashes(files []os.FileInfo) map[string]string {
|
||||
var hashes map[string]string
|
||||
var mutex sync.Mutex
|
||||
|
||||
wg := new(sync.WaitGroup)
|
||||
wg.Add(len(files))
|
||||
|
||||
hashes = make(map[string]string)
|
||||
|
||||
for _, file := range files {
|
||||
go calcHash(ArgInput, file, &hashes, &mutex, wg)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
return hashes
|
||||
}
|
||||
|
||||
func calcHash(path string, fileinfo os.FileInfo, hashes *map[string]string, mutex *sync.Mutex, wg *sync.WaitGroup) {
|
||||
file, err := os.Open(filepath.Join(path, fileinfo.Name()))
|
||||
if err != nil {
|
||||
output.Errorln(err)
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
image, _, err := image.Decode(file)
|
||||
if err != nil {
|
||||
output.Errorln(err)
|
||||
}
|
||||
hash := blockhash.NewBlockhash(image, 16).Hexdigest()
|
||||
mutex.Lock()
|
||||
(*hashes)[fileinfo.Name()] = hash
|
||||
mutex.Unlock()
|
||||
wg.Done()
|
||||
}
|
||||
|
||||
func calcRanges(hashes map[string]string, Threshold int) string {
|
||||
var ranges strings.Builder
|
||||
var prevHash string = ""
|
||||
var dist int = 0
|
||||
var startName = ""
|
||||
var rangeNoise string = ""
|
||||
|
||||
if ArgDefaultNoiseLevel >= 0 {
|
||||
rangeNoise = "\t" + string(ArgDefaultNoiseLevel)
|
||||
}
|
||||
|
||||
names := []string{}
|
||||
for key := range hashes {
|
||||
names = append(names, key)
|
||||
}
|
||||
sort.Strings(names)
|
||||
|
||||
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())
|
||||
|
||||
return ranges.String()
|
||||
}
|
||||
|
||||
func writeRanges(ranges string) {
|
||||
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)
|
||||
if err == nil {
|
||||
err = datawriter.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
output.Errorln(err)
|
||||
os.Exit(EXIT_FILE_WRITE_ERR)
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user