Working VMAF implementation

utilities/vmaf/vmaf.go

@@ -1,8 +1,9 @@
 package vmaf
 
 /*
-#cgo pkg-config: vmaf
+#cgo pkg-config: libvmaf
 #include "libvmaf.h"
+#include <stdlib.h>
 
 */
 import "C"
@@ -10,7 +11,6 @@ import (
 	"errors"
 	"fmt"
 	"git.gammaspectra.live/S.O.N.G/Ignite/frame"
-	"runtime"
 	"unsafe"
 )
 
@@ -51,28 +51,28 @@ func New(subsample uint, model *ModelConfiguration) (*VMAF, error) {
 	v := &VMAF{}
 
 	var cfg C.VmafConfiguration
-	cfg.n_threads = runtime.NumCPU()
+	cfg.n_threads = 0 //runtime.NumCPU()
 	cfg.n_subsample = C.uint(subsample)
 	cfg.cpumask = 0
-	cfg.log_level = C.VMAF_LOG_LEVEL_DEBUG
+	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(modelConfig.name)
+		defer C.free(unsafe.Pointer(modelConfig.name))
 		modelConfig.flags = C.uint64_t(model.Flags)
 
 		modelVersion := C.CString(model.Version)
-		defer C.free(modelVersion)
+		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 {
+		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)
 		}
 	}
@@ -82,7 +82,7 @@ func New(subsample uint, model *ModelConfiguration) (*VMAF, error) {
 
 func (v *VMAF) allocatePicture(properties frame.Properties) *C.VmafPicture {
 	//maybe calloc
-	pixFmt := C.VMAF_PIX_FMT_YUV420P
+	pixFmt := uint32(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
@@ -108,7 +108,7 @@ func (v *VMAF) deallocatePicture(p *C.VmafPicture) bool {
 	if p == nil {
 		return false
 	}
-	return C.vmaf_picture_unref(&p) == 0
+	return C.vmaf_picture_unref(p) == 0
 }
 
 func (v *VMAF) frameToPicture(f frame.Frame) *C.VmafPicture {
@@ -154,7 +154,7 @@ func (v *VMAF) ScoreAtIndex(index uint) (float64, error) {
 
 func (v *VMAF) FeatureScoreAtIndex(featureName string, index uint) (float64, error) {
 	name := C.CString(featureName)
-	defer C.free(name)
+	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)
@@ -163,7 +163,7 @@ func (v *VMAF) FeatureScoreAtIndex(featureName string, index uint) (float64, err
 	return score, nil
 }
 
-type PoolingMethod int
+type PoolingMethod uint32
 
 const (
 	PoolingMethodUnknown      PoolingMethod = C.VMAF_POOL_METHOD_UNKNOWN
@@ -176,7 +176,7 @@ const (
 
 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 {
+	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)
 	}
 
@@ -257,7 +257,7 @@ func (v *VMAF) Close() {
 	}
 }
 
-type OutputFormat int
+type OutputFormat uint32
 
 const (
 	OutputFormatNone OutputFormat = C.VMAF_OUTPUT_FORMAT_NONE
@@ -269,8 +269,8 @@ const (
 
 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 {
+	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)
 	}
 

utilities/vmaf/vmaf_test.go

@@ -0,0 +1,63 @@
+package vmaf
+
+import (
+	"git.gammaspectra.live/S.O.N.G/Ignite/decoder/libdav1d"
+	"git.gammaspectra.live/S.O.N.G/Ignite/decoder/y4m"
+	"git.gammaspectra.live/S.O.N.G/Ignite/testdata"
+	"os"
+	"testing"
+)
+
+func TestVersion(t *testing.T) {
+	t.Logf("vmaf version: %s", Version())
+}
+
+func TestDecodeYUV420_8bit(t *testing.T) {
+	referenceFile, err := os.Open(testdata.Y4M_Sintel_Trailer_720p24_YUV420_8bit)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer referenceFile.Close()
+	distortedFile, err := os.Open(testdata.AV1_Sintel_Trailer_720p24_YUV420_8bit_Low)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer distortedFile.Close()
+
+	referenceDecoder, err := y4m.New(referenceFile, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	//defer referenceDecoder.Close()
+
+	distortedDecoder, err := libdav1d.NewDecoder(distortedFile, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer distortedDecoder.Close()
+
+	if vmaf, err := New(0, NewModelDefault()); err != nil {
+		t.Fatal(err)
+	} else {
+		frameCount, err := vmaf.ReadStreams(referenceDecoder.DecodeStream(), distortedDecoder.DecodeStream())
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if frameCount != 1253 {
+			t.Fatalf("expected %d frames, got %d", 1253, frameCount)
+		}
+
+		score, err := vmaf.ScoreAtIndex(800)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Logf("frame 800 score: %f", score)
+
+		score, err = vmaf.ScorePooled(PoolingMethodMean, 0, frameCount-1)
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Logf("mean score: %f", score)
+	}
+}