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

//go:build cgo && !disable_codec_libaom

package main

import (
	"encoding/csv"
	"errors"
	"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.")

	flag.Parse()

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

	if *maximumSceneLength < *minimumSceneLength {
		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", *outputFolder)
	}

	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"
	}

	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([]int64, 0, 4096)

	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 errors.Is(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)
		}()
	}

	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

	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
		}
	}

	print("encoded 0/0 frames\n")

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

		startTime := time.Now()

		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(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(frameNumber) {
				break
			}
		}
	}()

	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, log.Default()); 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, int(startFrameNumber), int(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
}