feat: add YouTube subtitle download & embed support (Fixes #288) (#1430)

* feat(youtube): add support for video captions in youtubeDownload

* feat(youtube): filter subtitles based on Items option in Extract method

* feat(youtube): display available caption languages in printInfo

* feat(youtube): add option to embed subtitles into video during download

* refactor(downloader): use utils.ConvertXmlToSrt and remove redundant function

* test(utils): add unit test for ConvertXMLToSRT function

* feat(readme): update subtitle download options and embed feature for YouTube

* refactor(downloader): rename ConvertXmlToSrt to ConvertXMLFileToSRT and update implementation

* refactor(ffmpeg): enhance subtitle handling with ISO 639-2 mapping and codec selection

* refactor(utils): simplify XML unmarshalling in ConvertXMLToSRT function

Author: MelonHiker

README.md

@@ -571,7 +571,11 @@ $ lux -j "https://www.bilibili.com/video/av20203945"
 #### Subtitle:
 
 ```
-  -C	Download captions
+  -C	Download subtitles
+  -C -items en,zh
+    	Download specific languages (YouTube only)
+  -C -items en,zh -embed 
+    	Embed subtitles into the video (YouTube only)
 ```
 
 #### Youku:

app/app.go

@@ -119,6 +119,11 @@ func New() *cli.App {
 				Aliases: []string{"C"},
 				Usage:   "Download captions",
 			},
+			&cli.BoolFlag{
+				Name:    "embed-subtitle",
+				Aliases: []string{"embed"},
+				Usage:   "Embed subtitles into the video (requires ffmpeg)",
+			},
 
 			&cli.UintFlag{
 				Name:  "start",
@@ -311,6 +316,7 @@ func download(c *cli.Context, videoURL string) error {
 		OutputName:     c.String("output-name"),
 		FileNameLength: int(c.Uint("file-name-length")),
 		Caption:        c.Bool("caption"),
+		EmbedSubtitle:  c.Bool("embed-subtitle"),
 		MultiThread:    c.Bool("multi-thread"),
 		ThreadNumber:   int(c.Uint("thread")),
 		RetryTimes:     int(c.Uint("retry")),

downloader/downloader.go

@@ -12,6 +12,7 @@ import (
 	"path/filepath"
 	"regexp"
 	"sort"
+	"strings"
 	"sync"
 	"time"
 
@@ -34,6 +35,7 @@ type Options struct {
 	OutputName     string
 	FileNameLength int
 	Caption        bool
+	EmbedSubtitle  bool
 
 	MultiThread  bool
 	ThreadNumber int
@@ -605,12 +607,28 @@ func (downloader *Downloader) Download(data *extractors.Data) error {
 	}
 
 	// download caption
+	var subtitlePaths []string
+	var subtitleLangs []string
+	var subtitleFilesToDelete []string
 	if downloader.option.Caption && data.Captions != nil {
 		fmt.Println("\nDownloading captions...")
 		for k, v := range data.Captions {
 			if v != nil {
 				fmt.Printf("Downloading %s ...\n", k)
-				downloader.caption(v.URL, title, v.Ext, v.Transform) // nolint
+				if err := downloader.caption(v.URL, title, v.Ext, v.Transform); err != nil {
+					// nolint
+				} else if downloader.option.EmbedSubtitle {
+					subtitlePath, _ := utils.FilePath(title, v.Ext, downloader.option.FileNameLength, downloader.option.OutputPath, true)
+					subtitleFilesToDelete = append(subtitleFilesToDelete, subtitlePath)
+					if strings.HasSuffix(v.Ext, "xml") {
+						if srtPath, err := utils.ConvertXMLFileToSRT(subtitlePath); err == nil {
+							subtitlePath = srtPath
+							subtitleFilesToDelete = append(subtitleFilesToDelete, srtPath)
+						}
+					}
+					subtitlePaths = append(subtitlePaths, subtitlePath)
+					subtitleLangs = append(subtitleLangs, k)
+				}
 			}
 		}
 	}
@@ -652,6 +670,18 @@ func (downloader *Downloader) Download(data *extractors.Data) error {
 			return err
 		}
 		downloader.bar.Finish()
+
+		if downloader.option.EmbedSubtitle && len(subtitlePaths) > 0 {
+			if !downloader.option.Silent {
+				fmt.Println("Embedding subtitles...")
+			}
+			if err := utils.EmbedSubtitles(mergedFilePath, subtitlePaths, subtitleLangs); err != nil {
+				return err
+			}
+			for _, path := range subtitleFilesToDelete {
+				os.Remove(path)
+			}
+		}
 		return nil
 	}
 
@@ -706,7 +736,26 @@ func (downloader *Downloader) Download(data *extractors.Data) error {
 		fmt.Printf("Merging video parts into %s\n", mergedFilePath)
 	}
 	if stream.Ext != "mp4" || stream.NeedMux {
-		return utils.MergeFilesWithSameExtension(parts, mergedFilePath)
+		if err := utils.MergeFilesWithSameExtension(parts, mergedFilePath); err != nil {
+			return err
+		}
+	} else {
+		if err := utils.MergeToMP4(parts, mergedFilePath, title); err != nil {
+			return err
+		}
+	}
+
+	if downloader.option.EmbedSubtitle && len(subtitlePaths) > 0 {
+		if !downloader.option.Silent {
+			fmt.Println("Embedding subtitles...")
+		}
+		if err := utils.EmbedSubtitles(mergedFilePath, subtitlePaths, subtitleLangs); err != nil {
+			return err
+		}
+		for _, path := range subtitleFilesToDelete {
+			os.Remove(path)
+		}
 	}
-	return utils.MergeToMP4(parts, mergedFilePath, title)
+
+	return nil
 }

downloader/utils.go

@@ -51,7 +51,23 @@ func printStream(stream *extractors.Stream) {
 
 func printInfo(data *extractors.Data, sortedStreams []*extractors.Stream) {
 	printHeader(data)
-
+	if len(data.Captions) > 0 {
+		cyan.Printf(" Captions:  ") // nolint
+		languages := make([]string, 0, len(data.Captions))
+		for lang := range data.Captions {
+			languages = append(languages, lang)
+		}
+		sort.Strings(languages)
+		captionList := ""
+		for _, lang := range languages {
+			caption := data.Captions[lang]
+			if caption == nil {
+				continue
+			}
+			captionList += fmt.Sprintf("%s ", lang)
+		}
+		fmt.Println(captionList)
+	}
 	cyan.Printf(" Streams:   ") // nolint
 	fmt.Println("# All available quality")
 	for _, stream := range sortedStreams {

extractors/youtube/youtube.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"slices"
 	"strconv"
+	"strings"
 
 	"github.com/kkdai/youtube/v2"
 	"github.com/pkg/errors"
@@ -46,7 +47,19 @@ func (e *extractor) Extract(url string, option extractors.Options) ([]*extractor
 		if err != nil {
 			return nil, errors.WithStack(err)
 		}
-		return []*extractors.Data{e.youtubeDownload(url, video)}, nil
+		data := e.youtubeDownload(url, video)
+		if option.Items != "" {
+			// If it is not a playlist, we can use the Items option to filter the subtitles.
+			filteredCaptions := make(map[string]*extractors.CaptionPart)
+			items := strings.Split(option.Items, ",")
+			for k, v := range data.Captions {
+				if slices.Contains(items, k) {
+					filteredCaptions[k] = v
+				}
+			}
+			data.Captions = filteredCaptions
+		}
+		return []*extractors.Data{data}, nil
 	}
 
 	playlist, err := e.client.GetPlaylist(url)
@@ -127,12 +140,23 @@ func (e *extractor) youtubeDownload(url string, video *youtube.Video) *extractor
 		streams[itag] = stream
 	}
 
+	captions := make(map[string]*extractors.CaptionPart)
+	for _, c := range video.CaptionTracks {
+		captions[c.LanguageCode] = &extractors.CaptionPart{
+			Part: extractors.Part{
+				URL: c.BaseURL,
+				Ext: c.LanguageCode + ".xml",
+			},
+		}
+	}
+
 	return &extractors.Data{
-		Site:    "YouTube youtube.com",
-		Title:   video.Title,
-		Type:    "video",
-		Streams: streams,
-		URL:     url,
+		Site:     "YouTube youtube.com",
+		Title:    video.Title,
+		Type:     "video",
+		Streams:  streams,
+		Captions: captions,
+		URL:      url,
 	}
 }
 

utils/ffmpeg.go

@@ -7,6 +7,7 @@ import (
 	"os/exec"
 	"path/filepath"
 	"runtime"
+	"strings"
 
 	"github.com/pkg/errors"
 )
@@ -76,3 +77,81 @@ func MergeToMP4(paths []string, mergedFilePath string, filename string) error {
 	)
 	return runMergeCmd(cmd, paths, mergeFilePath)
 }
+
+// ISO 639-2 language code mapping
+var langToISO = map[string]string{
+	"zh": "chi", "en": "eng", "ja": "jpn", "ko": "kor",
+	"es": "spa", "fr": "fre", "de": "ger", "ru": "rus",
+	"pt": "por", "it": "ita", "nl": "nld", "sv": "swe",
+	"no": "nor", "fi": "fin", "da": "dan", "pl": "pol",
+}
+
+// toISO639 converts language code (e.g. "en-US") to ISO 639-2 (e.g. "eng")
+func toISO639(lang string) string {
+	base := strings.ToLower(lang)
+	if i := strings.IndexAny(base, "-_"); i != -1 {
+		base = base[:i]
+	}
+	if iso, ok := langToISO[base]; ok {
+		return iso
+	}
+	if len(base) == 3 {
+		return base
+	}
+	return "und"
+}
+
+// subtitleCodec returns the appropriate subtitle codec for the container format
+func subtitleCodec(ext string) string {
+	switch strings.ToLower(ext) {
+	case ".mp4":
+		return "mov_text"
+	case ".webm":
+		return "webvtt"
+	default:
+		return ""
+	}
+}
+
+// EmbedSubtitles embeds subtitles into the video.
+func EmbedSubtitles(videoPath string, subtitlePaths []string, langs []string) error {
+	ext := filepath.Ext(videoPath)
+	tempOutput := videoPath + ".temp" + ext
+
+	// Build ffmpeg command
+	cmds := []string{"-y", "-i", videoPath}
+	for _, subPath := range subtitlePaths {
+		cmds = append(cmds, "-i", subPath)
+	}
+
+	cmds = append(cmds,
+		"-map", "0", "-dn", "-ignore_unknown",
+		"-c", "copy",
+	)
+
+	if codec := subtitleCodec(ext); codec != "" {
+		cmds = append(cmds, "-c:s", codec)
+	}
+
+	// Exclude existing subtitles, then map new ones
+	cmds = append(cmds, "-map", "-0:s")
+	for i, lang := range langs {
+		if i >= len(subtitlePaths) {
+			break
+		}
+		iso := toISO639(lang)
+		cmds = append(cmds,
+			"-map", fmt.Sprintf("%d:0", i+1),
+			fmt.Sprintf("-metadata:s:s:%d", i), fmt.Sprintf("language=%s", iso),
+			fmt.Sprintf("-metadata:s:s:%d", i), fmt.Sprintf("handler_name=%s", lang),
+			fmt.Sprintf("-metadata:s:s:%d", i), fmt.Sprintf("title=%s", lang),
+		)
+	}
+
+	cmds = append(cmds, tempOutput)
+
+	if err := runMergeCmd(exec.Command(findFFmpegExecutable(), cmds...), []string{videoPath}, ""); err != nil {
+		return err
+	}
+	return os.Rename(tempOutput, videoPath)
+}

utils/utils.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"bytes"
 	"crypto/md5"
+	"encoding/xml"
 	"fmt"
 	"io"
 	"net/url"
@@ -19,6 +20,78 @@ import (
 	"github.com/iawia002/lux/request"
 )
 
+// ConvertXMLToSRT converts YouTube XML subtitles to SRT format
+func ConvertXMLToSRT(xmlContent []byte) (string, error) {
+	var data struct {
+		Body struct {
+			P []struct {
+				T    int    `xml:"t,attr"`
+				D    int    `xml:"d,attr"`
+				Text string `xml:",chardata"`
+				S    []struct {
+					T    int    `xml:"t,attr"`
+					Text string `xml:",chardata"`
+				} `xml:"s"`
+			} `xml:"p"`
+		} `xml:"body"`
+	}
+
+	if err := xml.Unmarshal(xmlContent, &data); err != nil {
+		return "", err
+	}
+
+	var srtBuilder strings.Builder
+	index := 1
+	for _, p := range data.Body.P {
+		startTime := formatSRTTime(p.T)
+		endTime := formatSRTTime(p.T + p.D)
+
+		// Handle text content
+		var text string
+		if len(p.S) > 0 {
+			for _, s := range p.S {
+				text += s.Text
+			}
+		} else {
+			text = p.Text
+		}
+		text = strings.TrimSpace(text)
+
+		// Skip empty lines
+		if text == "" {
+			continue
+		}
+
+		srtBuilder.WriteString(fmt.Sprintf("%d\n%s --> %s\n%s\n\n", index, startTime, endTime, text))
+		index++
+	}
+	return srtBuilder.String(), nil
+}
+
+// ConvertXMLFileToSRT converts XML subtitles file to SRT format
+func ConvertXMLFileToSRT(xmlPath string) (string, error) {
+	content, err := os.ReadFile(xmlPath)
+	if err != nil {
+		return "", err
+	}
+	srtContent, err := ConvertXMLToSRT(content)
+	if err != nil {
+		return "", err
+	}
+	srtPath := xmlPath[:len(xmlPath)-len("xml")] + "srt"
+	return srtPath, os.WriteFile(srtPath, []byte(srtContent), 0644)
+}
+
+func formatSRTTime(ms int) string {
+	hours := ms / 3600000
+	ms %= 3600000
+	minutes := ms / 60000
+	ms %= 60000
+	seconds := ms / 1000
+	ms %= 1000
+	return fmt.Sprintf("%02d:%02d:%02d,%03d", hours, minutes, seconds, ms)
+}
+
 // MatchOneOf match one of the patterns
 func MatchOneOf(text string, patterns ...string) []string {
 	var (