Skip to content

Commit

Permalink
Add media and extract routine
Browse files Browse the repository at this point in the history
  • Loading branch information
nandesh-dev committed Oct 10, 2024
1 parent a303036 commit eb33c59
Show file tree
Hide file tree
Showing 34 changed files with 588 additions and 678 deletions.
23 changes: 20 additions & 3 deletions cmd/subtle/main.go
Original file line number Diff line number Diff line change
@@ -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)
}
}
7 changes: 6 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 12 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,19 @@ 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=
github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U=
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=
Expand Down Expand Up @@ -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=
Expand All @@ -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=
169 changes: 169 additions & 0 deletions internal/routine/extract/extract.go
Original file line number Diff line number Diff line change
@@ -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
}
92 changes: 0 additions & 92 deletions internal/routine/library/library.go

This file was deleted.

Loading

0 comments on commit eb33c59

Please sign in to comment.