buttrate.go

package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/wcharczuk/go-chart"
)

var arg struct {
	expected float64
	fileName string
	outName  string
	outType  string
	interval int
}

func init() {
	flag.Float64Var(&arg.expected, "e", 0, "Expected bitrate (kbit/s)")
	flag.StringVar(&arg.fileName, "i", "", "Input filename")
	flag.StringVar(&arg.outName, "o", "", "Output filename")
	flag.StringVar(&arg.outType, "t", "svg", "Output file type (png, svg)")
	flag.IntVar(&arg.interval, "g", 3, "Bitrate interval (seconds)")
}

func main() {
	flag.Parse()

	if flag.NFlag() < 1 {
		fmt.Fprintf(os.Stderr, "Simple Bitrate Plotting Tool\n\n")
		fmt.Fprintf(os.Stderr, "Usage: %s -i <input> -o <output> -e <bitrate> [-t <type>] [-g <interval>]\n\n", os.Args[0])
		fmt.Fprintln(os.Stderr, "Parameters:")
		flag.PrintDefaults()
		return
	}

	if arg.expected <= 0 {
		fmt.Println("Invalid expected bitrate.")
		os.Exit(1)
	} else if arg.fileName == "" {
		fmt.Println("No input specified.")
		os.Exit(1)
	} else if arg.outName == "" {
		fmt.Println("No output specified")
		os.Exit(1)
	} else if arg.outType != "png" && arg.outType != "svg" {
		fmt.Println("Invalid output file type.")
		os.Exit(1)
	} else if arg.interval <= 0 {
		fmt.Println("Invalid bitrate interval")
		os.Exit(1)
	}

	ffprobe := exec.Command("ffprobe", "-of", "csv", "-show_packets", arg.fileName)
	ffprobeOut, err := ffprobe.StdoutPipe()
	if err != nil {
		fmt.Println(err)
		return
	}
	err = ffprobe.Start()
	if err != nil {
		fmt.Println(err)
		return
	}

	r := csv.NewReader(ffprobeOut)
	rec, err := r.ReadAll()
	if err != nil {
		fmt.Println(err)
		return
	}

	err = ffprobe.Wait()
	if err != nil {
		fmt.Println(err)
		return
	}

	durationSum := float64(0)
	sizeSum := int64(0)
	ptsCur := float64(0)
	pts := make([]float64, 0, 1000)
	keyframes := make([]float64, 0, 1000)
	bitrate := make([]float64, 0, 1000)
	totalBits := int64(0)
	totalDuration := float64(0)

	for _, v := range rec {
		if v[13] == "N/A" {
			continue
		}
		ptsTime, err := strconv.ParseFloat(v[4], 64)
		if err != nil {
			continue
		} else if ptsTime < 0 {
			continue
		}
		durationTime, err := strconv.ParseFloat(v[8], 64)
		if err != nil {
			continue
		}
		size, err := strconv.ParseInt(v[11], 10, 64)
		if err != nil {
			continue
		}
		if strings.Contains(v[13], "K") {
			keyframes = append(keyframes, ptsTime)
		}

		if durationSum > float64(arg.interval) {
			pts = append(pts, float64(int64(ptsCur*100))/100)
			bitrate = append(bitrate, float64(sizeSum)/durationSum/1000)
			ptsCur = ptsTime
			durationSum = 0
			sizeSum = 0
		}
		totalBits += size * 8
		if v[12] != "N/A" {
			_, err := strconv.Atoi(v[12])
			if err != nil {
				fmt.Println("ffprobe issue?")
				return
			}
			totalDuration += durationTime
			durationSum += durationTime
		}
		sizeSum += size * 8
	}
	if durationSum > 0 {
		pts = append(pts, float64(int64(ptsCur*100))/100)
		bitrate = append(bitrate, float64(sizeSum)/durationSum/1000)
	}

	avgBitrate := float64(totalBits) / totalDuration / 1000

	fmt.Printf("Duration: %f s\n", totalDuration)
	fmt.Printf("Size: %f MB\n", float64(totalBits)/8/1024/1024)
	fmt.Printf("Average bitrate: %f kbit/s\n", avgBitrate)

	graph := chart.Chart{
		Series: []chart.Series{
			chart.ContinuousSeries{
				Name:    filepath.Base(arg.fileName),
				XValues: pts,
				YValues: bitrate,
				Style: chart.Style{
					StrokeColor: chart.GetDefaultColor(0),
					FillColor:   chart.GetDefaultColor(0).WithAlpha(64),
				},
			},
			chart.ContinuousSeries{
				Name:    "Average Bitrate",
				XValues: []float64{0, pts[len(pts)-1]},
				YValues: []float64{avgBitrate, avgBitrate},
				Style: chart.Style{
					StrokeColor: chart.GetDefaultColor(2),
				},
			},
			chart.ContinuousSeries{
				Name:    "Expected Bitrate",
				XValues: []float64{0, pts[len(pts)-1]},
				YValues: []float64{arg.expected, arg.expected},
				Style: chart.Style{
					StrokeColor: chart.GetDefaultColor(3),
				},
			},
		},
		XAxis: chart.XAxis{
			Name:      "Time (seconds)",
			NameStyle: chart.Style{},
			Style: chart.Style{
				TextRotationDegrees: 90,
			},
		},
		YAxis: chart.YAxis{
			Name:      "Bitrate (kbit/s) over " + strconv.Itoa(arg.interval) + " seconds",
			NameStyle: chart.Style{},
			Style:     chart.Style{},
		},
	}

	graph.Width = int(float64(graph.GetWidth()) * 1.5)
	graph.Height = int(float64(graph.GetHeight()) * 1.5)
	graph.Elements = []chart.Renderable{
		chart.Legend(&graph),
	}

	outfile, err := os.Create(arg.outName)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer outfile.Close()
	outType := chart.SVG
	if arg.outType == "png" {
		outType = chart.PNG
	}
	err = graph.Render(outType, outfile)
	if err != nil {
		fmt.Println(err)
		return
	}
}