encodehls.go

encodehls.go

package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"

	"github.com/ssttevee/go-ffmpeg"
)

var streamConfigurations = []StreamVariant{
	{240, 500000, 64000},
	{360, 800000, 96000},
	{480, 2000000, 96000},
	{720, 4000000, 128000},
	{1080, 6000000, 128000},
}

type StreamVariant struct {
	height, videoBitrate, audioBitrate int
}

type StreamBuilder struct {
	outputDir, masterName, variantName, segmentName string

	gpu bool

	variants      []*StreamVariant
	initalOptions []ffmpeg.CliOption
}

func NewStreamBuilder(outputDir, masterName, variantName, segmentName string, gpu bool, initalOptions ...ffmpeg.CliOption) *StreamBuilder {
	return &StreamBuilder{
		outputDir:     outputDir,
		masterName:    masterName,
		variantName:   variantName,
		segmentName:   segmentName,
		gpu:           gpu,
		initalOptions: initalOptions,
	}
}

func (b *StreamBuilder) AddVariant(cfg *StreamVariant) {
	b.variants = append(b.variants, cfg)
}

func (b *StreamBuilder) Build(job *ffmpeg.Job) {
	options := b.initalOptions

	filter := "scale"
	if b.gpu {
		options = append(options, ffmpeg.Option("-c:v", "h264_nvenc"))
		filter += "_npp"
	}

	var streams []string
	for i, cfg := range b.variants {
		streamIndex := strconv.Itoa(i)

		options = append(
			options,
			ffmpeg.Option("-b:v:"+streamIndex, strconv.Itoa(cfg.videoBitrate)),
			ffmpeg.Option("-b:a:"+streamIndex, strconv.Itoa(cfg.audioBitrate)),
			ffmpeg.Option("-filter:v:"+streamIndex, filter+"=-2:"+strconv.Itoa(cfg.height)),
			ffmpeg.Option("-map", "0:v"),
			ffmpeg.Option("-map", "0:a"),
		)
		streams = append(streams, "v:"+streamIndex+",a:"+streamIndex)
	}

	job.AddOutputFile(
		path.Join(b.outputDir, b.variantName),
		append(
			options,
			ffmpeg.Option("-f", "hls"),
			ffmpeg.Option("-var_stream_map", strings.Join(streams, " ")),
			ffmpeg.Option("-hls_segment_filename", path.Join(b.outputDir, b.segmentName)),
			ffmpeg.Option("-hls_playlist_type", "vod"),
			ffmpeg.Option("-master_pl_name", b.masterName),
		)...,
	)
}

func main() {
	var ffmpegPath, ffprobePath, inputFile string

	var overwrite, gpu, debug bool

	flag.StringVar(&ffmpegPath, "ffmpeg", "", "path to ffmpeg")
	flag.StringVar(&ffprobePath, "ffprobe", "", "path to ffmpeg")
	flag.BoolVar(&overwrite, "y", false, "overwrite file")
	flag.StringVar(&inputFile, "i", "", "input file")
	flag.BoolVar(&gpu, "gpu", false, "user gpu")
	flag.BoolVar(&debug, "debug", false, "user gpu")

	flag.Parse()

	fi, _ := os.Stdin.Stat()
	if (fi.Mode()&os.ModeCharDevice) != 0 && inputFile == "" {
		log.Fatalln("missing input file")
	}

	outputFile := flag.Arg(0)
	if outputFile == "" {
		log.Fatalln("missing output file")
	}

	envPath := os.Getenv("FFMPEG_PATH")
	if ffmpegPath == "" && ffprobePath == "" && envPath != "" {
		ffmpegPath = path.Join(envPath, "ffmpeg")
		ffprobePath = path.Join(envPath, "ffprobe")
	}

	var cfg *ffmpeg.Configuration
	var err error
	if ffmpegPath != "" && ffprobePath != "" {
		cfg, err = ffmpeg.NewConfiguration(ffmpegPath, ffprobePath)
	} else {
		cfg, err = ffmpeg.DefaultConfiguration()
	}

	if err != nil {
		panic(err)
	}

	var globalOptions []ffmpeg.CliOption
	if gpu {
		globalOptions = append(globalOptions, ffmpeg.Option("-hwaccel", "cuvid"))
	}

	if overwrite {
		globalOptions = append(globalOptions, ffmpeg.Flag("-y"))
	}

	job := cfg.NewJob(globalOptions...)

	inputOptions := []ffmpeg.CliOption{
		ffmpeg.Option("-vsync", "0"),
	}

	if gpu {
		inputOptions = append(inputOptions, ffmpeg.Option("-c:v", "h264_cuvid"))
	}

	var metadata *ffmpeg.Metadata
	if inputFile == "" {
		metadata, err = job.AddInputReader(os.Stdin, inputOptions...)
	} else {
		metadata, err = job.AddInputFile(inputFile, inputOptions...)
	}

	if err != nil {
		panic(err)
	}

	totalFrames := metadata.Format.NbFrames
	totalTime, _ := strconv.ParseFloat(metadata.Format.Duration, 64)

	var videoStream *ffmpeg.Stream
	for i := range metadata.Streams {
		if metadata.Streams[i].CodecType == "video" {
			videoStream = &metadata.Streams[i]
			break
		}
	}

	if videoStream == nil {
		panic("no video stream found")
	}

	lis, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		panic(err)
	}

	go http.Serve(lis, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		outputFile := r.URL.Path[1:]
		f, err := os.Create(outputFile)
		if os.IsNotExist(err) {
			if err := os.MkdirAll(path.Dir(outputFile), os.ModePerm); err != nil {
				panic(err.Error())
			}

			f, err = os.Create(outputFile)
		}

		if err != nil {
			panic("error receiving file: " + err.Error())
		}

		defer f.Close()

		n, err := io.Copy(f, r.Body)
		if err != nil {
			fmt.Println("error receiving file:", err)
			return
		}

		fmt.Println("received", r.URL.Path, n)
	}))

	builder := NewStreamBuilder(
		fmt.Sprintf("http://%s/%s", lis.Addr().String(), outputFile),
		"video.m3u8",
		"%v.m3u8",
		"%v/%d.ts",
		gpu,
		ffmpeg.Option("-preset", "fast"),
		ffmpeg.Option("-crf", "51"),
	)

	for i := range streamConfigurations {
		if videoStream.Height >= streamConfigurations[i].height {
			builder.AddVariant(&streamConfigurations[i])
		}
	}

	builder.Build(job)

	var debugOutput *os.File
	if debug {
		debugOutput, err = os.Create("debug.log")
		if err != nil {
			panic(err)
		}

		defer debugOutput.Close()
	}

	statusChan, err := job.StartDebug(context.Background(), debugOutput)
	if err != nil {
		panic(err.Error())
	}

	for status := range statusChan {
		switch v := status.(type) {
		case *ffmpeg.Progress:
			var percent float64
			if totalFrames > 0 {
				percent = float64(v.Frame) / float64(totalFrames)
			} else if totalTime > 0 {
				percent = v.Time / totalTime
			} else {
				log.Printf("%#v", v)
				continue
			}

			log.Printf("%.2f%% %#v", percent*100, v)
		case *ffmpeg.Done:
			log.Printf("done")
			return
		case *ffmpeg.Error:
			log.Fatalf("%#v", v)
		}
	}
}