package main // TODO: combine this with line-conf-buckets, separating the parsing // out to a separate library, probably // see https://github.com/OCR-D/ocrd-train/issues/7 and https://github.com/OCR-D/ocrd-train/ // // TODO: Simplify things into functions more; this works well, but is a bit of a rush job // TODO: Parse line name to zero pad line numbers, so they come out in the correct order import ( "encoding/xml" "flag" "fmt" "image" "image/png" "io" "io/ioutil" "log" "os" "path/filepath" "regexp" "sort" "strconv" "strings" ) type LineDetail struct { name string avgconf float64 img image.Image text string hocrname string } type LineDetails []LineDetail // Used by sort.Sort. func (l LineDetails) Len() int { return len(l) } // Used by sort.Sort. func (l LineDetails) Less(i, j int) bool { return l[i].avgconf < l[j].avgconf } // Used by sort.Sort. func (l LineDetails) Swap(i, j int) { l[i], l[j] = l[j], l[i] } func copyline(filebase string, dirname string, basename string, avgconf string, outdir string, todir string) (err error) { outname := filepath.Join(outdir, todir, filepath.Base(dirname) + "_" + basename + "_" + avgconf) //log.Fatalf("I'd use '%s' as outname, and '%s' as filebase\n", outname, filebase) for _, extn := range []string{".bin.png", ".txt"} { infile, err := os.Open(filebase + extn) if err != nil { fmt.Fprintf(os.Stderr, "Failed to open %s\n", filebase + extn) return err } defer infile.Close() err = os.MkdirAll(filepath.Join(outdir, todir), 0700) if err != nil { return err } outfile, err := os.Create(outname + extn) if err != nil { fmt.Fprintf(os.Stderr, "Failed to create %s\n", outname + extn) return err } defer outfile.Close() _, err = io.Copy(outfile, infile) if err != nil { return err } } return err } type Hocr struct { Lines []OcrLine `xml:"body>div>div>p>span"` } type OcrLine struct { Class string `xml:"class,attr"` Id string `xml:"id,attr"` Title string `xml:"title,attr"` Words []OcrWord `xml:"span"` Text string `xml:",chardata"` } type OcrWord struct { Class string `xml:"class,attr"` Id string `xml:"id,attr"` Title string `xml:"title,attr"` Chars []OcrChar `xml:"span"` Text string `xml:",chardata"` } type OcrChar struct { Class string `xml:"class,attr"` Id string `xml:"id,attr"` Title string `xml:"title,attr"` Chars []OcrChar `xml:"span"` Text string `xml:",chardata"` } // Returns the confidence for a word based on its x_wconf value func wordConf(s string) (float64, error) { re, err := regexp.Compile(`x_wconf ([0-9.]+)`) if err != nil { return 0.0, err } conf := re.FindStringSubmatch(s) return strconv.ParseFloat(conf[1], 64) } func boxCoords(s string) ([4]int, error) { var coords [4]int re, err := regexp.Compile(`bbox ([0-9]+) ([0-9]+) ([0-9]+) ([0-9]+)`) if err != nil { return coords, err } coordstr := re.FindStringSubmatch(s) for i := range coords { c, err := strconv.Atoi(coordstr[i+1]) if err != nil { return coords, err } coords[i] = c } return coords, nil } func noText(s string) bool { t := strings.Trim(s, " \n") return len(t) == 0 } func main() { flag.Usage = func() { fmt.Fprintf(os.Stderr, "Usage: line-conf-buckets hocr1 [hocr2] [...]\n") fmt.Fprintf(os.Stderr, "Copies image-text line pairs into different directories according\n") fmt.Fprintf(os.Stderr, "to the average character probability for the line.\n") flag.PrintDefaults() } flag.Parse() if flag.NArg() < 1 { flag.Usage() os.Exit(1) } lines := make(LineDetails, 0) var hocr Hocr for _, f := range flag.Args() { file, err := ioutil.ReadFile(f) if err != nil { log.Fatal(err) } err = xml.Unmarshal(file, &hocr) if err != nil { log.Fatal(err) } pngfn := strings.Replace(f, ".hocr", ".png", 1) pngf, err := os.Open(pngfn) if err != nil { log.Fatal(err) } defer pngf.Close() img, err := png.Decode(pngf) if err != nil { log.Fatal(err) } for _, l := range hocr.Lines { totalconf := float64(0) num := 0 for _, w := range l.Words { c, err := wordConf(w.Title) if err != nil { log.Fatal(err) } num++ totalconf += c } coords, err := boxCoords(l.Title) if err != nil { log.Fatal(err) } var line LineDetail line.name = l.Id line.avgconf = (totalconf/float64(num)) / 100 linetext := "" linetext = l.Text if(noText(linetext)) { linetext = "" for _, w := range l.Words { if(w.Class != "ocrx_word") { continue } linetext += w.Text + " " } } if(noText(linetext)) { linetext = "" for _, w := range l.Words { if(w.Class != "ocrx_word") { continue } for _, c := range w.Chars { if(c.Class != "ocrx_cinfo") { continue } linetext += c.Text } linetext += " " } } line.text = strings.TrimRight(linetext, " ") line.text += "\n" line.hocrname = strings.Replace(filepath.Base(f), ".hocr", "", 1) line.img = img.(*image.Gray).SubImage(image.Rect(coords[0], coords[1], coords[2], coords[3])) lines = append(lines, line) } } sort.Sort(lines) worstnum := 0 mediumnum := 0 bestnum := 0 outdir := "buckets" // TODO: set this from cmdline todir := "" for _, l := range lines { switch { case l.avgconf < 0.95: todir = "bad" worstnum++ case l.avgconf < 0.98: todir = "95to98" mediumnum++ default: todir = "98plus" bestnum++ } avgstr := strconv.FormatFloat(l.avgconf, 'f', 5, 64) avgstr = avgstr[2:] fmt.Printf("Line: %s, avg: %f, avgstr: %s\n", l.name, l.avgconf, avgstr) outname := filepath.Join(outdir, todir, l.hocrname + "_" + l.name + "_" + avgstr + ".png") err := os.MkdirAll(filepath.Join(outdir, todir), 0700) if err != nil { log.Fatal(err) } outfile, err := os.Create(outname) if err != nil { fmt.Fprintf(os.Stderr, "Failed to create %s\n", outname) log.Fatal(err) } defer outfile.Close() err = png.Encode(outfile, l.img) if err != nil { log.Fatal(err) } outname = filepath.Join(outdir, todir, l.hocrname + "_" + l.name + "_" + avgstr + ".txt") outfile, err = os.Create(outname) if err != nil { fmt.Fprintf(os.Stderr, "Failed to create %s\n", outname) log.Fatal(err) } defer outfile.Close() _, err = io.WriteString(outfile, l.text) if err != nil { log.Fatal(err) } // TODO: test whether the line.img works properly with multiple hocrs, as it could be that as it's a pointer, it always points to the latest image (don't think so, but not sure) } total := worstnum + mediumnum + bestnum fmt.Printf("Copied lines to %s\n", outdir) fmt.Printf("---------------------------------\n") fmt.Printf("Lines 98%%+ quality: %d%%\n", 100 * bestnum / total) fmt.Printf("Lines 95-98%% quality: %d%%\n", 100 * mediumnum / total) fmt.Printf("Lines <95%% quality: %d%%\n", 100 * worstnum / total) }