diff --git a/cmd/subtle/main.go b/cmd/subtle/main.go index fa188e9..9c491c4 100644 --- a/cmd/subtle/main.go +++ b/cmd/subtle/main.go @@ -1,16 +1,33 @@ package main import ( + "fmt" "log" - "github.com/nandesh-dev/subtle/internal/routine/library" + "github.com/nandesh-dev/subtle/internal/routine/extract" + "github.com/nandesh-dev/subtle/internal/routine/media" "github.com/nandesh-dev/subtle/pkgs/config" + "github.com/nandesh-dev/subtle/pkgs/db" ) func main() { - if err := config.Init("/config"); err != nil { + fmt.Println("Initilizing config") + if err := config.Init("./config"); err != nil { log.Fatal(err) } - library.RunLibraryRoutine() + fmt.Println("Initilizing database") + if err := db.Init(); err != nil { + log.Fatal(err) + } + + fmt.Println("Running media routine") + media.Run() + + fmt.Println("Running extract routine") + warns := extract.Run() + + for _, warning := range warns.Warnings() { + fmt.Println(warning) + } } diff --git a/go.mod b/go.mod index aa6db72..01b57fb 100644 --- a/go.mod +++ b/go.mod @@ -5,15 +5,20 @@ go 1.22.6 require ( github.com/otiai10/gosseract/v2 v2.4.1 github.com/u2takey/ffmpeg-go v0.5.0 - golang.org/x/text v0.18.0 + golang.org/x/text v0.19.0 google.golang.org/grpc v1.67.0 google.golang.org/protobuf v1.34.2 gopkg.in/yaml.v3 v3.0.1 + gorm.io/driver/sqlite v1.5.6 + gorm.io/gorm v1.25.12 ) require ( github.com/aws/aws-sdk-go v1.38.20 // indirect + github.com/jinzhu/inflection v1.0.0 // indirect + github.com/jinzhu/now v1.1.5 // indirect github.com/jmespath/go-jmespath v0.4.0 // indirect + github.com/mattn/go-sqlite3 v1.14.24 // indirect github.com/u2takey/go-utils v0.3.1 // indirect golang.org/x/net v0.28.0 // indirect golang.org/x/sys v0.24.0 // indirect diff --git a/go.sum b/go.sum index 37a2865..37276b1 100644 --- a/go.sum +++ b/go.sum @@ -12,6 +12,10 @@ github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeN github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4= +github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= +github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc= +github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ= +github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8= github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg= github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo= github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8= @@ -19,6 +23,8 @@ github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfC github.com/json-iterator/go v1.1.10/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4= github.com/kisielk/errcheck v1.2.0/go.mod h1:/BMXB+zMLi60iA8Vv6Ksmxu/1UDYcXs4uQLJ+jE2L00= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= +github.com/mattn/go-sqlite3 v1.14.24 h1:tpSp2G2KyMnnQu99ngJ47EIkWVmliIizyZBfPrBWDRM= +github.com/mattn/go-sqlite3 v1.14.24/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0= @@ -59,8 +65,8 @@ golang.org/x/sys v0.24.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= -golang.org/x/text v0.18.0 h1:XvMDiNzPAl0jr17s6W9lcaIhGUfUORdGCNsuLmPG224= -golang.org/x/text v0.18.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= +golang.org/x/text v0.19.0 h1:kTxAhCbGbxhK0IwgSKiMO5awPoDQ0RpfiVYBfK860YM= +golang.org/x/text v0.19.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20181030221726-6c7e314b6563/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= @@ -78,4 +84,8 @@ gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10= gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gorm.io/driver/sqlite v1.5.6 h1:fO/X46qn5NUEEOZtnjJRWRzZMe8nqJiQ9E+0hi+hKQE= +gorm.io/driver/sqlite v1.5.6/go.mod h1:U+J8craQU6Fzkcvu8oLeAQmi50TkwPEhHDEjQZXDah4= +gorm.io/gorm v1.25.12 h1:I0u8i2hWQItBq1WfE0o2+WuL9+8L21K9e2HHSTE/0f8= +gorm.io/gorm v1.25.12/go.mod h1:xh7N7RHfYlNc5EmcI/El95gXusucDrQnHXe0+CgWcLQ= sigs.k8s.io/yaml v1.2.0/go.mod h1:yfXDCHCao9+ENCvLSE62v9VSji2MKu5jeNfTrofGhJc= diff --git a/internal/routine/extract/extract.go b/internal/routine/extract/extract.go new file mode 100644 index 0000000..3cbb347 --- /dev/null +++ b/internal/routine/extract/extract.go @@ -0,0 +1,169 @@ +package extract + +import ( + "bytes" + "fmt" + "image/png" + "log" + "slices" + "strings" + + "github.com/nandesh-dev/subtle/pkgs/ass" + "github.com/nandesh-dev/subtle/pkgs/config" + "github.com/nandesh-dev/subtle/pkgs/db" + "github.com/nandesh-dev/subtle/pkgs/filemanager" + "github.com/nandesh-dev/subtle/pkgs/pgs" + "github.com/nandesh-dev/subtle/pkgs/subtitle" + "github.com/nandesh-dev/subtle/pkgs/warning" + "golang.org/x/text/language" +) + +func Run() warning.WarningList { + warnings := warning.NewWarningList() + for _, rootDirectoryConfig := range config.Config().Media.RootDirectories { + dir, _, err := filemanager.ReadDirectory(rootDirectoryConfig.Path) + if err != nil { + warnings.AddWarning(fmt.Errorf("Error reading root directory: %v; %v", rootDirectoryConfig.Path, err)) + continue + } + + warns := extractSubtitleFromDirectory(*dir, rootDirectoryConfig.AutoExtract) + warnings.Append(warns) + } + + return *warnings +} + +func extractSubtitleFromDirectory(dir filemanager.Directory, autoExtractConfig config.AutoExtract) warning.WarningList { + warnings := warning.NewWarningList() + + for _, video := range dir.VideoFiles() { + var videoEntry db.Video + if err := db.DB().Where(&db.Video{DirectoryPath: video.DirectoryPath(), Filename: video.Filename()}). + Preload("Subtitles"). + Preload("Subtitles.Segments"). + First(&videoEntry).Error; err != nil { + log.Fatal("Error getting entry: ", err, video.DirectoryPath(), video.Filename()) + } + + rawStreams, err := video.RawStreams() + if err != nil { + warnings.AddWarning(fmt.Errorf("Error extracting raw stream from video: %v; %v", video.Filepath(), err)) + continue + } + + rawStreamRanks := map[filemanager.RawStream]int{} + + for _, rawStream := range *rawStreams { + if !slices.ContainsFunc(videoEntry.Subtitles, func(subtitleEntry db.Subtitle) bool { + subtitleEntryLanguageTag, _ := language.Parse(subtitleEntry.Language) + return subtitleEntryLanguageTag == rawStream.Language() + }) && slices.ContainsFunc(autoExtractConfig.Languages, func(lang language.Tag) bool { + return lang == rawStream.Language() + }) { + rawStreamRanks[rawStream] = 1 + + for _, titleKeyword := range autoExtractConfig.RawStreamTitleKeywords { + if strings.Contains(rawStream.Title(), titleKeyword) { + rawStreamRanks[rawStream]++ + } + } + } + } + + highestRank := 0 + var highestRankRawStream filemanager.RawStream + for rawStream, rank := range rawStreamRanks { + if rank >= highestRank { + highestRank = rank + highestRankRawStream = rawStream + } + } + + if highestRank == 0 { + continue + } + + rawStream := highestRankRawStream + var sub subtitle.Subtitle + + switch rawStream.Format() { + case subtitle.ASS: + s, warns, err := ass.DecodeSubtitle(rawStream) + warnings.Append(warns) + if err != nil { + warnings.AddWarning(fmt.Errorf("Error decoding subtitle for video: %v; %v", video.Filepath(), err)) + continue + } + + sub = *s + + case subtitle.PGS: + s, warns, err := pgs.DecodeSubtitle(rawStream) + warnings.Append(warns) + if err != nil { + warnings.AddWarning(fmt.Errorf("Error decoding subtitle for video: %v; %v", video.Filepath(), err)) + continue + } + + sub = *s + } + + if sub == nil { + continue + } + + switch sub := sub.(type) { + case subtitle.TextSubtitle: + subtitleEntry := db.Subtitle{ + Language: rawStream.Language().String(), + Filepath: "", + IsImage: false, + Segments: make([]db.Segment, 0), + } + + for _, segment := range sub.Segments() { + segmentEntry := db.Segment{ + StartTime: segment.Start(), + EndTime: segment.End(), + Text: segment.Text(), + } + + subtitleEntry.Segments = append(subtitleEntry.Segments, segmentEntry) + } + + videoEntry.Subtitles = append(videoEntry.Subtitles, subtitleEntry) + + db.DB().Save(&videoEntry) + case subtitle.ImageSubtitle: + subtitleEntry := db.Subtitle{ + Language: rawStream.Language().String(), + Filepath: "", + IsImage: true, + Segments: make([]db.Segment, 0), + } + + for _, segment := range sub.Segments() { + imageDataBuffer := new(bytes.Buffer) + if err := png.Encode(imageDataBuffer, segment.Image()); err != nil { + warnings.AddWarning(fmt.Errorf("Error encoding image to png for video: %v; %v", video.Filepath(), err)) + continue + } + + segmentEntry := db.Segment{ + StartTime: segment.Start(), + EndTime: segment.End(), + ImageData: imageDataBuffer.Bytes(), + } + + subtitleEntry.Segments = append(subtitleEntry.Segments, segmentEntry) + } + + videoEntry.Subtitles = append(videoEntry.Subtitles, subtitleEntry) + + db.DB().Save(&videoEntry) + } + } + + return *warnings +} diff --git a/internal/routine/library/library.go b/internal/routine/library/library.go deleted file mode 100644 index 60d7648..0000000 --- a/internal/routine/library/library.go +++ /dev/null @@ -1,92 +0,0 @@ -package library - -import ( - "fmt" - "os" - "path/filepath" - "slices" - - "github.com/nandesh-dev/subtle/pkgs/config" - "github.com/nandesh-dev/subtle/pkgs/filemanager" - "github.com/nandesh-dev/subtle/pkgs/subtitle" - "github.com/nandesh-dev/subtle/pkgs/warning" - "golang.org/x/text/language" -) - -func RunLibraryRoutine() { - warnings := warning.NewWarningList() - - for _, watchDirectoryConfig := range config.Config().Media.WatchDirectories { - dir, _, err := filemanager.ReadDirectory(watchDirectoryConfig.Path) - if err != nil { - return - } - - for _, videoFile := range dir.VideoFiles() { - missingSubtitleLanguages := make([]language.Tag, 0) - for _, languageCode := range watchDirectoryConfig.AutoExtract.Languages { - languageTag, err := language.Parse(languageCode) - if err != nil { - warnings.AddWarning(fmt.Errorf("Invalid language code in config: %v; %v", languageCode, err)) - continue - } - - hasSubtitleLanguage, wrn := videoFile.HasSubtitleLanguage(languageTag) - warnings.Append(wrn) - - if !hasSubtitleLanguage { - missingSubtitleLanguages = append(missingSubtitleLanguages, languageTag) - } - } - - if len(missingSubtitleLanguages) == 0 { - continue - } - - rawStreams, err := subtitle.ExtractRawStreams(videoFile.Filepath()) - if err != nil { - warnings.AddWarning(fmt.Errorf("Error extracting raw streams from video file: %v", err)) - } - - for _, formatCode := range watchDirectoryConfig.AutoExtract.Formats { - format, err := subtitle.ParseFormat(formatCode) - if err != nil { - warnings.AddWarning(fmt.Errorf("Invalid subtitle format in config: %v", formatCode)) - continue - } - - for _, rawStream := range rawStreams { - if format == rawStream.Format() && slices.Contains(missingSubtitleLanguages, rawStream.Language()) { - sub, wrn, err := subtitle.FromRawStream(rawStream) - warnings.Append(wrn) - if err != nil { - warnings.AddWarning(fmt.Errorf("Error extracting subtitle: %v", err)) - return - } - - outputFormat, err := subtitle.ParseFormat(watchDirectoryConfig.AutoExtract.OutputFormat) - if err != nil { - warnings.AddWarning(fmt.Errorf("Invalid output format in config: %v", err)) - return - } - - encodedSubtitle, _ := sub.Encode(outputFormat) - - for _, file := range encodedSubtitle.Files() { - path := filepath.Join(videoFile.DirectoryPath(), videoFile.Basename()+"."+rawStream.Language().String()) + file.Extension() - - if err := os.WriteFile(path, file.Content(), 0644); err != nil { - fmt.Print(err) - } - } - } - } - } - } - } - - fmt.Println("WARNINGS BEGIN") - for _, warning := range warnings.Warnings() { - fmt.Println(warning) - } -} diff --git a/internal/routine/media/media.go b/internal/routine/media/media.go new file mode 100644 index 0000000..77cd437 --- /dev/null +++ b/internal/routine/media/media.go @@ -0,0 +1,43 @@ +package media + +import ( + "fmt" + + "github.com/nandesh-dev/subtle/pkgs/config" + "github.com/nandesh-dev/subtle/pkgs/db" + "github.com/nandesh-dev/subtle/pkgs/filemanager" +) + +func Run() error { + for _, rootDirectoryConfig := range config.Config().Media.RootDirectories { + dir, _, err := filemanager.ReadDirectory(rootDirectoryConfig.Path) + if err != nil { + return fmt.Errorf("Error reading root directory: %v", err) + } + + if err := syncDirectoryVideos(dir); err != nil { + return fmt.Errorf("Error syncing directory: %v", err) + } + } + + return nil +} + +func syncDirectoryVideos(dir *filemanager.Directory) error { + for _, video := range dir.VideoFiles() { + if err := db.DB().Where(db.Video{DirectoryPath: video.DirectoryPath(), Filename: video.Filename()}).FirstOrCreate(&db.Video{}, db.Video{ + DirectoryPath: video.DirectoryPath(), + Filename: video.Filename(), + }).Error; err != nil { + return fmt.Errorf("Error creating video entry: %v", err) + } + } + + for _, dir := range dir.Children() { + if err := syncDirectoryVideos(&dir); err != nil { + return err + } + } + + return nil +} diff --git a/pkgs/subtitle/ass/decoder.go b/pkgs/ass/decoder.go similarity index 81% rename from pkgs/subtitle/ass/decoder.go rename to pkgs/ass/decoder.go index c728001..7492cea 100644 --- a/pkgs/subtitle/ass/decoder.go +++ b/pkgs/ass/decoder.go @@ -7,6 +7,8 @@ import ( "strings" "time" + "github.com/nandesh-dev/subtle/pkgs/filemanager" + "github.com/nandesh-dev/subtle/pkgs/subtitle" "github.com/nandesh-dev/subtle/pkgs/warning" ffmpeg "github.com/u2takey/ffmpeg-go" ) @@ -28,21 +30,20 @@ const ( Format ) -func DecodeSubtitle(path string, index int) (Stream, *warning.WarningList, error) { - stream := NewStream() +func DecodeSubtitle(rawStream filemanager.RawStream) (*subtitle.TextSubtitle, warning.WarningList, error) { warnings := warning.NewWarningList() var subtitleBuf, errorBuf bytes.Buffer ffmpeg.LogCompiledCommand = false - err := ffmpeg.Input(path). - Output("pipe:", ffmpeg.KwArgs{"map": fmt.Sprintf("0:%v", index), "f": "ass"}). + err := ffmpeg.Input(rawStream.Filepath()). + Output("pipe:", ffmpeg.KwArgs{"map": fmt.Sprintf("0:%v", rawStream.Index()), "f": "ass"}). WithOutput(&subtitleBuf). WithErrorOutput(&errorBuf). Run() if err != nil { - return *stream, warnings, fmt.Errorf("Error extracting subtitles: %v %v", err, errorBuf.String()) + return nil, *warnings, fmt.Errorf("Error extracting subtitles: %v %v", err, errorBuf.String()) } reader := NewReader(subtitleBuf.Bytes()) @@ -51,6 +52,8 @@ func DecodeSubtitle(path string, index int) (Stream, *warning.WarningList, error timeMultiplier := 1 timeOffset := time.Duration(0) + sub := subtitle.NewTextSubtitle() + for !reader.ReachedEnd() { line, _ := reader.Advance() @@ -68,32 +71,36 @@ func DecodeSubtitle(path string, index int) (Stream, *warning.WarningList, error currentFormat = append(currentFormat, strings.TrimSpace(pt)) } case Dialogue: - segment := NewSegment() + start, end, text := time.Second*0, time.Second*0, "" parts := strings.SplitN(suffix, ",", len(currentFormat)) for i, partName := range currentFormat { switch partName { case "Start": - start, err := parseTime(parts[i], timeMultiplier, timeOffset) + st, err := parseTime(parts[i], timeMultiplier, timeOffset) if err != nil { warnings.AddWarning(fmt.Errorf("Error parsing start timestamp: %v; %v", err, line)) } else { - segment.SetStart(start) + start = st } case "End": - end, err := parseTime(parts[i], timeMultiplier, timeOffset) + ed, err := parseTime(parts[i], timeMultiplier, timeOffset) if err != nil { warnings.AddWarning(fmt.Errorf("Error parsing end timestamp: %v; %v", err, line)) } else { - segment.SetEnd(end) + end = ed } case "Text": - segment.SetText(extractText(parts[i])) + text = extractText(parts[i]) } } - stream.AddSegment(*segment) + if text != "" { + segment := subtitle.NewTextSegment(start, end, text) + sub.AddSegment(*segment) + } + case Timer: multiplier, err := strconv.ParseFloat(suffix, 32) if err != nil { @@ -113,7 +120,7 @@ func DecodeSubtitle(path string, index int) (Stream, *warning.WarningList, error } } - return *stream, warnings, nil + return sub, *warnings, nil } func extractLineTypePrefix(line string) (LineType, string, error) { diff --git a/pkgs/subtitle/ass/reader.go b/pkgs/ass/reader.go similarity index 100% rename from pkgs/subtitle/ass/reader.go rename to pkgs/ass/reader.go diff --git a/pkgs/config/config.go b/pkgs/config/config.go index 254da2a..8ad1a0d 100644 --- a/pkgs/config/config.go +++ b/pkgs/config/config.go @@ -6,27 +6,37 @@ import ( "path/filepath" "sync" + "golang.org/x/text/language" "gopkg.in/yaml.v3" ) type Server struct { - Port int - GRPCReflection bool + Port int + GRPCReflection bool + DatabaseDirectory string +} + +type AutoExtractFormat struct { + ASS AutoExtractASS +} + +type AutoExtractASS struct { + Enabled bool } type AutoExtract struct { - Formats []string - Languages []string - OutputFormat string + Languages []language.Tag + Formats AutoExtractFormat + RawStreamTitleKeywords []string } -type WatchDirectory struct { +type RootDirectory struct { Path string AutoExtract AutoExtract } type Media struct { - WatchDirectories []WatchDirectory + RootDirectories []RootDirectory } type t struct { @@ -48,17 +58,22 @@ func Init(basepath string) (e error) { once.Do(func() { config = t{ Server: Server{ - Port: 3000, - GRPCReflection: false, + Port: 3000, + GRPCReflection: false, + DatabaseDirectory: filepath.Join(basepath, "db"), }, Media: Media{ - WatchDirectories: []WatchDirectory{ + RootDirectories: []RootDirectory{ { Path: "/media", AutoExtract: AutoExtract{ - Languages: []string{"eng"}, - Formats: []string{"ass"}, - OutputFormat: "srt", + Languages: []language.Tag{language.English}, + Formats: AutoExtractFormat{ + ASS: AutoExtractASS{ + Enabled: true, + }, + }, + RawStreamTitleKeywords: []string{"Full", "Dialogue"}, }, }, }, diff --git a/pkgs/db/db.go b/pkgs/db/db.go new file mode 100644 index 0000000..3bbcb73 --- /dev/null +++ b/pkgs/db/db.go @@ -0,0 +1,59 @@ +package db + +import ( + "fmt" + "time" + + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +var ( + database *gorm.DB +) + +type Video struct { + ID int `gorm:"primaryKey"` + + DirectoryPath string + Filename string + Subtitles []Subtitle `gorm:"foreignKey:VideoID"` +} + +type Subtitle struct { + ID int `gorm:"primaryKey"` + VideoID int + + Language string + Filepath string + IsImage bool + Segments []Segment `gorm:"foreignKey:SubtitleID"` +} + +type Segment struct { + ID int `gorm:"primaryKey"` + SubtitleID int + + StartTime time.Duration + EndTime time.Duration + ImageData []byte + Text string +} + +func DB() *gorm.DB { + return database +} + +func Init() error { + db, err := gorm.Open(sqlite.Open("config/database.db"), &gorm.Config{}) + if err != nil { + return fmt.Errorf("Error opening database: %v", err) + } + + if err = db.AutoMigrate(&Video{}, &Subtitle{}, &Segment{}); err != nil { + return fmt.Errorf("Error auto migrating database: %v", err) + } + + database = db + return nil +} diff --git a/pkgs/filemanager/directory.go b/pkgs/filemanager/directory.go index 8a305df..6d4060c 100644 --- a/pkgs/filemanager/directory.go +++ b/pkgs/filemanager/directory.go @@ -1,41 +1,24 @@ package filemanager type Directory struct { - path string - children []Directory - videos []VideoFile - extraSubtitles []SubtitleFile -} - -func NewDirectory(path string) *Directory { - return &Directory{ - path: path, - children: make([]Directory, 0), - videos: make([]VideoFile, 0), - extraSubtitles: make([]SubtitleFile, 0), - } -} - -func (d *Directory) AddVideoFile(file VideoFile) { - d.videos = append(d.videos, file) -} - -func (d *Directory) AddExtraSubtitleFile(file SubtitleFile) { - d.extraSubtitles = append(d.extraSubtitles, file) -} - -func (d *Directory) AddChild(child Directory) { - d.children = append(d.children, child) + path string + children []Directory + videos []VideoFile + subtitles []SubtitleFile } func (d *Directory) VideoFiles() []VideoFile { return d.videos } -func (d *Directory) ExtraSubtitleFiles() []SubtitleFile { - return d.extraSubtitles +func (d *Directory) SubtitleFiles() []SubtitleFile { + return d.subtitles } func (d *Directory) Children() []Directory { return d.children } + +func (d *Directory) Path() string { + return d.path +} diff --git a/pkgs/filemanager/file.go b/pkgs/filemanager/file_type.go similarity index 100% rename from pkgs/filemanager/file.go rename to pkgs/filemanager/file_type.go diff --git a/pkgs/subtitle/raw_stream.go b/pkgs/filemanager/raw_stream.go similarity index 71% rename from pkgs/subtitle/raw_stream.go rename to pkgs/filemanager/raw_stream.go index 50a4f6e..077bd75 100644 --- a/pkgs/subtitle/raw_stream.go +++ b/pkgs/filemanager/raw_stream.go @@ -1,9 +1,11 @@ -package subtitle +package filemanager import ( "encoding/json" "fmt" + "strings" + "github.com/nandesh-dev/subtle/pkgs/subtitle" "github.com/nandesh-dev/subtle/pkgs/warning" ffmpeg "github.com/u2takey/ffmpeg-go" @@ -13,46 +15,33 @@ import ( type RawStream struct { filepath string index int - format Format + format subtitle.Format language language.Tag -} - -func NewRawStream(filepath string) *RawStream { - return &RawStream{ - filepath: filepath, - } + title string } func (s *RawStream) Filepath() string { return s.filepath } -func (s *RawStream) SetIndex(index int) { - s.index = index -} - func (s *RawStream) Index() int { return s.index } -func (s *RawStream) SetFormat(format Format) { - s.format = format -} - -func (s *RawStream) Format() Format { +func (s *RawStream) Format() subtitle.Format { return s.format } -func (s *RawStream) SetLanguage(lang language.Tag) { - s.language = lang -} - func (s *RawStream) Language() language.Tag { return s.language } -func ExtractRawStreams(path string) ([]RawStream, error) { - rawResultData, err := ffmpeg.Probe(path) +func (s *RawStream) Title() string { + return s.title +} + +func (v *VideoFile) RawStreams() (*[]RawStream, error) { + rawResultData, err := ffmpeg.Probe(v.path) warnings := warning.NewWarningList() if err != nil { @@ -83,8 +72,6 @@ func ExtractRawStreams(path string) ([]RawStream, error) { continue } - rawStream := NewRawStream(path) - codecName, codecNameExist := rawStreamData.(map[string]any)["codec_name"].(string) if !codecNameExist { return nil, fmt.Errorf("Missing codec name in probe JSON: %v", rawStreamData) @@ -94,16 +81,14 @@ func ExtractRawStreams(path string) ([]RawStream, error) { if err != nil { return nil, err } - rawStream.SetFormat(format) rawIndex, indexExist := rawStreamData.(map[string]any)["index"].(float64) if !indexExist { return nil, fmt.Errorf("Missing index in probe JSON: %v", rawStreamData) } - rawStream.SetIndex(int(rawIndex)) - lang := language.English + title := "" tags, tagsExist := rawStreamData.(map[string]any)["tags"].(map[string]any) if tagsExist { @@ -117,23 +102,33 @@ func ExtractRawStreams(path string) ([]RawStream, error) { lang = langTag } } - } - rawStream.SetLanguage(lang) + if rawTitle, titleExist := tags["title"].(string); titleExist { + title = strings.TrimSpace(rawTitle) + } + } - rawStreams = append(rawStreams, *rawStream) + rawStreams = append(rawStreams, RawStream{ + filepath: v.path, + index: int(rawIndex), + format: format, + language: lang, + title: title, + }) } - return rawStreams, nil + return &rawStreams, nil } -func mapCodecName(cN string) (Format, error) { +func mapCodecName(cN string) (subtitle.Format, error) { switch cN { case "hdmv_pgs_subtitle": - return PGS, nil + return subtitle.PGS, nil case "ass": - return ASS, nil + return subtitle.ASS, nil + case "subrip": + return subtitle.SRT, nil } - return ASS, fmt.Errorf("Unsupported or invalid codec name: %v", cN) + return subtitle.ASS, fmt.Errorf("Unsupported or invalid codec name: %v", cN) } diff --git a/pkgs/filemanager/read_dir.go b/pkgs/filemanager/read_dir.go index 892151f..96014ff 100644 --- a/pkgs/filemanager/read_dir.go +++ b/pkgs/filemanager/read_dir.go @@ -15,9 +15,12 @@ func ReadDirectory(path string) (*Directory, warning.WarningList, error) { return nil, *warnings, err } - directory := NewDirectory(path) - videos := make([]*VideoFile, 0) - subtitles := make([]*SubtitleFile, 0) + directory := &Directory{ + path: path, + children: make([]Directory, 0), + videos: make([]VideoFile, 0), + subtitles: make([]SubtitleFile, 0), + } for _, entry := range files { entrypath := filepath.Join(path, entry.Name()) @@ -29,11 +32,13 @@ func ReadDirectory(path string) (*Directory, warning.WarningList, error) { return nil, *warnings, err } - directory.AddChild(*child) + directory.children = append(directory.children, *child) } if IsSubtitleFile(entrypath) { - subtitles = append(subtitles, NewSubtitleFile(entrypath)) + directory.subtitles = append(directory.subtitles, SubtitleFile{ + path: entrypath, + }) continue } @@ -45,33 +50,10 @@ func ReadDirectory(path string) (*Directory, warning.WarningList, error) { } if isVideoFile { - videos = append(videos, NewVideoFile(entrypath)) - } - } - - extraSubtitles := make([]*SubtitleFile, 0) - - for _, subtitle := range subtitles { - found := false - for _, video := range videos { - if subtitle.Basename() == video.Basename() { - found = true - video.AddSubtitleFile(*subtitle) - break - } + directory.videos = append(directory.videos, VideoFile{ + path: entrypath, + }) } - - if !found { - extraSubtitles = append(extraSubtitles, subtitle) - } - } - - for _, video := range videos { - directory.AddVideoFile(*video) - } - - for _, subtitle := range extraSubtitles { - directory.AddExtraSubtitleFile(*subtitle) } return directory, *warnings, nil diff --git a/pkgs/filemanager/subtitle.go b/pkgs/filemanager/subtitle_file.go similarity index 72% rename from pkgs/filemanager/subtitle.go rename to pkgs/filemanager/subtitle_file.go index e7c5338..16fd222 100644 --- a/pkgs/filemanager/subtitle.go +++ b/pkgs/filemanager/subtitle_file.go @@ -1,20 +1,17 @@ package filemanager import ( + "fmt" "path/filepath" "strings" + + "golang.org/x/text/language" ) type SubtitleFile struct { path string } -func NewSubtitleFile(path string) *SubtitleFile { - return &SubtitleFile{ - path: path, - } -} - func (s *SubtitleFile) Path() string { return s.path } @@ -29,14 +26,14 @@ func (s *SubtitleFile) Basename() string { return pt[0] } -func (s *SubtitleFile) LanguageCode() string { +func (s *SubtitleFile) Language() (language.Tag, error) { pt := strings.Split(filepath.Base(s.path), ".") if len(pt) >= 3 { - return pt[len(pt)-2] + return language.Parse(pt[len(pt)-2]) } - return "" + return language.English, fmt.Errorf("No language code found") } func (s *SubtitleFile) Extension() string { diff --git a/pkgs/filemanager/video.go b/pkgs/filemanager/video.go deleted file mode 100644 index 8207d6b..0000000 --- a/pkgs/filemanager/video.go +++ /dev/null @@ -1,62 +0,0 @@ -package filemanager - -import ( - "fmt" - "path/filepath" - "strings" - - "github.com/nandesh-dev/subtle/pkgs/warning" - "golang.org/x/text/language" -) - -type VideoFile struct { - path string - subtitles []SubtitleFile -} - -func NewVideoFile(path string) *VideoFile { - return &VideoFile{ - path: path, - } -} - -func (v *VideoFile) DirectoryPath() string { - return filepath.Dir(v.path) -} - -func (v *VideoFile) Filepath() string { - return v.path -} - -func (v *VideoFile) Extension() string { - return filepath.Ext(v.path) -} - -func (v *VideoFile) Basename() string { - return strings.TrimSuffix(filepath.Base(v.path), v.Extension()) -} - -func (v *VideoFile) HasSubtitleLanguage(tag language.Tag) (bool, warning.WarningList) { - warnings := warning.NewWarningList() - for _, subtitleFile := range v.SubtitleFiles() { - subtitleLanguageTag, err := language.Parse(subtitleFile.LanguageCode()) - if err != nil { - warnings.AddWarning(fmt.Errorf("Invalid language code: %v; %v", subtitleFile.LanguageCode(), err)) - continue - } - - if subtitleLanguageTag == tag { - return true, *warnings - } - } - - return false, *warnings -} - -func (v *VideoFile) AddSubtitleFile(file SubtitleFile) { - v.subtitles = append(v.subtitles, file) -} - -func (v *VideoFile) SubtitleFiles() []SubtitleFile { - return v.subtitles -} diff --git a/pkgs/filemanager/video_file.go b/pkgs/filemanager/video_file.go new file mode 100644 index 0000000..ffbf8d3 --- /dev/null +++ b/pkgs/filemanager/video_file.go @@ -0,0 +1,30 @@ +package filemanager + +import ( + "path/filepath" + "strings" +) + +type VideoFile struct { + path string +} + +func (v *VideoFile) DirectoryPath() string { + return filepath.Dir(v.path) +} + +func (v *VideoFile) Filepath() string { + return v.path +} + +func (v *VideoFile) Filename() string { + return filepath.Base(v.path) +} + +func (v *VideoFile) Extension() string { + return filepath.Ext(v.path) +} + +func (v *VideoFile) Basename() string { + return strings.TrimSuffix(filepath.Base(v.path), v.Extension()) +} diff --git a/pkgs/subtitle/pgs/decoder.go b/pkgs/pgs/decoder.go similarity index 77% rename from pkgs/subtitle/pgs/decoder.go rename to pkgs/pgs/decoder.go index bb28369..3d2fc54 100644 --- a/pkgs/subtitle/pgs/decoder.go +++ b/pkgs/pgs/decoder.go @@ -3,26 +3,28 @@ package pgs import ( "bytes" "fmt" + "time" + "github.com/nandesh-dev/subtle/pkgs/filemanager" + "github.com/nandesh-dev/subtle/pkgs/subtitle" "github.com/nandesh-dev/subtle/pkgs/warning" ffmpeg "github.com/u2takey/ffmpeg-go" ) -func DecodeSubtitle(path string, index int) (Stream, *warning.WarningList, error) { - stream := NewStream() +func DecodeSubtitle(rawStream filemanager.RawStream) (*subtitle.ImageSubtitle, warning.WarningList, error) { warnings := warning.NewWarningList() var subtitleBuf, errorBuf bytes.Buffer ffmpeg.LogCompiledCommand = false - err := ffmpeg.Input(path). - Output("pipe:", ffmpeg.KwArgs{"map": fmt.Sprintf("0:%v", index), "c:s": "copy", "f": "sup"}). + err := ffmpeg.Input(rawStream.Filepath()). + Output("pipe:", ffmpeg.KwArgs{"map": fmt.Sprintf("0:%v", rawStream.Index()), "c:s": "copy", "f": "sup"}). WithOutput(&subtitleBuf). WithErrorOutput(&errorBuf). Run() if err != nil { - return *stream, warnings, fmt.Errorf("Error extracting subtitles: %v %v", err, errorBuf) + return nil, *warnings, fmt.Errorf("Error extracting subtitles: %v %v", err, errorBuf) } reader := NewReader(subtitleBuf.Bytes()) @@ -34,7 +36,7 @@ func DecodeSubtitle(path string, index int) (Stream, *warning.WarningList, error header, err := ReadHeader(reader) if err != nil { - return *stream, warnings, fmt.Errorf("Error reading header: %v", err) + return nil, *warnings, fmt.Errorf("Error reading header: %v", err) } dS.Header = *header @@ -95,6 +97,8 @@ func DecodeSubtitle(path string, index int) (Stream, *warning.WarningList, error } } + sub := subtitle.NewImageSubtitle() + for _, displaySet := range displaySets { images, err := displaySet.parse() if err != nil { @@ -102,13 +106,11 @@ func DecodeSubtitle(path string, index int) (Stream, *warning.WarningList, error continue } - segment := NewSegment() - - segment.AddImages(images) - segment.SetStart(displaySet.Header.PTS) - - stream.AddSegment(*segment) + for _, image := range images { + segment := subtitle.NewImageSegment(displaySet.Header.PTS, time.Second*0, image) + sub.AddSegment(*segment) + } } - return *stream, warnings, nil + return sub, *warnings, nil } diff --git a/pkgs/subtitle/pgs/display_set.go b/pkgs/pgs/display_set.go similarity index 100% rename from pkgs/subtitle/pgs/display_set.go rename to pkgs/pgs/display_set.go diff --git a/pkgs/subtitle/pgs/header.go b/pkgs/pgs/header.go similarity index 100% rename from pkgs/subtitle/pgs/header.go rename to pkgs/pgs/header.go diff --git a/pkgs/subtitle/pgs/ods.go b/pkgs/pgs/ods.go similarity index 100% rename from pkgs/subtitle/pgs/ods.go rename to pkgs/pgs/ods.go diff --git a/pkgs/subtitle/pgs/pcs.go b/pkgs/pgs/pcs.go similarity index 100% rename from pkgs/subtitle/pgs/pcs.go rename to pkgs/pgs/pcs.go diff --git a/pkgs/subtitle/pgs/pds.go b/pkgs/pgs/pds.go similarity index 100% rename from pkgs/subtitle/pgs/pds.go rename to pkgs/pgs/pds.go diff --git a/pkgs/subtitle/pgs/reader.go b/pkgs/pgs/reader.go similarity index 100% rename from pkgs/subtitle/pgs/reader.go rename to pkgs/pgs/reader.go diff --git a/pkgs/subtitle/pgs/segments.go b/pkgs/pgs/segments.go similarity index 100% rename from pkgs/subtitle/pgs/segments.go rename to pkgs/pgs/segments.go diff --git a/pkgs/subtitle/pgs/wds.go b/pkgs/pgs/wds.go similarity index 100% rename from pkgs/subtitle/pgs/wds.go rename to pkgs/pgs/wds.go diff --git a/pkgs/subtitle/srt/encoder.go b/pkgs/srt/encoder.go similarity index 81% rename from pkgs/subtitle/srt/encoder.go rename to pkgs/srt/encoder.go index 55cb33a..ae06c97 100644 --- a/pkgs/subtitle/srt/encoder.go +++ b/pkgs/srt/encoder.go @@ -3,10 +3,12 @@ package srt import ( "fmt" "time" + + "github.com/nandesh-dev/subtle/pkgs/subtitle" ) -func EncodeSubtitle(stream Stream) string { - segments := stream.Segments() +func EncodeSubtitle(sub subtitle.TextSubtitle) string { + segments := sub.Segments() output := "" diff --git a/pkgs/subtitle/ass/ass.go b/pkgs/subtitle/ass/ass.go deleted file mode 100644 index 4ff6cc0..0000000 --- a/pkgs/subtitle/ass/ass.go +++ /dev/null @@ -1,55 +0,0 @@ -package ass - -import "time" - -type Segment struct { - start time.Duration - end time.Duration - text string -} - -func NewSegment() *Segment { - return &Segment{} -} - -func (s *Segment) Text() string { - return s.text -} - -func (s *Segment) SetText(text string) { - s.text = text -} - -func (s *Segment) Start() time.Duration { - return s.start -} - -func (s *Segment) SetStart(start time.Duration) { - s.start = start -} - -func (s *Segment) End() time.Duration { - return s.end -} - -func (s *Segment) SetEnd(end time.Duration) { - s.end = end -} - -type Stream struct { - segments []Segment -} - -func NewStream() *Stream { - return &Stream{ - segments: make([]Segment, 0), - } -} - -func (s *Stream) Segments() []Segment { - return s.segments -} - -func (s *Stream) AddSegment(segment Segment) { - s.segments = append(s.segments, segment) -} diff --git a/pkgs/subtitle/format.go b/pkgs/subtitle/format.go new file mode 100644 index 0000000..e65d49d --- /dev/null +++ b/pkgs/subtitle/format.go @@ -0,0 +1,24 @@ +package subtitle + +import "fmt" + +type Format int + +const ( + PGS Format = iota + ASS + SRT +) + +func ParseFormat(f string) (Format, error) { + switch f { + case "ass": + return ASS, nil + case "pgs": + return PGS, nil + case "srt": + return SRT, nil + } + + return ASS, fmt.Errorf("Invalid format: %v", f) +} diff --git a/pkgs/subtitle/image.go b/pkgs/subtitle/image.go new file mode 100644 index 0000000..8b3655a --- /dev/null +++ b/pkgs/subtitle/image.go @@ -0,0 +1,50 @@ +package subtitle + +import ( + "image" + "time" +) + +type ImageSubtitle struct { + segments []ImageSegment +} + +type ImageSegment struct { + start time.Duration + end time.Duration + image image.Image +} + +func NewImageSubtitle() *ImageSubtitle { + return &ImageSubtitle{ + segments: make([]ImageSegment, 0), + } +} + +func (s *ImageSubtitle) AddSegment(segment ImageSegment) { + s.segments = append(s.segments, segment) +} + +func (s *ImageSubtitle) Segments() []ImageSegment { + return s.segments +} + +func NewImageSegment(start time.Duration, end time.Duration, image image.Image) *ImageSegment { + return &ImageSegment{ + start: start, + end: end, + image: image, + } +} + +func (s *ImageSegment) Start() time.Duration { + return s.start +} + +func (s *ImageSegment) End() time.Duration { + return s.end +} + +func (s *ImageSegment) Image() image.Image { + return s.image +} diff --git a/pkgs/subtitle/pgs/pgs.go b/pkgs/subtitle/pgs/pgs.go deleted file mode 100644 index 28e096a..0000000 --- a/pkgs/subtitle/pgs/pgs.go +++ /dev/null @@ -1,50 +0,0 @@ -package pgs - -import ( - "image" - "time" -) - -type Segment struct { - start time.Duration - images []image.Image -} - -func NewSegment() *Segment { - return &Segment{} -} - -func (s *Segment) Images() ([]image.Image, error) { - return s.images, nil -} - -func (s *Segment) AddImages(images []image.Image) error { - s.images = append(s.images, images...) - return nil -} - -func (s *Segment) Start() time.Duration { - return s.start -} - -func (s *Segment) SetStart(start time.Duration) { - s.start = start -} - -type Stream struct { - segments []Segment -} - -func NewStream() *Stream { - return &Stream{ - segments: make([]Segment, 0), - } -} - -func (s *Stream) AddSegment(segment Segment) { - s.segments = append(s.segments, segment) -} - -func (s *Stream) Segments() []Segment { - return s.segments -} diff --git a/pkgs/subtitle/srt/srt.go b/pkgs/subtitle/srt/srt.go deleted file mode 100644 index 622c873..0000000 --- a/pkgs/subtitle/srt/srt.go +++ /dev/null @@ -1,55 +0,0 @@ -package srt - -import "time" - -type Segment struct { - start time.Duration - end time.Duration - text string -} - -func NewSegment() *Segment { - return &Segment{} -} - -func (s *Segment) Text() string { - return s.text -} - -func (s *Segment) SetText(text string) { - s.text = text -} - -func (s *Segment) Start() time.Duration { - return s.start -} - -func (s *Segment) SetStart(start time.Duration) { - s.start = start -} - -func (s *Segment) End() time.Duration { - return s.end -} - -func (s *Segment) SetEnd(end time.Duration) { - s.end = end -} - -type Stream struct { - segments []Segment -} - -func NewStream() *Stream { - return &Stream{ - segments: make([]Segment, 0), - } -} - -func (s *Stream) AddSegment(segment Segment) { - s.segments = append(s.segments, segment) -} - -func (s *Stream) Segments() []Segment { - return s.segments -} diff --git a/pkgs/subtitle/subtitle.go b/pkgs/subtitle/subtitle.go index b269bfc..7413180 100644 --- a/pkgs/subtitle/subtitle.go +++ b/pkgs/subtitle/subtitle.go @@ -1,216 +1,3 @@ package subtitle -import ( - "fmt" - "time" - - "github.com/nandesh-dev/subtle/pkgs/subtitle/ass" - "github.com/nandesh-dev/subtle/pkgs/subtitle/pgs" - "github.com/nandesh-dev/subtle/pkgs/subtitle/srt" - "github.com/nandesh-dev/subtle/pkgs/tesseract" - "github.com/nandesh-dev/subtle/pkgs/warning" -) - -type Format int - -const ( - PGS Format = iota - ASS - SRT -) - -func ParseFormat(f string) (Format, error) { - switch f { - case "ass": - return ASS, nil - case "pgs": - return PGS, nil - case "srt": - return SRT, nil - } - - return ASS, fmt.Errorf("Invalid format: %v", f) -} - -type Segment struct { - start time.Duration - end time.Duration - text string -} - -func NewSegment() *Segment { - return &Segment{} -} - -func (s *Segment) Start() time.Duration { - return s.start -} - -func (s *Segment) End() time.Duration { - return s.end -} - -func (s *Segment) Text() string { - return s.text -} - -func (s *Segment) SetStart(start time.Duration) { - s.start = start -} - -func (s *Segment) SetEnd(end time.Duration) { - s.end = end -} - -func (s *Segment) SetText(text string) { - s.text = text -} - -type Subtitle struct { - segments []Segment -} - -func NewSubtitle() *Subtitle { - return &Subtitle{ - segments: make([]Segment, 0), - } -} - -func (s *Subtitle) Segments() []Segment { - return s.segments -} - -func (s *Subtitle) AddSegment(segment Segment) { - s.segments = append(s.segments, segment) -} - -func FromRawStream(rawStream RawStream) (*Subtitle, warning.WarningList, error) { - warnings := warning.NewWarningList() - subtitle := NewSubtitle() - - switch rawStream.Format() { - case ASS: - assSubtitle, wrn, err := ass.DecodeSubtitle(rawStream.Filepath(), rawStream.Index()) - warnings.Append(*wrn) - if err != nil { - return nil, *warnings, fmt.Errorf("Error decoding ass subtitle: %v", err) - } - - for _, assSegment := range assSubtitle.Segments() { - segment := NewSegment() - segment.SetStart(assSegment.Start()) - segment.SetEnd(assSegment.End()) - segment.SetText(assSegment.Text()) - - subtitle.AddSegment(*segment) - } - - case PGS: - pgsSubtitle, wrn, err := pgs.DecodeSubtitle(rawStream.Filepath(), rawStream.Index()) - warnings.Append(*wrn) - if err != nil { - return nil, *warnings, fmt.Errorf("Error decoding pgs subtitle: %v", err) - } - - tes := tesseract.NewClient() - defer tes.Close() - - for i, pgsSegment := range pgsSubtitle.Segments() { - segment := NewSegment() - segment.SetStart(pgsSegment.Start()) - - if i+1 < len(pgsSubtitle.Segments()) { - segment.SetEnd(pgsSubtitle.Segments()[i+1].Start()) - } else { - segment.SetEnd(pgsSegment.Start() + time.Second*10) - } - txt := "" - - images, err := pgsSegment.Images() - if err != nil { - return nil, *warnings, fmt.Errorf("Error getting images from pgs subtitle: %v", err) - } - - for _, img := range images { - line, err := tes.ExtractTextFromImage(img, rawStream.Language()) - if err != nil { - return nil, *warnings, fmt.Errorf("Error extracting text from image: %v", err) - } - - txt += line - } - - segment.SetText(txt) - - subtitle.AddSegment(*segment) - } - - } - - return subtitle, *warnings, nil -} - -type EncodedSubtitleFile struct { - extension string - content []byte -} - -func NewEncodedSubtitleFile(content []byte, extension string) *EncodedSubtitleFile { - return &EncodedSubtitleFile{ - extension: extension, - content: content, - } -} - -func (e *EncodedSubtitleFile) Extension() string { - return e.extension -} - -func (e *EncodedSubtitleFile) Content() []byte { - return e.content -} - -type EncodedSubtitle struct { - files []EncodedSubtitleFile -} - -func NewEncodedSubtitle() *EncodedSubtitle { - return &EncodedSubtitle{ - files: make([]EncodedSubtitleFile, 0), - } -} - -func (e *EncodedSubtitle) AddFile(file EncodedSubtitleFile) { - e.files = append(e.files, file) -} - -func (e *EncodedSubtitle) Files() []EncodedSubtitleFile { - return e.files -} - -func (s *Subtitle) Encode(format Format) (*EncodedSubtitle, error) { - switch format { - case SRT: - srtStream := srt.NewStream() - - for _, segment := range s.Segments() { - srtSegment := srt.NewSegment() - - srtSegment.SetStart(segment.Start()) - srtSegment.SetEnd(segment.End()) - srtSegment.SetText(segment.Text()) - - srtStream.AddSegment(*srtSegment) - } - - str := srt.EncodeSubtitle(*srtStream) - - srtFile := NewEncodedSubtitleFile([]byte(str), ".srt") - encodedSubtitle := NewEncodedSubtitle() - encodedSubtitle.AddFile(*srtFile) - - return encodedSubtitle, nil - } - - return nil, fmt.Errorf("Unsupported output format: %v", format) -} +type Subtitle interface{} diff --git a/pkgs/subtitle/text.go b/pkgs/subtitle/text.go new file mode 100644 index 0000000..6670246 --- /dev/null +++ b/pkgs/subtitle/text.go @@ -0,0 +1,47 @@ +package subtitle + +import "time" + +type TextSubtitle struct { + segments []TextSegment +} + +type TextSegment struct { + start time.Duration + end time.Duration + text string +} + +func NewTextSubtitle() *TextSubtitle { + return &TextSubtitle{ + segments: make([]TextSegment, 0), + } +} + +func (s *TextSubtitle) AddSegment(segment TextSegment) { + s.segments = append(s.segments, segment) +} + +func (s *TextSubtitle) Segments() []TextSegment { + return s.segments +} + +func NewTextSegment(start time.Duration, end time.Duration, text string) *TextSegment { + return &TextSegment{ + start: start, + end: end, + text: text, + } +} + +func (s *TextSegment) Start() time.Duration { + return s.start +} + +func (s *TextSegment) End() time.Duration { + return s.end +} + +func (s *TextSegment) Text() string { + return s.text +}