Add basic analyse step, working but incomplete

1 files changed, 107 insertions, 7 deletions

diff --git a/bookpipeline/main.go b/bookpipeline/main.go
index 22bf7f5..4596171 100644
--- a/bookpipeline/main.go
+++ b/bookpipeline/main.go
@@ -15,6 +15,7 @@ import (
 	"strings"
 	"time"
 
+	"rescribe.xyz/go.git/lib/hocr"
 	"rescribe.xyz/go.git/preproc"
 )
 
@@ -130,6 +131,86 @@ func ocr(training string) func(chan string, chan string, chan error, *log.Logger
 	}
 }
 
+type Conf struct {
+	path, code string
+	conf float64
+}
+
+func analyse(toanalyse chan string, up chan string, errc chan error, logger *log.Logger) {
+	confs := make(map[string][]*Conf)
+	bestconfs := make(map[string]*Conf)
+	savedir := ""
+
+	for path := range toanalyse {
+		if savedir == "" {
+			savedir = filepath.Dir(path)
+		}
+		logger.Println("Calculating confidence for", path)
+		avg, err := hocr.GetAvgConf(path)
+		if err != nil {
+			close(up)
+			errc <- errors.New(fmt.Sprintf("Error retreiving confidence for %s: %s", path, err))
+			return
+		}
+		base := filepath.Base(path)
+		codestart := strings.Index(base, "_bin")
+		name := base[0:codestart]
+		var c Conf
+		c.path = path
+		c.code = base[codestart:]
+		c.conf = avg
+		confs[name] = append(confs[name], &c)
+
+	}
+
+	conffn := filepath.Join(savedir, "conf")
+	logger.Println("Saving confidences in file", conffn)
+	f, err := os.Create(conffn)
+	if err != nil {
+		close(up)
+		errc <- errors.New(fmt.Sprintf("Error creating conf file %s: %s", conffn, err))
+		return
+	}
+	defer f.Close()
+
+	logger.Println("Finding best confidence for each page, and saving all confidences")
+	for base, conf := range confs {
+		var best float64
+		for _, c := range conf {
+			if c.conf > best {
+				best = c.conf
+				bestconfs[base] = c
+			}
+			_, err = fmt.Fprintf(f, "%s\t%02.f\n", c.path, c.conf)
+			if err != nil {
+				close(up)
+				errc <- errors.New(fmt.Sprintf("Error writing confidences file: %s", err))
+				return
+			}
+		}
+	}
+	up <- conffn
+
+	logger.Println("Creating best file listing the best file for each page")
+	bestfn := filepath.Join(savedir, "best")
+	f, err = os.Create(bestfn)
+	if err != nil {
+		close(up)
+		errc <- errors.New(fmt.Sprintf("Error creating best file %s: %s", bestfn, err))
+		return
+	}
+	defer f.Close()
+	for _, conf := range bestconfs {
+		_, err = fmt.Fprintf(f, "%s\n", filepath.Base(conf.path))
+	}
+	up <- bestfn
+
+	// TODO: plot a graph with the confs, using https://github.com/gonum/plot, send to up
+	// TODO: generate a general report.txt with statistics etc for the book, send to up
+
+	close(up)
+}
+
 func processBook(msg Qmsg, conn Pipeliner, process func(chan string, chan string, chan error, *log.Logger), match *regexp.Regexp, fromQueue string, toQueue string) error {
 	bookname := msg.Body
 
@@ -183,12 +264,14 @@ func processBook(msg Qmsg, conn Pipeliner, process func(chan string, chan string
 	case <-done:
 	}
 
-	conn.Logger().Println("Sending", bookname, "to queue")
-	err = conn.AddToQueue(toQueue, bookname)
-	if err != nil {
-		t.Stop()
-		_ = os.RemoveAll(d)
-		return errors.New(fmt.Sprintf("Error adding to queue %s: %s", bookname, err))
+	if toQueue != "" {
+		conn.Logger().Println("Sending", bookname, "to queue")
+		err = conn.AddToQueue(toQueue, bookname)
+		if err != nil {
+			t.Stop()
+			_ = os.RemoveAll(d)
+			return errors.New(fmt.Sprintf("Error adding to queue %s: %s", bookname, err))
+		}
 	}
 
 	t.Stop()
@@ -228,7 +311,7 @@ func main() {
 	// TODO: match jpg too
 	origPattern := regexp.MustCompile(`[0-9]{4}.jpg$`) // TODO: match other file naming
 	preprocessedPattern := regexp.MustCompile(`_bin[0-9].[0-9].png$`)
-	//ocredPattern := regexp.MustCompile(`.hocr$`)
+	ocredPattern := regexp.MustCompile(`.hocr$`)
 
 	var conn Pipeliner
 	conn = &awsConn{region: "eu-west-2", logger: verboselog}
@@ -242,8 +325,10 @@ func main() {
 
 	var checkPreQueue <-chan time.Time
 	var checkOCRQueue <-chan time.Time
+	var checkAnalyseQueue <-chan time.Time
 	checkPreQueue = time.After(0)
 	checkOCRQueue = time.After(0)
+	checkAnalyseQueue = time.After(0)
 
 	for {
 		select {
@@ -277,6 +362,21 @@ func main() {
 			if err != nil {
 				log.Println("Error during OCR process", err)
 			}
+		case <-checkAnalyseQueue:
+			msg, err := conn.CheckQueue(conn.AnalyseQueueId())
+			checkAnalyseQueue = time.After(PauseBetweenChecks)
+			if err != nil {
+				log.Println("Error checking analyse queue", err)
+				continue
+			}
+			if msg.Handle == "" {
+				verboselog.Println("No message received on analyse queue, sleeping")
+				continue
+			}
+			err = processBook(msg, conn, analyse, ocredPattern, conn.AnalyseQueueId(), "")
+			if err != nil {
+				log.Println("Error during analysis", err)
+			}
 		}
 	}
 }