vqmplot.go

// Copyright ©2022 Evolution. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// ease tool's vqmplot subcommand implementation.
package main

import (
	"flag"
	"fmt"
	"os"
	"path"
	"strings"

	"github.com/evolution-gaming/ease/internal/analysis"
	"github.com/evolution-gaming/ease/internal/logging"
	"github.com/evolution-gaming/ease/internal/vqm"
)

// Support these metrics for plotting.
var supportedMetrics = "VMAF, PSNR, MS-SSIM"

// CreateVQMPlotCommand will create instance of VQMPlotApp.
func CreateVQMPlotCommand() *VQMPlotApp {
	longHelp := `Subcommand "vqmplot" will create plot for given metric from JSON report as
generated by libvmaf.

Examples:

  ease vqmplot -i libvmaf.json -o vmaf.png
  ease vqmplot -m PSNR -i libvmaf.json -o psnr.png`

	app := &VQMPlotApp{
		fs: flag.NewFlagSet("vqmplot", flag.ContinueOnError),
		gf: globalFlags{},
	}
	app.gf.Register(app.fs)
	app.fs.StringVar(&app.flSrcFile, "i", "", "Input libvmaf JSON file (mandatory)")
	app.fs.StringVar(&app.flOutFile, "o", "", "Output file")
	app.fs.StringVar(&app.flMetric, "m", "VMAF", fmt.Sprintf("Metric to plot (%s)", supportedMetrics))
	app.fs.Float64Var(&app.flFPS, "fps", 0, "Source video file FPS")

	app.fs.Usage = func() {
		printSubCommandUsage(longHelp, app.fs)
	}

	return app
}

// VQMPlotApp is vqmplot subcommand context that implements Commander interface.
type VQMPlotApp struct {
	// FlagSet instance
	fs *flag.FlagSet
	// Source file containing per-frame libvmaf JSON data
	flSrcFile string
	// Output file to save plot to
	flOutFile string
	// Selected metric to plot
	flMetric string
	// Video file fps
	flFPS float64
	// Global flags
	gf globalFlags
}

// Run is entry point to VQMPlotApp command execution.
func (a *VQMPlotApp) Run(args []string) error {
	if err := a.fs.Parse(args); err != nil {
		return &AppError{
			exitCode: 2,
			msg:      "usage error",
		}
	}

	if a.gf.Debug {
		logging.EnableDebugLogger()
	}

	// Flag specifying libvmaf JSON metrics is mandatory.
	if a.flSrcFile == "" {
		a.fs.Usage()
		return &AppError{
			exitCode: 2,
			msg:      "mandatory option -i is missing",
		}
	}

	// libvmaf JSON metrics file must exist.
	if _, err := os.Stat(a.flSrcFile); os.IsNotExist(err) {
		return &AppError{
			exitCode: 2,
			msg:      fmt.Sprintf("input file missing: %s", err.Error()),
		}
	}

	// Output file will be constructed if not specified.
	if a.flOutFile == "" {
		base := path.Base(a.flSrcFile)
		base = strings.TrimSuffix(base, path.Ext(base))
		a.flOutFile = base + ".png"
	}

	if !strings.Contains(supportedMetrics, a.flMetric) {
		a.fs.Usage()
		return &AppError{
			exitCode: 2,
			msg:      fmt.Sprintf("unsupported metric, should be one of: %s\n", supportedMetrics),
		}
	}

	if a.flFPS == 0 {
		a.fs.Usage()
		return &AppError{
			exitCode: 2,
			msg:      "mandatory option -fps is missing",
		}
	}

	logging.Info("Starting...")

	jsonFd, err := os.Open(a.flSrcFile)
	if err != nil {
		return &AppError{
			exitCode: 1,
			msg:      err.Error(),
		}
	}
	defer jsonFd.Close()

	var frameMetrics vqm.FrameMetrics
	if err := frameMetrics.FromFfmpegVMAF(jsonFd); err != nil {
		return &AppError{
			exitCode: 1,
			msg:      err.Error(),
		}
	}

	vqms := make(metricXYs, 0, len(frameMetrics))
	switch a.flMetric {
	case "VMAF":
		for _, v := range frameMetrics {
			vqms = append(vqms, metricXY{X: float64(v.FrameNum) / a.flFPS, Y: v.VMAF})
		}
	case "PSNR":
		for _, v := range frameMetrics {
			vqms = append(vqms, metricXY{X: float64(v.FrameNum) / a.flFPS, Y: v.PSNR})
		}
	case "MS-SSIM":
		for _, v := range frameMetrics {
			vqms = append(vqms, metricXY{X: float64(v.FrameNum) / a.flFPS, Y: v.MS_SSIM})
		}
	}
	if len(vqms) == 0 {
		return &AppError{
			exitCode: 1,
			msg:      fmt.Sprintf("no records for %s in %s", a.flMetric, a.flSrcFile),
		}
	}

	if err := analysis.MultiPlotVqm(vqms, a.flMetric, path.Base(a.flSrcFile), a.flOutFile); err != nil {
		return &AppError{
			exitCode: 1,
			msg:      err.Error(),
		}
	}
	logging.Info("Done")
	return nil
}