filter: improve ocr filter

Author: martinlindhe

cleaner.go

@@ -44,7 +44,7 @@ var (
 		"seriessub",
 		"addic7ed", "addicted.com", "vaioholics",
 		"sdimedia", "sdi media",
-		"allsubs.org", "hdbits.org", "bierdopje.com", "subcentral",
+		"allsubs.org", "hdbits.org", "bierdopje.com", "subcentral", "mkvcage",
 		"cssubs", "tvsub", "uksubtitles",
 		"ragbear.com", "ydy.com", "yyets.net", "indivx.net", "sub-way.fr", "blogspot",
 		"forom.com", "forom. com", "facebook.com", "hdvietnam.com", "sapo.pt", "softhome.net",
@@ -57,7 +57,7 @@ var (
 		"thepiratebay", "anoxmous", "verdikt", "la fisher team", "red bee media",
 		"mkv player", "best watched using", "advertise your product", "remove all ads",
 		"memoryonsmells", "1st-booking",
-		":[gwc]:", "ripped with subrip", "titra film",
+		":[gwc]:", "ripped by", "ripped with subrip", "titra film",
 
 		// swedish:
 		"swedish subtitles", "svenska undertexter", "internationella undertexter",

filter_ocr.go

@@ -3,6 +3,7 @@ package subtitles
 import (
 	"strings"
 	"unicode"
+	"unicode/utf8"
 
 	log "github.com/Sirupsen/logrus"
 )
@@ -14,29 +15,32 @@ var (
 		"0n ":       "on ",
 		"c0uld":     "could",
 		"s0mething": "something",
-		"l've":      "i've",
+		"l've":      "I've",
+		"1 Oth":     "10th",
 	}
 )
 
 // filterOCR corrects some OCR mistakes
 func (subtitle *Subtitle) filterOCR() *Subtitle {
 	for _, cap := range subtitle.Captions {
 		for i, org := range cap.Text {
+			s := cap.Text[i]
 			for bad, good := range ocrErrors {
 				// lower case
-				cap.Text[i] = strings.Replace(cap.Text[i], bad, good, -1)
+				s = strings.Replace(s, bad, good, -1)
 
 				// upper case
-				cap.Text[i] = strings.Replace(cap.Text[i], strings.ToUpper(bad), strings.ToUpper(good), -1)
+				s = strings.Replace(s, strings.ToUpper(bad), strings.ToUpper(good), -1)
 
 				// ucfirst
-				cap.Text[i] = strings.Replace(cap.Text[i], strings.Title(bad), strings.Title(good), -1)
+				s = strings.Replace(s, 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])
+			s = fixOCRLineCapitalization(s)
+			if org != s {
+				log.Println("[ocr]", org, "->", s)
 			}
+			cap.Text[i] = s
 		}
 	}
 	return subtitle
@@ -52,26 +56,109 @@ func fixOCRLineCapitalization(s string) string {
 
 // fix capitalization errors due to ocr, GAsPs => GASPS
 func fixOCRWordCapitalization(s string) string {
-	if len(s) <= 3 {
+	if len(s) <= 3 || !isASCIIOnly(s) {
 		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++
-		}
+	// don't touch group of lowercase + uppercase such as in "macOS"
+	cases := countCaseInLetters(s)
+	if len(cases) < 4 {
+		return s
 	}
-	if upper >= 2 {
+
+	if countUppercaseLetters(s) >= 2 {
 		return strings.ToUpper(s)
 	}
-	if ucStart {
+	if startsWithUppercase(s) {
 		return strings.Title(s)
 	}
 	return strings.ToLower(s)
 }
+
+func countUppercaseLetters(s string) int {
+	upper := 0
+	for _, c := range s {
+		if unicode.IsUpper(c) {
+			upper++
+		}
+	}
+	return upper
+}
+
+func countLowercaseLetters(s string) int {
+	lower := 0
+	for _, c := range s {
+		if unicode.IsLower(c) {
+			lower++
+		}
+	}
+	return lower
+}
+
+func startsWithUppercase(s string) bool {
+	r, size := utf8.DecodeRuneInString(s)
+	if r == utf8.RuneError {
+		return false
+	}
+	if size > 0 && unicode.IsUpper(r) {
+		return true
+	}
+	return false
+}
+
+func isASCIIOnly(s string) bool {
+	for _, c := range s {
+		if (c < 'a' || c > 'z') && (c < 'A' || c > 'Z') {
+			return false
+		}
+	}
+	return true
+}
+
+type caseCount struct {
+	kind caseType
+	n    int
+}
+
+type caseType int
+
+const (
+	none caseType = iota
+	lower
+	upper
+)
+
+func getCase(c rune) caseType {
+	if unicode.IsUpper(c) {
+		return upper
+	}
+	if unicode.IsLower(c) {
+		return lower
+	}
+	return none
+}
+
+func countCaseInLetters(s string) []caseCount {
+	res := []caseCount{}
+	currentCount := 0
+	lastCase := none
+	for _, c := range s {
+		currentCase := getCase(c)
+		if lastCase == none {
+			lastCase = currentCase
+		}
+		if lastCase != currentCase {
+			if currentCount > 0 {
+				res = append(res, caseCount{lastCase, currentCount})
+				currentCount = 1
+				lastCase = currentCase
+			}
+		} else {
+			currentCount++
+		}
+	}
+	if currentCount > 0 {
+		res = append(res, caseCount{lastCase, currentCount})
+	}
+	return res
+}

filter_ocr_test.go

@@ -65,7 +65,7 @@ func TestFilterOCREnglish(t *testing.T) {
 		1,
 		makeTime(0, 0, 4, 630),
 		makeTime(0, 0, 6, 18),
-		[]string{"i've got a feeling"},
+		[]string{"I've got a feeling"},
 	}}}
 	assert.Equal(t, &expected, in.filterOCR())
 }
@@ -85,3 +85,23 @@ func TestFilterOCRCapitalization(t *testing.T) {
 	}}}
 	assert.Equal(t, &expected, in.filterOCR())
 }
+
+func TestFixOCRWordCapitalization(t *testing.T) {
+	assert.Equal(t, "He's", fixOCRWordCapitalization("He's"))
+	assert.Equal(t, "GASPS", fixOCRWordCapitalization("GAsPs"))
+
+	assert.Equal(t, "macOS", fixOCRWordCapitalization("macOS"))
+	assert.Equal(t, "WindowsXP", fixOCRWordCapitalization("WindowsXP"))
+}
+
+func TestStartsWithUppercase(t *testing.T) {
+	assert.Equal(t, true, startsWithUppercase("Allo"))
+	assert.Equal(t, true, startsWithUppercase("Ällo"))
+	assert.Equal(t, false, startsWithUppercase("allo"))
+}
+
+func TestCountCaseInLetters(t *testing.T) {
+	assert.Equal(t, []caseCount{{upper, 2}}, countCaseInLetters("GA"))
+	assert.Equal(t, []caseCount{{lower, 2}}, countCaseInLetters("ga"))
+	assert.Equal(t, []caseCount{{upper, 2}, {lower, 1}, {upper, 1}, {lower, 1}}, countCaseInLetters("GAsPs"))
+}