This commit is contained in:
parent
fe40f4da7e
commit
60a6306a43
270
utilities/vmaf/vmaf.go
Normal file
270
utilities/vmaf/vmaf.go
Normal file
|
@ -0,0 +1,270 @@
|
|||
package vmaf
|
||||
|
||||
/*
|
||||
#cgo pkg-config: vmaf
|
||||
#include "libvmaf.h"
|
||||
|
||||
*/
|
||||
import "C"
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"git.gammaspectra.live/S.O.N.G/Ignite/frame"
|
||||
"runtime"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
func Version() string {
|
||||
return "vmaf " + C.GoString(C.vmaf_version())
|
||||
}
|
||||
|
||||
type VMAF struct {
|
||||
context *C.VmafContext
|
||||
model *C.VmafModel
|
||||
pictures []*C.VmafPicture
|
||||
}
|
||||
|
||||
type ModelName string
|
||||
|
||||
const (
|
||||
ModelA ModelName = "test"
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
func New(subsample uint, model ModelName, flags ModelFlags, version string) (*VMAF, error) {
|
||||
v := &VMAF{}
|
||||
|
||||
var cfg C.VmafConfiguration
|
||||
cfg.n_threads = runtime.NumCPU()
|
||||
cfg.n_subsample = C.uint(subsample)
|
||||
cfg.cpumask = 0
|
||||
cfg.log_level = C.VMAF_LOG_LEVEL_DEBUG
|
||||
|
||||
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)
|
||||
defer C.free(modelConfig.name)
|
||||
modelConfig.flags = C.uint64_t(flags)
|
||||
|
||||
modelVersion := C.CString(version)
|
||||
defer C.free(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 {
|
||||
//maybe calloc
|
||||
pixFmt := C.VMAF_PIX_FMT_YUV420P
|
||||
switch true {
|
||||
case properties.ColorFormat.Subsampling.J == 4 && properties.ColorFormat.Subsampling.A == 4 && properties.ColorFormat.Subsampling.B == 4:
|
||||
pixFmt = C.VMAF_PIX_FMT_YUV444P
|
||||
case properties.ColorFormat.Subsampling.J == 4 && properties.ColorFormat.Subsampling.A == 2 && properties.ColorFormat.Subsampling.B == 2:
|
||||
pixFmt = C.VMAF_PIX_FMT_YUV422P
|
||||
case properties.ColorFormat.Subsampling.J == 4 && properties.ColorFormat.Subsampling.A == 2 && properties.ColorFormat.Subsampling.B == 0:
|
||||
pixFmt = C.VMAF_PIX_FMT_YUV420P
|
||||
case properties.ColorFormat.Subsampling.J == 4 && properties.ColorFormat.Subsampling.A == 0 && properties.ColorFormat.Subsampling.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.ColorFormat.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(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 int
|
||||
|
||||
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, 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.
|
||||
func (v *VMAF) ReadStreams(reference, distorted *frame.Stream) 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 err
|
||||
}
|
||||
return nil
|
||||
} else if !refOk {
|
||||
return errors.New("reference ended before distorted")
|
||||
} else if !distOk {
|
||||
return errors.New("distorted ended before reference")
|
||||
} else {
|
||||
if err = v.ReadPictures(refFrame, distFrame, index); err != nil {
|
||||
return 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_use_features_from_model error %d", ret)
|
||||
} else {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
ref := v.frameToPicture(reference)
|
||||
dist := v.frameToPicture(distorted)
|
||||
if ref == nil || dist == nil {
|
||||
v.deallocatePicture(ref)
|
||||
v.deallocatePicture(dist)
|
||||
return errors.New("could not allocate pictures")
|
||||
}
|
||||
if ret := C.vmaf_read_pictures(v.context, ref, dist, C.uint(index)); ret != 0 {
|
||||
v.deallocatePicture(ref)
|
||||
v.deallocatePicture(dist)
|
||||
return fmt.Errorf("vmaf_use_features_from_model error %d", ret)
|
||||
} else {
|
||||
//save memory
|
||||
v.pictures = append(v.pictures, ref)
|
||||
v.pictures = append(v.pictures, dist)
|
||||
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 int
|
||||
|
||||
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(cPath)
|
||||
if ret := C.vmaf_write_output(v.context, cPath, format); ret != 0 {
|
||||
return fmt.Errorf("vmaf_use_features_from_model error %d", ret)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
Loading…
Reference in a new issue