Ignite/utilities/libvmaf/libvmaf.go

//go:build cgo && !disable_library_libvmaf

package libvmaf

/*
#cgo pkg-config: libvmaf
#include "libvmaf.h"
#include <stdlib.h>

*/
import "C"
import (
	"errors"
	"fmt"
	"git.gammaspectra.live/S.O.N.G/Ignite/frame"
	"unsafe"
)

var vmafVersion = "vmaf " + C.GoString(C.vmaf_version())

func Version() string {
	return vmafVersion
}

type VMAF struct {
	context *C.VmafContext
	model   *C.VmafModel
}

type ModelFlags uint64

const (
	ModelFlagsDefault         ModelFlags = C.VMAF_MODEL_FLAGS_DEFAULT
	ModelFlagDisableClip      ModelFlags = C.VMAF_MODEL_FLAG_DISABLE_CLIP
	ModelFlagEnableTransform  ModelFlags = C.VMAF_MODEL_FLAG_ENABLE_TRANSFORM
	ModelFlagDisableTransform ModelFlags = C.VMAF_MODEL_FLAG_DISABLE_TRANSFORM
)

type ModelConfiguration struct {
	Name    string
	Version string
	Flags   ModelFlags
}

func NewModelDefault() *ModelConfiguration {
	return &ModelConfiguration{
		Name:    "vmaf",
		Version: "vmaf_v0.6.1",
		Flags:   ModelFlagsDefault,
	}
}

func New(subsample uint, model *ModelConfiguration) (*VMAF, error) {
	v := &VMAF{}

	var cfg C.VmafConfiguration
	cfg.n_threads = 0 //runtime.NumCPU()
	cfg.n_subsample = C.uint(subsample)
	cfg.cpumask = 0
	cfg.log_level = C.VMAF_LOG_LEVEL_NONE

	if ret := C.vmaf_init(&v.context, cfg); ret != 0 {
		return nil, fmt.Errorf("vmaf_init error %d", ret)
	} else {
		var modelConfig C.VmafModelConfig
		modelConfig.name = C.CString(model.Name)
		defer C.free(unsafe.Pointer(modelConfig.name))
		modelConfig.flags = C.uint64_t(model.Flags)

		modelVersion := C.CString(model.Version)
		defer C.free(unsafe.Pointer(modelVersion))

		if ret = C.vmaf_model_load(&v.model, &modelConfig, modelVersion); ret != 0 {
			v.Close()
			return nil, fmt.Errorf("vmaf_model_load error %d", ret)
		}

		if ret = C.vmaf_use_features_from_model(v.context, v.model); ret != 0 {
			return nil, fmt.Errorf("vmaf_use_features_from_model error %d", ret)
		}
	}

	return v, nil
}

func (v *VMAF) allocatePicture(properties frame.Properties) *C.VmafPicture {
	//todo: reuse these pictures
	pixFmt := uint32(C.VMAF_PIX_FMT_YUV420P)
	switch true {
	case properties.ColorSpace.ChromaSampling.J == 4 && properties.ColorSpace.ChromaSampling.A == 4 && properties.ColorSpace.ChromaSampling.B == 4:
		pixFmt = C.VMAF_PIX_FMT_YUV444P
	case properties.ColorSpace.ChromaSampling.J == 4 && properties.ColorSpace.ChromaSampling.A == 2 && properties.ColorSpace.ChromaSampling.B == 2:
		pixFmt = C.VMAF_PIX_FMT_YUV422P
	case properties.ColorSpace.ChromaSampling.J == 4 && properties.ColorSpace.ChromaSampling.A == 2 && properties.ColorSpace.ChromaSampling.B == 0:
		pixFmt = C.VMAF_PIX_FMT_YUV420P
	case properties.ColorSpace.ChromaSampling.J == 4 && properties.ColorSpace.ChromaSampling.A == 0 && properties.ColorSpace.ChromaSampling.B == 0:
		pixFmt = C.VMAF_PIX_FMT_YUV400P
	default:
		return nil
	}

	var p C.VmafPicture
	if ret := C.vmaf_picture_alloc(&p, pixFmt, C.uint(properties.ColorSpace.BitDepth), C.uint(properties.Width), C.uint(properties.Height)); ret != 0 {
		return nil
	}

	return &p
}

func (v *VMAF) deallocatePicture(p *C.VmafPicture) bool {
	if p == nil {
		return false
	}
	return C.vmaf_picture_unref(p) == 0
}

func (v *VMAF) frameToPicture(f frame.Frame) *C.VmafPicture {
	if f == nil {
		return nil
	}
	if p := v.allocatePicture(f.Properties()); p == nil {
		return nil
	} else {
		//TODO: check validity of lengths
		if f16, ok := f.(frame.TypedFrame[uint16]); ok {
			yPlane := unsafe.Slice((*uint16)(p.data[0]), len(f16.GetNativeLuma()))
			uPlane := unsafe.Slice((*uint16)(p.data[1]), len(f16.GetNativeCb()))
			vPlane := unsafe.Slice((*uint16)(p.data[2]), len(f16.GetNativeCr()))
			copy(yPlane, f16.GetNativeLuma())
			copy(uPlane, f16.GetNativeCb())
			copy(vPlane, f16.GetNativeCr())
		} else if f8, ok := f.(frame.TypedFrame[uint8]); ok {
			yPlane := unsafe.Slice((*uint8)(p.data[0]), len(f8.GetNativeLuma()))
			uPlane := unsafe.Slice((*uint8)(p.data[1]), len(f8.GetNativeCb()))
			vPlane := unsafe.Slice((*uint8)(p.data[2]), len(f8.GetNativeCr()))
			copy(yPlane, f8.GetNativeLuma())
			copy(uPlane, f8.GetNativeCb())
			copy(vPlane, f8.GetNativeCr())
		} else {
			// not supported frame
			v.deallocatePicture(p)
			return nil
		}

		return p
	}
}

