Add pagegraph tool

2 files changed, 103 insertions, 7 deletions

diff --git a/cmd/pagegraph/main.go b/cmd/pagegraph/main.go
new file mode 100644
index 0000000..fb2648f
--- /dev/null
+++ b/cmd/pagegraph/main.go
@@ -0,0 +1,82 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"log"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"rescribe.xyz/bookpipeline"
+	"rescribe.xyz/utils/pkg/hocr"
+)
+
+func walker(confs *[]*bookpipeline.Conf) filepath.WalkFunc {
+	return func(path string, info os.FileInfo, err error) error {
+		if info.IsDir() {
+			return nil
+		}
+		if !strings.HasSuffix(path, ".hocr") {
+			return nil
+		}
+		avg, err := hocr.GetAvgConf(path)
+		if err != nil {
+			if err.Error() == "No words found" {
+				return nil
+			}
+			return err
+		}
+		c := bookpipeline.Conf{
+			Conf: avg,
+			Path: path,
+		}
+		*confs = append(*confs, &c)
+		return nil
+	}
+}
+
+func main() {
+	flag.Usage = func() {
+		fmt.Fprintln(flag.CommandLine.Output(), "Usage: pagegraph file.hocr graph.png")
+		flag.PrintDefaults()
+	}
+	flag.Parse()
+
+	if flag.NArg() != 2 {
+		flag.Usage()
+		return
+	}
+
+	wordconfs, err := hocr.GetWordConfs(flag.Arg(0))
+	if err != nil {
+		log.Fatal(err)
+	}
+	var confs []*bookpipeline.Conf
+	for n, wc := range wordconfs {
+		c := bookpipeline.Conf{
+			Conf: wc,
+			//Path: "fakepath",
+			Path: fmt.Sprintf("word_%d", n),
+		}
+		confs = append(confs, &c)
+	}
+
+	// Structure to fit what bookpipeline.Graph needs
+	// TODO: probably reorganise bookpipeline to just need []*Conf
+	cconfs := make(map[string]*bookpipeline.Conf)
+	for _, c := range confs {
+		cconfs[c.Path] = c
+	}
+
+	fn := flag.Arg(1)
+	f, err := os.Create(fn)
+	if err != nil {
+		log.Fatalln("Error creating file", fn, err)
+	}
+	defer f.Close()
+	err = bookpipeline.GraphOpts(cconfs, filepath.Base(flag.Arg(0)), "Word number", false, f)
+	if err != nil {
+		log.Fatalln("Error creating graph", err)
+	}
+}
diff --git a/graph.go b/graph.go
index 44e650b..4f8f186 100644
--- a/graph.go
+++ b/graph.go
@@ -43,6 +43,10 @@ func createLine(xvalues []float64, y float64, c drawing.Color) chart.ContinuousS
 }
 
 func Graph(confs map[string]*Conf, bookname string, w io.Writer) error {
+	return GraphOpts(confs, bookname, "Page number", true, w)
+}
+
+func GraphOpts(confs map[string]*Conf, bookname string, xaxis string, guidelines bool, w io.Writer) error {
 	if len(confs) == 0 {
 		return errors.New("No valid confidences")
 	}
@@ -98,12 +102,16 @@ func Graph(confs map[string]*Conf, bookname string, w io.Writer) error {
 	// Make last tick the final page
 	final := graphconf[len(graphconf)-1]
 	ticks[len(ticks)-1] = chart.Tick{final.Pgnum, fmt.Sprintf("%.0f", final.Pgnum)}
-	for i := 1; i <= yticknum; i++ {
+	for i := 0; i <= yticknum; i++ {
 		n := float64(i * 100) / yticknum
 		yticks = append(yticks, chart.Tick{n, fmt.Sprintf("%.1f", n)})
 	}
 
 	mainSeries := chart.ContinuousSeries{
+		Style: chart.Style{
+			StrokeColor:     chart.ColorBlue,
+			FillColor:       chart.ColorAlternateBlue,
+		},
 		XValues: xvalues,
 		YValues: yvalues,
 	}
@@ -145,7 +153,7 @@ func Graph(confs map[string]*Conf, bookname string, w io.Writer) error {
 	// Create annotations
 	var annotations []chart.Value2
 	for _, c := range graphconf {
-		if c.Conf > highconf || c.Conf < lowconf {
+		if !guidelines || (c.Conf > highconf || c.Conf < lowconf) {
 			annotations = append(annotations, chart.Value2{Label: fmt.Sprintf("%.0f", c.Pgnum), XValue: c.Pgnum, YValue: c.Conf})
 		}
 	}
@@ -157,7 +165,7 @@ func Graph(confs map[string]*Conf, bookname string, w io.Writer) error {
 		Width:      3840,
 		Height:     2160,
 		XAxis: chart.XAxis{
-			Name:      "Page number",
+			Name:      xaxis,
 			Range: &chart.ContinuousRange{
 				Min: 0.0,
 			},
@@ -173,15 +181,21 @@ func Graph(confs map[string]*Conf, bookname string, w io.Writer) error {
 		},
 		Series: []chart.Series{
 			mainSeries,
+			chart.AnnotationSeries{
+				Annotations: annotations,
+			},
+		},
+	}
+	if guidelines {
+		for _, s := range []chart.Series{
 			minSeries,
 			maxSeries,
 			goodCutoffSeries,
 			mediumCutoffSeries,
 			badCutoffSeries,
-			chart.AnnotationSeries{
-				Annotations: annotations,
-			},
-		},
+		} {
+			graph.Series = append(graph.Series, s)
+		}
 	}
 	return graph.Render(chart.PNG, w)
 }