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

cli/encode-libaom/encode-libaom.go

@@ -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)
+	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)
-	frameCount := 0
+	var frameNumber int64
 
-	scenes := make([]int, 1, 4096)
+	scenes := make([]int64, 0, 4096)
-	//pre-process offsets
-	for f := range allFrames.Channel() {
-		//todo: scenecut detection
-		_ = f
-		log.Printf("processing frame %d/???", frameCount)
 
-		if (frameCount - scenes[len(scenes)-1]) >= int(*maximumSceneLength) {
+	if _, err := os.Stat(*externalScenes); *externalScenes != "" && err == nil {
-			scenes = append(scenes, frameCount)
+		fmt.Printf("loading scenes from %s", *externalScenes)
-		}
+		func() {
-		frameCount++
+			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)
+				timeLeft := timePerFrame * time.Duration(uint64(frameNumber)-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())
+				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()

frame/stream.go

@@ -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 {
+			for _, sliceItem := range sliceChannels {
-				close(sliceItem.channel)
+				close(sliceItem)
 			}
 		}()
 
 		for f := range s.channel {
-			for _, sliceItem := range slice {
+			for _, sliceItem := range sliceChannels {
-				sliceItem.channel <- f
+				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 {
+			for _, sC := range sliceChannels {
-				close(sC.channel)
+				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++
 		}

utilities/frameserver/frameserver.go

@@ -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()

utilities/frameserver/pool.go

@@ -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) {
+			if path.Ext(p.sourcePath) == ".xz" {
-					f.Close()
+				if s := NewXZCompressed(f, p.cachedFrameSeekTable); s != nil {
-				})
+					runtime.SetFinalizer(s, func(s *FrameServer) {
-				return s
+						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
+				}
 			}
 		}
 	}