ocr: normalize capitalization

Author: martinlindhe

filter_ocr.go

@@ -2,6 +2,7 @@ package subtitles
 
 import (
 	"strings"
+	"unicode"
 
 	log "github.com/Sirupsen/logrus"
 )
@@ -32,10 +33,45 @@ func (subtitle *Subtitle) filterOCR() *Subtitle {
 				cap.Text[i] = strings.Replace(cap.Text[i], strings.Title(bad), strings.Title(good), -1)
 			}
 
+			cap.Text[i] = fixOCRLineCapitalization(cap.Text[i])
 			if org != cap.Text[i] {
 				log.Println("[ocr]", org, "->", cap.Text[i])
 			}
 		}
 	}
 	return subtitle
 }
+
+func fixOCRLineCapitalization(s string) string {
+	words := strings.Split(s, " ")
+	for i := range words {
+		words[i] = fixOCRWordCapitalization(words[i])
+	}
+	return strings.Join(words, " ")
+}
+
+// fix capitalization errors due to ocr, GAsPs => GASPS
+func fixOCRWordCapitalization(s string) string {
+	if len(s) <= 3 {
+		return s
+	}
+
+	// if starts with uc, or at least 2 letters is upper, make all upper
+	upper := 0
+	ucStart := false
+	for i, char := range s {
+		if i == 0 && unicode.IsUpper(char) {
+			ucStart = true
+		}
+		if unicode.IsUpper(char) {
+			upper++
+		}
+	}
+	if upper >= 2 {
+		return strings.ToUpper(s)
+	}
+	if ucStart {
+		return strings.Title(s)
+	}
+	return strings.ToLower(s)
+}

filter_ocr_test.go

@@ -81,3 +81,22 @@ func TestFilterOCREnglish(t *testing.T) {
 
 	assert.Equal(t, &expected, in.filterOCR())
 }
+
+func TestFilterOCRCapitalization(t *testing.T) {
+
+	in := Subtitle{[]Caption{{
+		1,
+		makeTime(0, 0, 4, 630),
+		makeTime(0, 0, 6, 18),
+		[]string{"GAsPs slowly"},
+	}}}
+
+	expected := Subtitle{[]Caption{{
+		1,
+		makeTime(0, 0, 4, 630),
+		makeTime(0, 0, 6, 18),
+		[]string{"GASPS slowly"},
+	}}}
+
+	assert.Equal(t, &expected, in.filterOCR())
+}