diff --git a/.gitignore b/.gitignore index 51d43ba..11390d2 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1 @@ -internal/pb +generated/* diff --git a/Makefile b/Makefile index 85e7a73..b8cafe0 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: proto +.PHONY: proto run proto: protoc ./proto/**/*.proto --go_out=. --go-grpc_out=. diff --git a/cmd/subtle/main.go b/cmd/subtle/main.go index 59628e3..fa188e9 100644 --- a/cmd/subtle/main.go +++ b/cmd/subtle/main.go @@ -1,67 +1,16 @@ package main import ( - "fmt" "log" - "net" - "github.com/nandesh-dev/subtle/internal/filemanager" - "github.com/nandesh-dev/subtle/internal/pb/library" - "github.com/nandesh-dev/subtle/internal/pgs" - "github.com/nandesh-dev/subtle/internal/services" - "github.com/nandesh-dev/subtle/internal/subtitle" - "github.com/nandesh-dev/subtle/internal/tesseract" - "golang.org/x/text/language" - "google.golang.org/grpc" - "google.golang.org/grpc/reflection" + "github.com/nandesh-dev/subtle/internal/routine/library" + "github.com/nandesh-dev/subtle/pkgs/config" ) func main() { - dir, err := filemanager.ReadDirectory("/media") - if err != nil { + if err := config.Init("/config"); err != nil { log.Fatal(err) } - fmt.Println(dir) - - videos, warnings := dir.VideoFiles() - fmt.Println(warnings) - - rawStreams, err := subtitle.ExtractRawStreams(&videos[0]) - if err != nil { - log.Fatal(err) - } - - subtitle, _, err := pgs.DecodeSubtitle(&rawStreams[0]) - if err != nil { - log.Fatal(err) - } - - tes := tesseract.NewClient() - - for _, segment := range subtitle.Segments() { - for _, img := range segment.Images() { - text, err := tes.ExtractTextFromImage(img, language.English) - if err != nil { - log.Fatal(err) - } - - fmt.Println(text) - } - } - - listener, err := net.Listen("tcp", ":3000") - if err != nil { - log.Fatalln("failed to create listener: ", err) - } - - s := grpc.NewServer() - reflection.Register(s) - - libraryService := services.Services().Library - library.RegisterLibraryServiceServer(s, &libraryService) - - if err := s.Serve(listener); err != nil { - log.Fatalln("failed to serve: ", err) - } + library.RunLibraryRoutine() } diff --git a/go.mod b/go.mod index bd62f49..aa6db72 100644 --- a/go.mod +++ b/go.mod @@ -8,15 +8,14 @@ require ( golang.org/x/text v0.18.0 google.golang.org/grpc v1.67.0 google.golang.org/protobuf v1.34.2 + gopkg.in/yaml.v3 v3.0.1 ) require ( github.com/aws/aws-sdk-go v1.38.20 // indirect github.com/jmespath/go-jmespath v0.4.0 // indirect - github.com/mattn/go-sqlite3 v1.14.23 // 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 google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 // indirect - gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 42602d6..37a2865 100644 --- a/go.sum +++ b/go.sum @@ -19,8 +19,6 @@ 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.23 h1:gbShiuAP1W5j9UOksQ06aiiqPMxYecovVGwmTxWtuw0= -github.com/mattn/go-sqlite3 v1.14.23/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= @@ -72,6 +70,7 @@ google.golang.org/grpc v1.67.0 h1:IdH9y6PF5MPSdAntIcpjQ+tXO41pcQsfZV2RxtQgVcw= google.golang.org/grpc v1.67.0/go.mod h1:1gLDyUQU7CTLJI90u3nXZ9ekeghjeM7pTDZlqFNg2AA= google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg= google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.7/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= diff --git a/internal/filemanager/directory.go b/internal/filemanager/directory.go deleted file mode 100644 index c1746f1..0000000 --- a/internal/filemanager/directory.go +++ /dev/null @@ -1,65 +0,0 @@ -package filemanager - -import ( - "fmt" - - "github.com/nandesh-dev/subtle/pkgs/warning" -) - -type Directory struct { - path string - children []Directory - files []File -} - -func NewDirectory(path string) *Directory { - return &Directory{ - path: path, - children: make([]Directory, 0), - files: make([]File, 0), - } -} - -func (d *Directory) AddFile(file File) { - d.files = append(d.files, file) -} - -func (d *Directory) AddChild(child Directory) { - d.children = append(d.children, child) -} - -func (d *Directory) VideoFiles() ([]File, warning.WarningList) { - videos := make([]File, 0) - warnings := warning.NewWarningList() - - for _, file := range d.files { - isVideoFile, err := file.IsVideoFile() - - if err != nil { - warnings.AddWarning(fmt.Errorf("Error checking is file is of type video: %v", err)) - continue - } - - if isVideoFile { - videos = append(videos, file) - } - } - - return videos, *warnings -} - -func (d *Directory) SubtitleFiles() []File { - subtitles := make([]File, 0) - - for _, file := range d.files { - if file.IsSubtitleFile() { - subtitles = append(subtitles, file) - } - } - - return subtitles -} - -func (d *Directory) Children() []Directory { - return d.children -} diff --git a/internal/filemanager/read_dir.go b/internal/filemanager/read_dir.go deleted file mode 100644 index 73ae4fe..0000000 --- a/internal/filemanager/read_dir.go +++ /dev/null @@ -1,31 +0,0 @@ -package filemanager - -import ( - "os" - "path/filepath" -) - -func ReadDirectory(path string) (*Directory, error) { - files, err := os.ReadDir(path) - if err != nil { - return nil, err - } - - directory := NewDirectory(path) - - for _, entry := range files { - if entry.IsDir() { - child, err := ReadDirectory(filepath.Join(path, entry.Name())) - if err != nil { - return nil, err - } - - directory.AddChild(*child) - } - - file := NewFile(filepath.Join(path, entry.Name())) - directory.AddFile(*file) - } - - return directory, nil -} diff --git a/internal/routine/library/library.go b/internal/routine/library/library.go new file mode 100644 index 0000000..60d7648 --- /dev/null +++ b/internal/routine/library/library.go @@ -0,0 +1,92 @@ +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/server/library/library.go b/internal/server/library/library.go new file mode 100644 index 0000000..fc05713 --- /dev/null +++ b/internal/server/library/library.go @@ -0,0 +1,15 @@ +package library + +import ( + "context" + + "github.com/nandesh-dev/subtle/generated/api/library" +) + +type LibraryServiceServer struct { + library.UnsafeLibraryServiceServer +} + +func (s *LibraryServiceServer) GetMedia(ctx context.Context, req *library.GetMediaRequest) (*library.GetMediaResponse, error) { + return nil, nil +} diff --git a/internal/server/server.go b/internal/server/server.go new file mode 100644 index 0000000..f0c5dad --- /dev/null +++ b/internal/server/server.go @@ -0,0 +1,44 @@ +package server + +import ( + "fmt" + "net" + + "github.com/nandesh-dev/subtle/generated/api/library" + library_service "github.com/nandesh-dev/subtle/internal/server/library" + "google.golang.org/grpc" + "google.golang.org/grpc/reflection" +) + +type server struct { + listener *net.Listener + grpcServer *grpc.Server +} + +func New() *server { + return &server{} +} + +func (s *server) Listen(port int, enableReflection bool) error { + listener, err := net.Listen("tcp", ":3000") + if err != nil { + return fmt.Errorf("failed to create listener: %v", err) + } + + s.listener = &listener + + s.grpcServer = grpc.NewServer() + + if enableReflection { + reflection.Register(s.grpcServer) + } + + libraryService := library_service.LibraryServiceServer{} + library.RegisterLibraryServiceServer(s.grpcServer, &libraryService) + + if err := s.grpcServer.Serve(listener); err != nil { + return fmt.Errorf("failed to serve: %v", err) + } + + return nil +} diff --git a/internal/services/library/library.go b/internal/services/library/library.go deleted file mode 100644 index f7a7f30..0000000 --- a/internal/services/library/library.go +++ /dev/null @@ -1,44 +0,0 @@ -package library - -import ( - "context" - - "github.com/nandesh-dev/subtle/internal/filemanager" - "github.com/nandesh-dev/subtle/internal/pb/library" -) - -type LibraryServiceServer struct { - library.UnsafeLibraryServiceServer -} - -func (s *LibraryServiceServer) GetMedia(ctx context.Context, req *library.GetMediaRequest) (*library.GetMediaResponse, error) { - media, _ := filemanager.ReadDirectory("./media") - - var loop func(filemanager.Directory) *library.Directory - loop = func(dir filemanager.Directory) *library.Directory { - children := make([]*library.Directory, len(dir.Children())) - - for i, child := range dir.Children() { - children[i] = loop(child) - } - - videoFiles, _ := dir.VideoFiles() - - videos := make([]*library.Video, len(videoFiles)) - - for i, _ := range videoFiles { - videos[i] = &library.Video{ - Id: "1", - } - } - - return &library.Directory{ - Children: children, - Videos: videos, - } - } - - return &library.GetMediaResponse{ - Directories: []*library.Directory{loop(*media)}, - }, nil -} diff --git a/internal/services/services.go b/internal/services/services.go deleted file mode 100644 index 7f6b982..0000000 --- a/internal/services/services.go +++ /dev/null @@ -1,13 +0,0 @@ -package services - -import "github.com/nandesh-dev/subtle/internal/services/library" - -type services struct { - Library library.LibraryServiceServer -} - -func Services() services { - return services{ - Library: library.LibraryServiceServer{}, - } -} diff --git a/internal/srt/routines/library/library.go b/internal/srt/routines/library/library.go deleted file mode 100644 index e196664..0000000 --- a/internal/srt/routines/library/library.go +++ /dev/null @@ -1 +0,0 @@ -package library diff --git a/internal/config/config.go b/pkgs/config/config.go similarity index 54% rename from internal/config/config.go rename to pkgs/config/config.go index 83044d0..254da2a 100644 --- a/internal/config/config.go +++ b/pkgs/config/config.go @@ -9,7 +9,29 @@ import ( "gopkg.in/yaml.v3" ) +type Server struct { + Port int + GRPCReflection bool +} + +type AutoExtract struct { + Formats []string + Languages []string + OutputFormat string +} + +type WatchDirectory struct { + Path string + AutoExtract AutoExtract +} + +type Media struct { + WatchDirectories []WatchDirectory +} + type t struct { + Server Server + Media Media } var ( @@ -24,6 +46,25 @@ func Config() *t { func Init(basepath string) (e error) { once.Do(func() { + config = t{ + Server: Server{ + Port: 3000, + GRPCReflection: false, + }, + Media: Media{ + WatchDirectories: []WatchDirectory{ + { + Path: "/media", + AutoExtract: AutoExtract{ + Languages: []string{"eng"}, + Formats: []string{"ass"}, + OutputFormat: "srt", + }, + }, + }, + }, + } + path = filepath.Join(basepath, "config.yaml") file, err := os.ReadFile(path) @@ -35,12 +76,13 @@ func Init(basepath string) (e error) { return } file.Close() - } else { - e = fmt.Errorf("Error reading config file: %v", err) + + e = Write() return } - file = make([]byte, 0) + e = fmt.Errorf("Error reading config file: %v", err) + return } if err := yaml.Unmarshal(file, &config); err != nil { @@ -61,7 +103,7 @@ func Write() error { return fmt.Errorf("Error marshaling file: %v", err) } - if err := os.WriteFile(path, output, 644); err != nil { + if err := os.WriteFile(path, output, 0644); err != nil { return fmt.Errorf("Error writing config: %v", err) } diff --git a/pkgs/filemanager/directory.go b/pkgs/filemanager/directory.go new file mode 100644 index 0000000..8a305df --- /dev/null +++ b/pkgs/filemanager/directory.go @@ -0,0 +1,41 @@ +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) +} + +func (d *Directory) VideoFiles() []VideoFile { + return d.videos +} + +func (d *Directory) ExtraSubtitleFiles() []SubtitleFile { + return d.extraSubtitles +} + +func (d *Directory) Children() []Directory { + return d.children +} diff --git a/internal/filemanager/file.go b/pkgs/filemanager/file.go similarity index 62% rename from internal/filemanager/file.go rename to pkgs/filemanager/file.go index c47f4e8..a886073 100644 --- a/internal/filemanager/file.go +++ b/pkgs/filemanager/file.go @@ -5,35 +5,12 @@ import ( "fmt" "path/filepath" "slices" - "strings" ffmpeg "github.com/u2takey/ffmpeg-go" ) -type File struct { - path string -} - -func NewFile(path string) *File { - return &File{ - path: path, - } -} - -func (f *File) Path() string { - return f.path -} - -func (f *File) Extension() string { - return filepath.Ext(f.path) -} - -func (f *File) Basename() string { - return strings.TrimSuffix(filepath.Base(f.path), f.Extension()) -} - -func (f *File) IsVideoFile() (bool, error) { - rawProbeResult, err := ffmpeg.Probe(f.path) +func IsVideoFile(path string) (bool, error) { + rawProbeResult, err := ffmpeg.Probe(path) if err != nil { return false, fmt.Errorf("Failed to probe file stats: %v", err) } @@ -64,6 +41,6 @@ func (f *File) IsVideoFile() (bool, error) { return false, nil } -func (f *File) IsSubtitleFile() bool { - return slices.Contains([]string{}, f.Extension()) +func IsSubtitleFile(path string) bool { + return slices.Contains([]string{".srt"}, filepath.Ext(path)) } diff --git a/pkgs/filemanager/read_dir.go b/pkgs/filemanager/read_dir.go new file mode 100644 index 0000000..892151f --- /dev/null +++ b/pkgs/filemanager/read_dir.go @@ -0,0 +1,78 @@ +package filemanager + +import ( + "fmt" + "os" + "path/filepath" + + "github.com/nandesh-dev/subtle/pkgs/warning" +) + +func ReadDirectory(path string) (*Directory, warning.WarningList, error) { + warnings := warning.NewWarningList() + files, err := os.ReadDir(path) + if err != nil { + return nil, *warnings, err + } + + directory := NewDirectory(path) + videos := make([]*VideoFile, 0) + subtitles := make([]*SubtitleFile, 0) + + for _, entry := range files { + entrypath := filepath.Join(path, entry.Name()) + if entry.IsDir() { + child, w, err := ReadDirectory(entrypath) + warnings.Append(w) + + if err != nil { + return nil, *warnings, err + } + + directory.AddChild(*child) + } + + if IsSubtitleFile(entrypath) { + subtitles = append(subtitles, NewSubtitleFile(entrypath)) + continue + } + + isVideoFile, err := IsVideoFile(entrypath) + + if err != nil { + warnings.AddWarning(fmt.Errorf("Error checking if file is video: %v", err)) + continue + } + + 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 + } + } + + 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.go new file mode 100644 index 0000000..e7c5338 --- /dev/null +++ b/pkgs/filemanager/subtitle.go @@ -0,0 +1,44 @@ +package filemanager + +import ( + "path/filepath" + "strings" +) + +type SubtitleFile struct { + path string +} + +func NewSubtitleFile(path string) *SubtitleFile { + return &SubtitleFile{ + path: path, + } +} + +func (s *SubtitleFile) Path() string { + return s.path +} + +func (s *SubtitleFile) Basename() string { + pt := strings.Split(strings.TrimSuffix(filepath.Base(s.path), s.Extension()), ".") + + if len(pt) >= 2 { + return strings.Join(pt[:len(pt)-1], ".") + } + + return pt[0] +} + +func (s *SubtitleFile) LanguageCode() string { + pt := strings.Split(filepath.Base(s.path), ".") + + if len(pt) >= 3 { + return pt[len(pt)-2] + } + + return "" +} + +func (s *SubtitleFile) Extension() string { + return filepath.Ext(s.path) +} diff --git a/pkgs/filemanager/video.go b/pkgs/filemanager/video.go new file mode 100644 index 0000000..8207d6b --- /dev/null +++ b/pkgs/filemanager/video.go @@ -0,0 +1,62 @@ +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/internal/ass/ass.go b/pkgs/subtitle/ass/ass.go similarity index 88% rename from internal/ass/ass.go rename to pkgs/subtitle/ass/ass.go index 45c29dc..4ff6cc0 100644 --- a/internal/ass/ass.go +++ b/pkgs/subtitle/ass/ass.go @@ -49,3 +49,7 @@ func NewStream() *Stream { func (s *Stream) Segments() []Segment { return s.segments } + +func (s *Stream) AddSegment(segment Segment) { + s.segments = append(s.segments, segment) +} diff --git a/internal/ass/decoder.go b/pkgs/subtitle/ass/decoder.go similarity index 91% rename from internal/ass/decoder.go rename to pkgs/subtitle/ass/decoder.go index 07d8579..0c862ad 100644 --- a/internal/ass/decoder.go +++ b/pkgs/subtitle/ass/decoder.go @@ -28,24 +28,15 @@ const ( Format ) -type RawStream interface { - Index() int - Filepath() string -} - -func (s *Stream) AddSegment(segment Segment) { - s.segments = append(s.segments, segment) -} - -func DecodeSubtitle(rawStream RawStream) (Stream, *warning.WarningList, error) { +func DecodeSubtitle(path string, index int) (Stream, *warning.WarningList, error) { stream := NewStream() warnings := warning.NewWarningList() var subtitleBuf, errorBuf bytes.Buffer ffmpeg.LogCompiledCommand = false - err := ffmpeg.Input(rawStream.Filepath()). - Output("pipe:", ffmpeg.KwArgs{"map": fmt.Sprintf("0:%v", rawStream.Index()), "f": "ass"}). + err := ffmpeg.Input(path). + Output("pipe:", ffmpeg.KwArgs{"map": fmt.Sprintf("0:%v", index), "f": "ass"}). WithOutput(&subtitleBuf). WithErrorOutput(&errorBuf). Run() diff --git a/internal/ass/reader.go b/pkgs/subtitle/ass/reader.go similarity index 100% rename from internal/ass/reader.go rename to pkgs/subtitle/ass/reader.go diff --git a/internal/pgs/decoder.go b/pkgs/subtitle/pgs/decoder.go similarity index 90% rename from internal/pgs/decoder.go rename to pkgs/subtitle/pgs/decoder.go index a355989..bb28369 100644 --- a/internal/pgs/decoder.go +++ b/pkgs/subtitle/pgs/decoder.go @@ -8,20 +8,15 @@ import ( ffmpeg "github.com/u2takey/ffmpeg-go" ) -type RawStream interface { - Index() int - Filepath() string -} - -func DecodeSubtitle(rawStream RawStream) (Stream, *warning.WarningList, error) { +func DecodeSubtitle(path string, index int) (Stream, *warning.WarningList, error) { stream := NewStream() warnings := warning.NewWarningList() var subtitleBuf, errorBuf bytes.Buffer ffmpeg.LogCompiledCommand = false - err := ffmpeg.Input(rawStream.Filepath()). - Output("pipe:", ffmpeg.KwArgs{"map": fmt.Sprintf("0:%v", rawStream.Index()), "c:s": "copy", "f": "sup"}). + err := ffmpeg.Input(path). + Output("pipe:", ffmpeg.KwArgs{"map": fmt.Sprintf("0:%v", index), "c:s": "copy", "f": "sup"}). WithOutput(&subtitleBuf). WithErrorOutput(&errorBuf). Run() diff --git a/internal/pgs/display_set.go b/pkgs/subtitle/pgs/display_set.go similarity index 100% rename from internal/pgs/display_set.go rename to pkgs/subtitle/pgs/display_set.go diff --git a/internal/pgs/header.go b/pkgs/subtitle/pgs/header.go similarity index 100% rename from internal/pgs/header.go rename to pkgs/subtitle/pgs/header.go diff --git a/internal/pgs/ods.go b/pkgs/subtitle/pgs/ods.go similarity index 100% rename from internal/pgs/ods.go rename to pkgs/subtitle/pgs/ods.go diff --git a/internal/pgs/pcs.go b/pkgs/subtitle/pgs/pcs.go similarity index 100% rename from internal/pgs/pcs.go rename to pkgs/subtitle/pgs/pcs.go diff --git a/internal/pgs/pds.go b/pkgs/subtitle/pgs/pds.go similarity index 100% rename from internal/pgs/pds.go rename to pkgs/subtitle/pgs/pds.go diff --git a/internal/pgs/pgs.go b/pkgs/subtitle/pgs/pgs.go similarity index 75% rename from internal/pgs/pgs.go rename to pkgs/subtitle/pgs/pgs.go index ad35c40..28e096a 100644 --- a/internal/pgs/pgs.go +++ b/pkgs/subtitle/pgs/pgs.go @@ -14,12 +14,13 @@ func NewSegment() *Segment { return &Segment{} } -func (s *Segment) Images() []image.Image { - return s.images +func (s *Segment) Images() ([]image.Image, error) { + return s.images, nil } -func (s *Segment) AddImages(img []image.Image) { - s.images = append(s.images, img...) +func (s *Segment) AddImages(images []image.Image) error { + s.images = append(s.images, images...) + return nil } func (s *Segment) Start() time.Duration { diff --git a/internal/pgs/reader.go b/pkgs/subtitle/pgs/reader.go similarity index 100% rename from internal/pgs/reader.go rename to pkgs/subtitle/pgs/reader.go diff --git a/internal/pgs/segments.go b/pkgs/subtitle/pgs/segments.go similarity index 100% rename from internal/pgs/segments.go rename to pkgs/subtitle/pgs/segments.go diff --git a/internal/pgs/wds.go b/pkgs/subtitle/pgs/wds.go similarity index 100% rename from internal/pgs/wds.go rename to pkgs/subtitle/pgs/wds.go diff --git a/internal/subtitle/raw_stream.go b/pkgs/subtitle/raw_stream.go similarity index 92% rename from internal/subtitle/raw_stream.go rename to pkgs/subtitle/raw_stream.go index 6f59465..50a4f6e 100644 --- a/internal/subtitle/raw_stream.go +++ b/pkgs/subtitle/raw_stream.go @@ -10,13 +10,6 @@ import ( "golang.org/x/text/language" ) -type Format int - -const ( - PGS Format = iota - ASS -) - type RawStream struct { filepath string index int @@ -58,12 +51,8 @@ func (s *RawStream) Language() language.Tag { return s.language } -type file interface { - Path() string -} - -func ExtractRawStreams(file file) ([]RawStream, error) { - rawResultData, err := ffmpeg.Probe(file.Path()) +func ExtractRawStreams(path string) ([]RawStream, error) { + rawResultData, err := ffmpeg.Probe(path) warnings := warning.NewWarningList() if err != nil { @@ -94,7 +83,7 @@ func ExtractRawStreams(file file) ([]RawStream, error) { continue } - rawStream := NewRawStream(file.Path()) + rawStream := NewRawStream(path) codecName, codecNameExist := rawStreamData.(map[string]any)["codec_name"].(string) if !codecNameExist { diff --git a/internal/srt/encoder.go b/pkgs/subtitle/srt/encoder.go similarity index 100% rename from internal/srt/encoder.go rename to pkgs/subtitle/srt/encoder.go diff --git a/internal/srt/srt.go b/pkgs/subtitle/srt/srt.go similarity index 100% rename from internal/srt/srt.go rename to pkgs/subtitle/srt/srt.go diff --git a/pkgs/subtitle/subtitle.go b/pkgs/subtitle/subtitle.go new file mode 100644 index 0000000..b269bfc --- /dev/null +++ b/pkgs/subtitle/subtitle.go @@ -0,0 +1,216 @@ +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) +} diff --git a/internal/tesseract/tesseract.go b/pkgs/tesseract/tesseract.go similarity index 100% rename from internal/tesseract/tesseract.go rename to pkgs/tesseract/tesseract.go diff --git a/proto/library/library.proto b/proto/library/library.proto index 9681d8f..be667f7 100644 --- a/proto/library/library.proto +++ b/proto/library/library.proto @@ -2,7 +2,7 @@ syntax = "proto3"; package library; -option go_package = "internal/pb/library"; +option go_package = "generated/api/library"; service LibraryService { rpc GetMedia(GetMediaRequest) returns (GetMediaResponse); diff --git a/shell.nix b/shell.nix index 15e5dce..92c380c 100644 --- a/shell.nix +++ b/shell.nix @@ -10,5 +10,6 @@ pkgs.mkShell { protobuf protoc-gen-go-grpc protoc-gen-go + graphviz ]; }