Add configuration struct for command-line flags and refactor main logic to use it
This commit is contained in:
54
main.go
54
main.go
@@ -15,6 +15,13 @@ import (
|
|||||||
"github.com/corona10/goimagehash"
|
"github.com/corona10/goimagehash"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
type Config struct {
|
||||||
|
Threshold int
|
||||||
|
Workers int
|
||||||
|
NotifyInterval int
|
||||||
|
NoGroupByFolder bool
|
||||||
|
}
|
||||||
|
|
||||||
type HashResult struct {
|
type HashResult struct {
|
||||||
file string
|
file string
|
||||||
hash *goimagehash.ExtImageHash
|
hash *goimagehash.ExtImageHash
|
||||||
@@ -26,14 +33,20 @@ type GroupResult struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
thresh := flag.Int("thresh", 10, "Threshold for distance")
|
cfg := Config{
|
||||||
workers := flag.Int("workers", 100, "Number of workers")
|
Threshold: *flag.Int("thresh", 10, "Threshold for distance"),
|
||||||
notifyInterval := flag.Int("notify", 1000, "Notify interval")
|
Workers: *flag.Int("workers", 100, "Number of workers"),
|
||||||
|
NotifyInterval: *flag.Int("notify", 1000, "Notify interval"),
|
||||||
|
NoGroupByFolder: *flag.Bool("nogroup", false, "Don't group by folder"),
|
||||||
|
}
|
||||||
flag.Parse()
|
flag.Parse()
|
||||||
logger.InitFlag()
|
logger.InitFlag()
|
||||||
logger.Info("Starting")
|
logger.Info("Starting")
|
||||||
logger.Info("Threshold: %v", *thresh)
|
logger.Info("Threshold: %v", cfg.Threshold)
|
||||||
logger.Info("Patterns: %d", len(flag.Args()))
|
logger.Info("Patterns: %d", len(flag.Args()))
|
||||||
|
logger.Info("Workers: %v", cfg.Workers)
|
||||||
|
logger.Info("Notify interval: %v", cfg.NotifyInterval)
|
||||||
|
logger.Info("Group by folder: %v", !cfg.NoGroupByFolder)
|
||||||
|
|
||||||
// Collect files
|
// Collect files
|
||||||
files := make([]string, 0)
|
files := make([]string, 0)
|
||||||
@@ -54,15 +67,38 @@ func main() {
|
|||||||
}
|
}
|
||||||
logger.Info("Patterns expanded to %d files", len(files))
|
logger.Info("Patterns expanded to %d files", len(files))
|
||||||
|
|
||||||
|
if !cfg.NoGroupByFolder {
|
||||||
|
// Group files by folder
|
||||||
|
folderMap := make(map[string][]string)
|
||||||
|
for _, file := range files {
|
||||||
|
dir := filepath.Dir(file)
|
||||||
|
folderMap[dir] = append(folderMap[dir], file)
|
||||||
|
}
|
||||||
|
logger.Info("Processing %d folders", len(folderMap))
|
||||||
|
|
||||||
|
// Process each folder separately
|
||||||
|
for dir, group := range folderMap {
|
||||||
|
logger.Info("Processing folder: %s (%d files)", dir, len(group))
|
||||||
|
processFiles(group, cfg)
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// Process all files together
|
||||||
|
processFiles(files, cfg)
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.Info("Done")
|
||||||
|
}
|
||||||
|
|
||||||
|
func processFiles(files []string, cfg Config) {
|
||||||
// Phase 1: Hashing
|
// Phase 1: Hashing
|
||||||
hashChan := make(chan string, len(files))
|
hashChan := make(chan string, len(files))
|
||||||
hashResults := make(chan HashResult, len(files))
|
hashResults := make(chan HashResult, len(files))
|
||||||
var hashWg sync.WaitGroup
|
var hashWg sync.WaitGroup
|
||||||
|
|
||||||
// Start hash workers
|
// Start hash workers
|
||||||
for i := 0; i < *workers; i++ {
|
for i := 0; i < cfg.Workers; i++ {
|
||||||
hashWg.Add(1)
|
hashWg.Add(1)
|
||||||
go hashWorker(hashChan, hashResults, &hashWg, *notifyInterval, len(files))
|
go hashWorker(hashChan, hashResults, &hashWg, cfg.NotifyInterval, len(files))
|
||||||
}
|
}
|
||||||
|
|
||||||
// Send files to hash workers
|
// Send files to hash workers
|
||||||
@@ -88,9 +124,9 @@ func main() {
|
|||||||
var groupWg sync.WaitGroup
|
var groupWg sync.WaitGroup
|
||||||
|
|
||||||
// Start group workers
|
// Start group workers
|
||||||
for i := 0; i < *workers; i++ {
|
for i := 0; i < cfg.Workers; i++ {
|
||||||
groupWg.Add(1)
|
groupWg.Add(1)
|
||||||
go groupWorker(groupChan, groupResults, &groupWg, *thresh, *notifyInterval, len(allHashes), allHashes)
|
go groupWorker(groupChan, groupResults, &groupWg, cfg.Threshold, cfg.NotifyInterval, len(allHashes), allHashes)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Send hash results to group workers
|
// Send hash results to group workers
|
||||||
@@ -138,8 +174,6 @@ func main() {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
logger.Info("Done")
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func hashWorker(in <-chan string, out chan<- HashResult, wg *sync.WaitGroup, notifyInterval int, totalFiles int) {
|
func hashWorker(in <-chan string, out chan<- HashResult, wg *sync.WaitGroup, notifyInterval int, totalFiles int) {
|
||||||
|
Reference in New Issue
Block a user