summaryrefslogtreecommitdiff
path: root/pare-gt/main.go
blob: d07e141d2d71aa4a2b064c0c0c8ad94bc4cc8dbe (plain)
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
122
123
124
125
126
127
128
129
130
131
132
133
package main

import (
	"flag"
	"fmt"
	"log"
	"math/rand"
	"os"
	"path"
	"path/filepath"
	"strings"
)

const usage = `Usage: pare-gt [-n num] gtdir movedir

Moves some of the ground truth from gt-dir into movedir,
ensuring that the same proportions of each ground truth
source are represented in the moved section. Proportion of
ground truth source is calculated by taking the prefix of
the filename up to the first '-' character.
`

// Prefixes is a map of the prefix string to a list of filenames
type Prefixes = map[string][]string

// walker adds any .txt path to prefixes map, under the appropriate
// prefix (blank if no '-' separator was found)
func walker(prefixes *Prefixes) filepath.WalkFunc{
	return func(fpath string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}
		ext := path.Ext(fpath)
		if ext != ".txt" {
			return nil
		}
		base := path.Base(fpath)
		idx := strings.Index(base, "-")
		var prefix string
		if idx > -1 {
			prefix = base[0:idx]
		}
		noext := strings.TrimSuffix(fpath, ext)
		(*prefixes)[prefix] = append((*prefixes)[prefix], noext)
		return nil
	}
}

// inStrSlice checks whether a given string is part of a slice of
// strings
func inStrSlice(sl []string, s string) bool {
	for _, v := range sl {
		if s == v {
			return true
		}
	}
	return false
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(flag.CommandLine.Output(), usage)
		flag.PrintDefaults()
	}
	numtopare := flag.Int("n", 10, "Percentage of the ground truth to pare away.")
	flag.Parse()
	if flag.NArg() != 2 {
		flag.Usage()
		os.Exit(1)
	}

	for _, d := range flag.Args() {
		info, err := os.Stat(d)
		if err != nil || !info.IsDir() {
			log.Fatalln("Error accessing directory", flag.Arg(0), err)
		}
	}

	var prefixes Prefixes
	prefixes = make(Prefixes)
	err := filepath.Walk(flag.Arg(0), walker(&prefixes))
	if err != nil {
		log.Fatalln("Failed to walk", flag.Arg(0), err)
	}

	var total, sample int
	for _, v := range prefixes {
		total += len(v)
	//	fmt.Printf("\n%s:\n%s\n", i, v)
	}

	sample = total / *numtopare

	// filestomove contains the names of files to move minus file extension
	var filestomove []string

	// select random samples for each prefix, proportional to
	// the amount of that prefix there are in the whole set
	for _, prefix := range prefixes {
		len := len(prefix)
		if len == 1 {
			continue
		}
		numtoget := int(float64(sample) / float64(total) * float64(len))
		if numtoget < 1 {
			numtoget = 1
		}
		for i:=0; i<numtoget; i++ {
			var selected string
			selected = prefix[rand.Int()%len]
			// pick a different random selection if the first one is
			// already in the filestomove slice
			for inStrSlice(filestomove, selected) {
				selected = prefix[rand.Int()%len]
			}
			filestomove = append(filestomove, selected)
		}
	}

	for _, f := range filestomove {
		fmt.Println("Moving ground truth", f)
		b := path.Base(f)
		for _, ext := range []string{".txt", ".png"} {
			err = os.Rename(f + ext, path.Join(flag.Arg(1), b + ext))
			if err != nil {
				log.Fatalln("Error moving file", f + ext, err)
			}
		}
	}
}