support y4m.xz on frameserver, accept external csv with frames

This commit is contained in:
DataHoarder 2023-08-11 04:01:35 +02:00
parent e8ca9b7064
commit bd735ad262
Signed by: DataHoarder
SSH key fingerprint: SHA256:OLTRf6Fl87G52SiR7sWLGNzlJt4WOX+tfI2yxo0z7xk
4 changed files with 162 additions and 44 deletions

View file

@ -3,27 +3,38 @@
package main
import (
"encoding/csv"
"flag"
"fmt"
"git.gammaspectra.live/S.O.N.G/Ignite/encoder/libaom"
"git.gammaspectra.live/S.O.N.G/Ignite/frame"
"git.gammaspectra.live/S.O.N.G/Ignite/utilities/frameserver"
"io"
"log"
"math"
"os"
"path"
"runtime"
"slices"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
)
const ResetLine = "\033[1A\033[K"
func resetLine() {
print(ResetLine)
}
func main() {
y4mInput := flag.String("input", "", ".y4m input stream")
outputFolder := flag.String("output", "", "Output folder")
maximumSceneLength := flag.Uint64("max-scene-len", 240, "Maximum scene length")
minimumSceneLength := flag.Uint64("min-scene-len", 24, "Minimum scene length")
externalScenes := flag.String("scenes", "", "List of external scene points to split at, CSV file")
workerNumberOption := flag.Uint64("workers", 0, "Number of workers to spawn [0 = automatic]")
videoParams := flag.String("video-params", "", "Extra parameters for video encoder. This can also be used to override defaults.")
@ -34,7 +45,7 @@ func main() {
}
if *maximumSceneLength < *minimumSceneLength {
log.Panic("max-scene-len < min-scene-len")
panic("max-scene-len < min-scene-len")
}
if i, err := os.Stat(*y4mInput); err != nil || i.IsDir() {
@ -46,7 +57,7 @@ func main() {
}
if i, err := os.Stat(*outputFolder); err != nil || !i.IsDir() {
log.Panicf("output = %s must exist and be a directory", *y4mInput)
log.Panicf("output = %s must exist and be a directory", *outputFolder)
}
frameServerPool := frameserver.NewPool(*y4mInput, nil)
@ -61,25 +72,97 @@ func main() {
colorRange = "pc/full"
}
log.Printf("input %s : %dx%d (PAR %s) (FPS %s %.04f) %s %s", *y4mInput, prop.Width, prop.Height, prop.PixelAspectRatio.String(), prop.FrameRate.String(), prop.FrameRate.Float64(), prop.ColorSpace.String(), colorRange)
frameCount := 0
fmt.Printf("input %s : %dx%d (PAR %s) (FPS %s %.04f) %s %s\n", *y4mInput, prop.Width, prop.Height, prop.PixelAspectRatio.String(), prop.FrameRate.String(), prop.FrameRate.Float64(), prop.ColorSpace.String(), colorRange)
var frameNumber int64
scenes := make([]int, 1, 4096)
//pre-process offsets
for f := range allFrames.Channel() {
//todo: scenecut detection
_ = f
log.Printf("processing frame %d/???", frameCount)
scenes := make([]int64, 0, 4096)
if (frameCount - scenes[len(scenes)-1]) >= int(*maximumSceneLength) {
scenes = append(scenes, frameCount)
}
frameCount++
if _, err := os.Stat(*externalScenes); *externalScenes != "" && err == nil {
fmt.Printf("loading scenes from %s", *externalScenes)
func() {
f, err := os.Open(*externalScenes)
if err != nil {
panic(err)
}
defer f.Close()
c := csv.NewReader(f)
var keyStartFrame, keyEndFrame int
var lastSceneEnd int64
for {
record, err := c.Read()
if err != nil {
if err == io.EOF {
break
}
panic(err)
} else if len(record) > 0 && strings.HasPrefix(record[0], "Timecode List:") {
// skip initial header in csv
c.FieldsPerRecord = 0
} else if keyStartFrame == 0 && keyEndFrame == 0 {
for i, k := range record {
if k == "Start Frame" {
keyStartFrame = i
} else if k == "End Frame" {
keyEndFrame = i
}
}
if keyStartFrame == 0 || keyEndFrame == 0 {
panic("could not find Start Frame or End Frame keys")
}
} else {
start := record[keyStartFrame]
end := record[keyEndFrame]
if startNumber, err := strconv.ParseInt(start, 10, 0); err != nil {
panic(err)
} else {
scenes = append(scenes, startNumber-1)
}
if endNumber, err := strconv.ParseInt(end, 10, 0); err != nil {
panic(err)
} else {
lastSceneEnd = endNumber
}
}
}
scenes = append(scenes, lastSceneEnd)
slices.Sort(scenes)
}()
}
//add last entry
scenes = append(scenes, frameCount)
log.Printf("input %s : %d total frames, %d scenes", *y4mInput, frameCount, len(scenes)-1)
if len(scenes) == 0 {
scenes = append(scenes, 0)
print("detecting scenes frame 0/???\n")
//detector := scenedetect.NewThresholdDetectorDefault(int64(*minimumSceneLength))
//pre-process offsets
for f := range allFrames.Channel() {
resetLine()
fmt.Printf("detecting scenes frame %d/???\n", frameNumber)
_ = f
if (frameNumber - scenes[len(scenes)-1]) >= int64(*maximumSceneLength) {
scenes = append(scenes, frameNumber)
}
frameNumber++
}
//add last entries
//scenes = append(scenes, detector.PostProcess(frameNumber)...)
scenes = append(scenes, frameNumber)
} else {
print("counting frames 0/???\n")
for range allFrames.Channel() {
resetLine()
fmt.Printf("counting frames %d/???\n", frameNumber)
frameNumber++
}
}
fmt.Printf("input %s : %d total frames, %d scenes\n", *y4mInput, frameNumber, len(scenes)-1)
//todo: save seekTable
var wg sync.WaitGroup
@ -112,25 +195,30 @@ func main() {
}
}
print("encoded 0/0 frames\n")
wg.Add(1)
go func() {
defer wg.Done()
startTime := time.Now()
for {
t := time.NewTicker(time.Second * 1)
defer t.Stop()
for range t.C {
encoded := framesEncoded.Load()
if encoded > 0 {
//todo: take into account lag-in-frames (start counting after workers * lag-in-frames setting)
timeTaken := time.Now().Sub(startTime)
timePerFrame := timeTaken / time.Duration(encoded)
timeLeft := timePerFrame * time.Duration(uint64(frameCount)-encoded)
log.Printf("encoded %d/%d frames, %.03f%%, ~%.01f seconds left @ rate ~%.03f fps", encoded, frameCount, float64(encoded*100)/float64(frameCount), timeLeft.Seconds(), 1/timePerFrame.Seconds())
timeLeft := timePerFrame * time.Duration(uint64(frameNumber)-encoded)
resetLine()
fmt.Printf("encoded %d/%d frames, %.04f%%, ~%.01f seconds left @ rate ~%.04f fps / ~%.04f spf\n", encoded, frameNumber, float64(encoded*100)/float64(frameNumber), timeLeft.Seconds(), 1/timePerFrame.Seconds(), timePerFrame.Seconds())
}
if encoded >= uint64(frameCount) {
if encoded >= uint64(frameNumber) {
break
}
time.Sleep(time.Second * 1)
}
}()
@ -181,7 +269,7 @@ func main() {
}
}
}
}(index, startFrameNumber, endFrameNumber)
}(index, int(startFrameNumber), int(endFrameNumber))
}
wg.Wait()

