Ignite/cli/encode-libaom/encode-libaom.go

package main

import (
	"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"
	"log"
	"math"
	"os"
	"path"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

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")
	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.")

	flag.Parse()

	if *minimumSceneLength < 2 {
		log.Panic("min-scene-len < 2")
	}

	if *maximumSceneLength < *minimumSceneLength {
		log.Panic("max-scene-len < min-scene-len")
	}

	if i, err := os.Stat(*y4mInput); err != nil || i.IsDir() {
		log.Panicf("%s must exist and be a file", *y4mInput)
	}

	if i, err := os.Stat(*y4mInput); err != nil || i.IsDir() {
		log.Panicf("input = %s must exist and be a file", *y4mInput)
	}

	if i, err := os.Stat(*outputFolder); err != nil || !i.IsDir() {
		log.Panicf("output = %s must exist and be a directory", *y4mInput)
	}

	frameServerPool := frameserver.NewPool(*y4mInput, nil)

	allFrames := frameServerPool.GetAllFrames()
	if allFrames == nil {
		log.Panicf("could not open %s", *y4mInput)
	}
	prop := allFrames.Properties()
	colorRange := "tv/partial"
	if prop.FullColorRange {
		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

	scenes := make([]int, 1, 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) {
			scenes = append(scenes, frameCount)
		}
		frameCount++
	}
	//add last entry
	scenes = append(scenes, frameCount)

	log.Printf("input %s : %d total frames, %d scenes", *y4mInput, frameCount, len(scenes)-1)
	//todo: save seekTable

	var wg sync.WaitGroup

	var framesEncoded atomic.Uint64

	workerNumber := uint64(runtime.NumCPU())
	if *workerNumberOption != 0 {
		workerNumber = *workerNumberOption
	}

	//max one worker per scene, use this to bump threads
	if uint64(len(scenes)-1) < workerNumber {
		workerNumber = uint64(len(scenes) - 1)
	}

	workerChannel := make(chan struct{}, workerNumber)
	for i := uint64(0); i < workerNumber; i++ {
		workerChannel <- struct{}{}
	}

	defaultParams := getDefaultParameters(prop, int(workerNumber))

	for _, p := range strings.Split(*videoParams, "") {
		splits := strings.Split(p, "=")
		if len(splits) > 1 {
			defaultParams[splits[0]] = strings.Trim(strings.Join(splits[1:], "="), "'\"")
		} else {
			defaultParams[splits[0]] = 1
		}
	}

	wg.Add(1)
	go func() {
		defer wg.Done()

		startTime := time.Now()

		for {
			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())
			}
			if encoded >= uint64(frameCount) {
				break
			}
			time.Sleep(time.Second * 1)
		}
	}()

	for index, startFrameNumber := range scenes[:len(scenes)-1] {
		//todo: make this streamable
		endFrameNumber := scenes[index+1]

		//take lock
		<-workerChannel

		wg.Add(1)
		go func(index, startFrameNumber, endFrameNumber int) {
			defer wg.Done()
			defer func() {
				//return lock
				workerChannel <- struct{}{}
			}()

			outputPath := path.Join(*outputFolder, fmt.Sprintf("scene_%d_%d_%d.ivf", index, startFrameNumber, endFrameNumber-1))
			if f, err := os.Create(outputPath); err != nil {
				log.Panicf("could not create %s: %s", outputPath, err.Error())
			} else {
				defer f.Close()

				settings := make(map[string]any)

				for k, v := range defaultParams {
					settings[k] = v
				}

				stream := frameServerPool.GetFrames(startFrameNumber, endFrameNumber)

				if encoder, err := libaom.NewEncoder(f, stream.Properties(), settings); err != nil {
					log.Panicf("could not create encoder for %s: %s", outputPath, err.Error())
				} else {
					defer encoder.Close()
					frameNumber := 0
					for f := range stream.Channel() {
						if err = encoder.Encode(f); err != nil {
							log.Panicf("error while encoding %s at frame %d: %s", outputPath, frameNumber, err.Error())
						}
						//log.Printf("scene %d: encoded %d/%d", index, frameNumber+1, endFrameNumber-startFrameNumber)
						frameNumber++
						framesEncoded.Add(1)
					}
					if err = encoder.Flush(); err != nil {
						log.Panicf("error while flushing %s: %s", outputPath, err.Error())
					}
				}
			}
		}(index, startFrameNumber, endFrameNumber)
	}

	wg.Wait()
}

func getDefaultParameters(properties frame.StreamProperties, workerNumber int) map[string]any {
	defaultParams := make(map[string]any)
	defaultParams["row-mt"] = 1
	defaultParams["frame-parallel"] = 1
	defaultParams["cpu-used"] = 6

	//crf
	defaultParams["end-usage"] = "q"
	defaultParams["cq-level"] = 30

	threadCount := runtime.NumCPU() / workerNumber
	if threadCount < 1 {
		threadCount = 1
	}
	defaultParams["threads"] = threadCount

	//todo: maybe change to thread-independent method? example below
	/*
		tilesHorizontal := utilities.Max((properties.Width-1)/720, 1)
		tilesVertical := utilities.Max((properties.Height-1)/720, 1)

		defaultParams["tile-columns"] = utilities.Log2(tilesHorizontal)
		defaultParams["tile-rows"] = utilities.Log2(tilesVertical)
	*/

	if threadCount <= 6 {
		defaultParams["tile-columns"] = 2
		defaultParams["tile-rows"] = 1
	} else {
		po2 := math.Log2(float64(threadCount))
		_ = po2
		defaultParams["tile-columns"] = int(math.Ceil(po2 * 0.75))
		defaultParams["tile-rows"] = int(math.Ceil(po2 * 0.25))
	}

	return defaultParams
}