feat: add .srt subtitles support (#2592)

files/file.go

@@ -14,6 +14,7 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"regexp"
 	"strings"
 	"time"
 
@@ -290,27 +291,31 @@ func (i *FileInfo) detectSubtitles() {
 	}
 
 	i.Subtitles = []string{}
-	ext := filepath.Ext(i.Path)
 
-	// detect multiple languages. Base*.vtt
+	// detect multiple languages. Base *.vtt/*.srt
 	// TODO: give subtitles descriptive names (lang) and track attributes
 	parentDir := strings.TrimRight(i.Path, i.Name)
-	var dir []os.FileInfo
+	// pattern to match against <video_base_name>.*\.(vtt|srt)
-	if len(i.currentDir) > 0 {
+	subsRegex := regexp.MustCompile("(?i)" + strings.TrimSuffix(filepath.Base(i.Name), filepath.Ext(i.Name)) + ".*\\.(vtt|srt)$")
-		dir = i.currentDir
+	subsFoldersRegex := regexp.MustCompile(`(?i)\bsub(s|titles)$`)
-	} else {
+
-		var err error
+	// find .vtt/.srt files in current dir or nested folders, Subtitles/Subs/subtitles/subs
-		dir, err = afero.ReadDir(i.Fs, parentDir)
+	subsWalker := func(path string, fileInfo os.FileInfo, err error) error {
 		if err != nil {
-			return
+			return err
-		}
 		}
 
-	base := strings.TrimSuffix(i.Name, ext)
+		if fileInfo.IsDir() && path != parentDir && !subsFoldersRegex.MatchString(path) {
-	for _, f := range dir {
+			return filepath.SkipDir
-		if !f.IsDir() && strings.HasPrefix(f.Name(), base) && strings.HasSuffix(f.Name(), ".vtt") {
+		} else if subsRegex.MatchString(filepath.Base(path)) {
-			i.Subtitles = append(i.Subtitles, path.Join(parentDir, f.Name()))
+			i.Subtitles = append(i.Subtitles, path)
 		}
+		return nil
+	}
+
+	err := afero.Walk(i.Fs, parentDir, subsWalker)
+	if err != nil {
+		return
 	}
 }
 

files/file_test.go

@@ -0,0 +1,74 @@
+package files
+
+import (
+	"testing"
+
+	"github.com/spf13/afero"
+)
+
+func TestDetectSubtitles(t *testing.T) {
+	testCases := []struct {
+		path string
+		want bool // is file detected as subtitles?
+	}{
+		{path: "/media/movie.mkv", want: false},
+		{path: "/media/movie.vtt", want: true},
+		{path: "/media/movie.en.srt", want: true},
+		{path: "/media/Subs/movie.pt.srt", want: true},
+		{path: "/media/subs/movie.zh-tw.srt", want: true},
+		{path: "/media/movie.es-es.srt", want: true},
+		{path: "/media/movie.fr.srt", want: true},
+		{path: "/media/srt", want: false},
+		{path: "/media/movie.dir.vtt", want: false},
+		{path: "/media/subs/movie.dir.srt", want: false},
+	}
+
+	fs := afero.NewMemMapFs()
+	err0 := fs.MkdirAll("/media", PermDir)
+	if err0 != nil {
+		t.Fatalf("Failed to create directory: %v", err0)
+	}
+	err1 := fs.MkdirAll("/media/movie.dir.vtt", PermDir)
+	if err1 != nil {
+		t.Fatalf("Failed to create directory: %v", err1)
+	}
+	err2 := fs.MkdirAll("/media/subs/movie.dir.srt", PermDir)
+	if err2 != nil {
+		t.Fatalf("Failed to create directory: %v", err2)
+	}
+
+	for _, path := range testCases {
+		err := afero.WriteFile(fs, path.path, []byte("data"), PermFile)
+		if err != nil {
+			return
+		}
+	}
+
+	file := &FileInfo{
+		Fs:        fs,
+		Path:      "/media/movie.mkv",
+		Name:      "movie.mkv",
+		Type:      "video",
+		Size:      42,
+		Extension: ".mkv",
+	}
+
+	file.detectSubtitles()
+
+	for _, tt := range testCases {
+		t.Run(tt.path, func(t *testing.T) {
+			if got := contains(file.Subtitles, tt.path); got != tt.want {
+				t.Errorf("detectSubtitles() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func contains(s []string, e string) bool {
+	for _, v := range s {
+		if v == e {
+			return true
+		}
+	}
+	return false
+}

frontend/src/views/files/Preview.vue

@@ -76,8 +76,8 @@
             kind="captions"
             v-for="(sub, index) in subtitles"
             :key="index"
-            :src="sub"
+            :src="sub.src"
-            :label="'Subtitle ' + index"
+            :label="sub.label"
             :default="index === 0"
           />
           Sorry, your browser doesn't support embedded videos, but don't worry,
@@ -176,6 +176,19 @@ export default {
       nextRaw: "",
     };
   },
+  asyncComputed: {
+    async subtitles() {
+      if (this.req.subtitles) {
+        const subs = await Promise.all(
+          api
+            .getSubtitlesURL(this.req)
+            .map(async (s) => await this.loadSubtitle(s))
+        );
+        return subs.filter((s) => s.src != undefined);
+      }
+      return [];
+    },
+  },
   computed: {
     ...mapState(["req", "user", "oldReq", "jwt", "loading"]),
     ...mapGetters(["currentPrompt"]),
@@ -201,12 +214,6 @@ export default {
     isResizeEnabled() {
       return resizePreview;
     },
-    subtitles() {
-      if (this.req.subtitles) {
-        return api.getSubtitlesURL(this.req);
-      }
-      return [];
-    },
   },
   watch: {
     $route: function () {
@@ -350,6 +357,76 @@ export default {
     download() {
       window.open(this.downloadUrl);
     },
+    async loadSubtitle(subUrl) {
+      const _self = this;
+      const url = new URL(subUrl);
+      const label = url.pathname
+        .split("/")
+        .pop()
+        .replace(/\.[^/.]+$/, "");
+      let src;
+      if (url.pathname.toLowerCase().endsWith(".srt")) {
+        try {
+          const resp = await fetch(subUrl);
+          if (!resp.ok) {
+            throw new Error(`Failed to fetch subtitle from ${subUrl}!`);
+          }
+          const vtt = _self.srtToVttBlob(await resp.text());
+          src = URL.createObjectURL(vtt);
+        } catch (error) {
+          console.error(error);
+        }
+      } else {
+        src = subUrl;
+      }
+
+      return { src, label };
+    },
+    srtToVttBlob(srtData) {
+      const VTT_HEAD = `WEBVTT
+
+STYLE
+::cue {font-size: 3vh; color:#ffffff; border:3px solid black; text-align:center;text-shadow: -1px -1px 1px rgba(255,255,255,.1), 1px 1px 1px rgba(0,0,0,.5), 2px 2px 3px rgba(206,89,55,0)}
+::cue(b) {color: rgb(51, 216, 18); border: 2px solid black;}
+::cue(i) {color: #00bafd;}
+::cue(u) {color: #ff00ee;}
+`;
+      // Replace line breaks with \n
+      let subtitles = srtData.replace(/\r\n|\r|\n/g, "\n");
+      // commas -> dots in timestamps
+      subtitles = subtitles.replace(/(\d\d:\d\d:\d\d),(\d\d\d)/g, "$1.$2");
+      // map SRT font colors to VTT cue span classes
+      var colorMap = {};
+
+      // font tags -> ::cue span tags
+      subtitles = subtitles.replace(
+        /<font color="([^"]+)">([\s\S]*?)<\/font>/g,
+        function (_match, color, text) {
+          let key =
+            "c_" + color.replace(/^rgb/, "").replace(/\W/g, "").toLowerCase();
+          colorMap[key] = color;
+          return `<c.${key}>${text.replace("\n", "").trim()}</c>`;
+        }
+      );
+      subtitles = subtitles.replace(/<br\s*\/?>/g, "\n");
+
+      let vttSubtitles = VTT_HEAD;
+
+      if (Object.keys(colorMap).length) {
+        let vttStyles = "";
+
+        for (let cssClass in colorMap) {
+          let color = colorMap[cssClass];
+          // add cue style declaration
+          vttStyles += `::cue(.${cssClass}) {color: ${color};}\n`;
+        }
+        vttSubtitles += vttStyles;
+      }
+      vttSubtitles += "\n"; // an empty line MUST separate styles from subs
+      vttSubtitles += subtitles;
+
+      return new Blob([vttSubtitles], { type: "text/vtt" });
+    },
   },
 };
 </script>