View file

@ -8,16 +8,17 @@ import (
type Stream struct {
properties StreamProperties
channel chan Frame
channel <-chan Frame
lock atomic.Bool
}
func NewStream(properties StreamProperties) (*Stream, chan Frame) {
func NewStream(properties StreamProperties) (*Stream, chan<- Frame) {
c := make(chan Frame)
s := &Stream{
properties: properties,
channel: make(chan Frame),
channel: c,
}
return s, s.channel
return s, c
}
type StreamProperties struct {
@ -48,7 +49,7 @@ func (s *Stream) Properties() StreamProperties {
}
// Channel gets the Frame channel, and locks the input
func (s *Stream) Channel() chan Frame {
func (s *Stream) Channel() <-chan Frame {
if s.Lock() {
return s.channel
}
@ -72,20 +73,21 @@ func (s *Stream) Copy(n int) []*Stream {
return nil
}
slice := make([]*Stream, n)
sliceChannels := make([]chan<- Frame, n)
for i := range slice {
slice[i], _ = NewStream(s.Properties())
slice[i], sliceChannels[i] = NewStream(s.Properties())
}
go func() {
defer func() {
for _, sliceItem := range slice {
close(sliceItem.channel)
for _, sliceItem := range sliceChannels {
close(sliceItem)
}
}()
for f := range s.channel {
for _, sliceItem := range slice {
sliceItem.channel <- f
for _, sliceItem := range sliceChannels {
sliceItem <- f
}
}
}()
@ -136,8 +138,9 @@ func (s *Stream) Split(splits ...int) []*Stream {
return nil
}
slice := make([]*Stream, len(splits)+1)
sliceChannels := make([]chan<- Frame, len(splits)+1)
for i := range slice {
slice[i], _ = NewStream(s.Properties())
slice[i], sliceChannels[i] = NewStream(s.Properties())
}
var index int
@ -149,19 +152,20 @@ func (s *Stream) Split(splits ...int) []*Stream {
}
}()
defer func() {
for _, sC := range slice {
close(sC.channel)
for _, sC := range sliceChannels {
close(sC)
}
}()
for f := range s.channel {
for len(splits) > 0 && index >= splits[0] {
close(slice[0].channel)
close(sliceChannels[0])
splits = splits[1:]
sliceChannels = sliceChannels[1:]
slice = slice[1:]
}
slice[0].channel <- f
sliceChannels[0] <- f
index++
}

View file

@ -25,6 +25,19 @@ func New(reader io.ReadSeeker, cachedFrameSeekTable []int64) *FrameServer {
return s
}
func NewXZCompressed(reader io.ReadSeeker, cachedFrameSeekTable []int64) *FrameServer {
s := &FrameServer{}
var err error
settings := make(map[string]any)
settings["seek_table"] = cachedFrameSeekTable
if s.decoder, err = y4m.NewXZCompressedDecoder(reader, settings); err != nil {
return nil
}
return s
}
func (s *FrameServer) GetFrameSeekTable() []int64 {
s.readerLock.Lock()
defer s.readerLock.Unlock()

View file

@ -3,6 +3,7 @@ package frameserver
import (
"git.gammaspectra.live/S.O.N.G/Ignite/frame"
"os"
"path"
"runtime"
"sync"
)
@ -26,13 +27,25 @@ func NewPool(sourcePath string, cachedFrameSeekTable []int64) *Pool {
} else {
p.seekTableLock.RLock()
defer p.seekTableLock.RUnlock()
if s := New(f, p.cachedFrameSeekTable); s != nil {
runtime.SetFinalizer(s, func(s *FrameServer) {
f.Close()
})
return s
if path.Ext(p.sourcePath) == ".xz" {
if s := NewXZCompressed(f, p.cachedFrameSeekTable); s != nil {
runtime.SetFinalizer(s, func(s *FrameServer) {
f.Close()
})
return s
} else {
return nil
}
} else {
return nil
if s := New(f, p.cachedFrameSeekTable); s != nil {
runtime.SetFinalizer(s, func(s *FrameServer) {
f.Close()
})
return s
} else {
return nil
}
}
}
}