func (v *VMAF) ScoreAtIndex(index uint) (float64, error) {
	var score float64
	if ret := C.vmaf_score_at_index(v.context, v.model, (*C.double)(&score), C.uint(index)); ret != 0 {
		return 0, fmt.Errorf("vmaf_score_at_index error %d", ret)
	}

	return score, nil
}

func (v *VMAF) FeatureScoreAtIndex(featureName string, index uint) (float64, error) {
	name := C.CString(featureName)
	defer C.free(unsafe.Pointer(name))
	var score float64
	if ret := C.vmaf_feature_score_at_index(v.context, name, (*C.double)(&score), C.uint(index)); ret != 0 {
		return 0, fmt.Errorf("vmaf_feature_score_at_index error %d", ret)
	}

	return score, nil
}

type PoolingMethod uint32

const (
	PoolingMethodUnknown      PoolingMethod = C.VMAF_POOL_METHOD_UNKNOWN
	PoolingMethodMinimum      PoolingMethod = C.VMAF_POOL_METHOD_MIN
	PoolingMethodMaximum      PoolingMethod = C.VMAF_POOL_METHOD_MAX
	PoolingMethodMean         PoolingMethod = C.VMAF_POOL_METHOD_MEAN
	PoolingMethodHarmonicMean PoolingMethod = C.VMAF_POOL_METHOD_HARMONIC_MEAN
	PoolingMethodNB           PoolingMethod = C.VMAF_POOL_METHOD_NB
)

func (v *VMAF) ScorePooled(method PoolingMethod, indexLow, indexHigh uint) (float64, error) {
	var score float64
	if ret := C.vmaf_score_pooled(v.context, v.model, uint32(method), (*C.double)(&score), C.uint(indexLow), C.uint(indexHigh)); ret != 0 {
		return 0, fmt.Errorf("vmaf_score_at_index error %d", ret)
	}

	return score, nil
}

// ReadStreams Read a pair of frame.Stream and queues each frame.Frame via ReadPictures for eventual feature extraction, and flushes at the end. Returns the count of frame.Frame used (indexes will be up to count - 1).
func (v *VMAF) ReadStreams(reference, distorted *frame.Stream) (uint, error) {
	var index uint

	var refFrame, distFrame frame.Frame
	var refOk, distOk bool
	var err error

	refChannel, distChannel := reference.Channel(), distorted.Channel()
	for {
		refFrame, refOk = <-refChannel
		distFrame, distOk = <-distChannel
		if !refOk && !distOk {
			//flush
			if err = v.ReadPictures(nil, nil, index); err != nil {
				return index, err
			}
			return index, nil
		} else if !refOk {
			return index, errors.New("reference ended before distorted")
		} else if !distOk {
			return index, errors.New("distorted ended before reference")
		} else {
			if err = v.ReadPictures(refFrame, distFrame, index); err != nil {
				return index, err
			}
		}

		index++
	}
}

// ReadPictures Read a pair of frame.Frame and queue them for eventual feature extraction.
// When you're done reading pictures call this function again with both reference and distorted set to nil
func (v *VMAF) ReadPictures(reference, distorted frame.Frame, index uint) error {
	if reference == nil && distorted == nil {
		//flush
		if ret := C.vmaf_read_pictures(v.context, nil, nil, C.uint(index)); ret != 0 {
			return fmt.Errorf("vmaf_read_pictures error %d", ret)
		} else {
			return nil
		}
	}

	ref := v.frameToPicture(reference)
	dist := v.frameToPicture(distorted)
	defer v.deallocatePicture(ref)
	defer v.deallocatePicture(dist)

	if ref == nil || dist == nil {
		return errors.New("could not allocate pictures")
	}
	if ret := C.vmaf_read_pictures(v.context, ref, dist, C.uint(index)); ret != 0 {
		return fmt.Errorf("vmaf_read_pictures error %d", ret)
	} else {
		return nil
	}
}

func (v *VMAF) Close() {
	if v.model != nil {
		C.vmaf_model_destroy(v.model)
		v.model = nil
	}
	if v.context != nil {
		C.vmaf_close(v.context)
		v.context = nil
	}
}

type OutputFormat uint32

const (
	OutputFormatNone OutputFormat = C.VMAF_OUTPUT_FORMAT_NONE
	OutputFormatXml  OutputFormat = C.VMAF_OUTPUT_FORMAT_XML
	OutputFormatJson OutputFormat = C.VMAF_OUTPUT_FORMAT_JSON
	OutputFormatCsv  OutputFormat = C.VMAF_OUTPUT_FORMAT_CSV
	OutputFormatSub  OutputFormat = C.VMAF_OUTPUT_FORMAT_SUB
)

func (v *VMAF) Output(path string, format OutputFormat) error {
	cPath := C.CString(path)
	defer C.free(unsafe.Pointer(cPath))
	if ret := C.vmaf_write_output(v.context, cPath, uint32(format)); ret != 0 {
		return fmt.Errorf("vmaf_write_output error %d", ret)
	}

	return nil
}