Make bucket-lines and related packages more robust

bucket-lines would crash for any line that didn't have a corresponding image.

Lines which weren't grayscale would also cause crashes; now they are just
converted to grayscale if necessary.

As a bonus, lines in jpeg can also be decoded successfull.

2 files changed, 15 insertions, 6 deletions

diff --git a/bucket-lines/main.go b/bucket-lines/main.go
index 6ae506a..af81b44 100644
--- a/bucket-lines/main.go
+++ b/bucket-lines/main.go
@@ -74,7 +74,9 @@ func main() {
 		}
 
 		for _, l := range newlines {
-			lines = append(lines, l)
+			if l.Img != nil {
+				lines = append(lines, l)
+			}
 		}
 	}
 
diff --git a/pkg/hocr/lines.go b/pkg/hocr/lines.go
index e90b0a8..d2983be 100644
--- a/pkg/hocr/lines.go
+++ b/pkg/hocr/lines.go
@@ -5,7 +5,9 @@ package hocr
 
 import (
 	"image"
-	"image/png"
+	"image/draw"
+	_ "image/jpeg"
+	_ "image/png"
 	"io/ioutil"
 	"log"
 	"os"
@@ -48,7 +50,7 @@ func getLineText(l OcrLine) (string) {
 	return linetext
 }
 
-func parseLineDetails(h Hocr, i image.Image, name string) (line.Details, error) {
+func parseLineDetails(h Hocr, i *image.Gray, name string) (line.Details, error) {
 	lines := make(line.Details, 0)
 
 	for _, l := range h.Lines {
@@ -75,7 +77,7 @@ func parseLineDetails(h Hocr, i image.Image, name string) (line.Details, error)
 		ln.OcrName = name
 		if i != nil {
 			var imgd line.ImgDirect
-			imgd.Img = i.(*image.Gray).SubImage(image.Rect(coords[0], coords[1], coords[2], coords[3]))
+			imgd.Img = i.SubImage(image.Rect(coords[0], coords[1], coords[2], coords[3]))
 			ln.Img = imgd
 		}
 		lines = append(lines, ln)
@@ -97,20 +99,25 @@ func GetLineDetails(hocrfn string) (line.Details, error) {
 	}
 
 	var img image.Image
+	var gray *image.Gray
 	pngfn := strings.Replace(hocrfn, ".hocr", ".png", 1)
 	pngf, err := os.Open(pngfn)
 	if err != nil {
 		log.Println("Warning: can't open image %s\n", pngfn)
 	} else {
 		defer pngf.Close()
-		img, err = png.Decode(pngf)
+		img, _, err = image.Decode(pngf)
 		if err != nil {
 			log.Println("Warning: can't load image %s\n", pngfn)
+		} else {
+			b := img.Bounds()
+			gray = image.NewGray(image.Rect(0, 0, b.Dx(), b.Dy()))
+			draw.Draw(gray, b, img, b.Min, draw.Src)
 		}
 	}
 
 	n := strings.Replace(filepath.Base(hocrfn), ".hocr", "", 1)
-	return parseLineDetails(h, img, n)
+	return parseLineDetails(h, gray, n)
 }
 
 func GetLineBasics(hocrfn string) (line.Details, error) {