2022-11-13 12:42:03 +00:00
//go:build cgo && !disable_library_libvmaf
package libvmaf
2022-11-11 06:52:44 +00:00
/ *
2022-11-11 10:13:19 +00:00
# cgo pkg - config : libvmaf
2022-11-11 06:52:44 +00:00
# include "libvmaf.h"
2022-11-11 10:13:19 +00:00
# include < stdlib . h >
2022-11-11 06:52:44 +00:00
* /
import "C"
import (
"errors"
"fmt"
"git.gammaspectra.live/S.O.N.G/Ignite/frame"
"unsafe"
)
2022-11-14 11:41:30 +00:00
var vmafVersion = "vmaf " + C . GoString ( C . vmaf_version ( ) )
2022-11-11 06:52:44 +00:00
func Version ( ) string {
2022-11-14 11:41:30 +00:00
return vmafVersion
2022-11-11 06:52:44 +00:00
}
type VMAF struct {
2022-11-12 10:04:17 +00:00
context * C . VmafContext
model * C . VmafModel
2022-11-11 06:52:44 +00:00
}
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
)
2022-11-11 07:15:44 +00:00
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 ) {
2022-11-11 06:52:44 +00:00
v := & VMAF { }
var cfg C . VmafConfiguration
2022-11-11 10:13:19 +00:00
cfg . n_threads = 0 //runtime.NumCPU()
2022-11-11 06:52:44 +00:00
cfg . n_subsample = C . uint ( subsample )
cfg . cpumask = 0
2022-11-11 10:13:19 +00:00
cfg . log_level = C . VMAF_LOG_LEVEL_NONE
2022-11-11 06:52:44 +00:00
if ret := C . vmaf_init ( & v . context , cfg ) ; ret != 0 {
return nil , fmt . Errorf ( "vmaf_init error %d" , ret )
} else {
var modelConfig C . VmafModelConfig
2022-11-11 07:15:44 +00:00
modelConfig . name = C . CString ( model . Name )
2022-11-11 10:13:19 +00:00
defer C . free ( unsafe . Pointer ( modelConfig . name ) )
2022-11-11 07:15:44 +00:00
modelConfig . flags = C . uint64_t ( model . Flags )
2022-11-11 06:52:44 +00:00
2022-11-11 07:15:44 +00:00
modelVersion := C . CString ( model . Version )
2022-11-11 10:13:19 +00:00
defer C . free ( unsafe . Pointer ( modelVersion ) )
2022-11-11 06:52:44 +00:00
if ret = C . vmaf_model_load ( & v . model , & modelConfig , modelVersion ) ; ret != 0 {
v . Close ( )
return nil , fmt . Errorf ( "vmaf_model_load error %d" , ret )
}
2022-11-11 10:13:19 +00:00
if ret = C . vmaf_use_features_from_model ( v . context , v . model ) ; ret != 0 {
2022-11-11 06:52:44 +00:00
return nil , fmt . Errorf ( "vmaf_use_features_from_model error %d" , ret )
}
}
return v , nil
}
func ( v * VMAF ) allocatePicture ( properties frame . Properties ) * C . VmafPicture {
2022-11-13 12:42:03 +00:00
//todo: reuse these pictures
2022-11-11 10:13:19 +00:00
pixFmt := uint32 ( C . VMAF_PIX_FMT_YUV420P )
2022-11-11 06:52:44 +00:00
switch true {
2022-11-14 11:41:30 +00:00
case properties . ColorSpace . ChromaSampling . J == 4 && properties . ColorSpace . ChromaSampling . A == 4 && properties . ColorSpace . ChromaSampling . B == 4 :
2022-11-11 06:52:44 +00:00
pixFmt = C . VMAF_PIX_FMT_YUV444P
2022-11-14 11:41:30 +00:00
case properties . ColorSpace . ChromaSampling . J == 4 && properties . ColorSpace . ChromaSampling . A == 2 && properties . ColorSpace . ChromaSampling . B == 2 :
2022-11-11 06:52:44 +00:00
pixFmt = C . VMAF_PIX_FMT_YUV422P
2022-11-14 11:41:30 +00:00
case properties . ColorSpace . ChromaSampling . J == 4 && properties . ColorSpace . ChromaSampling . A == 2 && properties . ColorSpace . ChromaSampling . B == 0 :
2022-11-11 06:52:44 +00:00
pixFmt = C . VMAF_PIX_FMT_YUV420P
2022-11-14 11:41:30 +00:00
case properties . ColorSpace . ChromaSampling . J == 4 && properties . ColorSpace . ChromaSampling . A == 0 && properties . ColorSpace . ChromaSampling . B == 0 :
2022-11-11 06:52:44 +00:00
pixFmt = C . VMAF_PIX_FMT_YUV400P
default :
return nil
}
var p C . VmafPicture
2022-11-14 11:41:30 +00:00
if ret := C . vmaf_picture_alloc ( & p , pixFmt , C . uint ( properties . ColorSpace . BitDepth ) , C . uint ( properties . Width ) , C . uint ( properties . Height ) ) ; ret != 0 {
2022-11-11 06:52:44 +00:00
return nil
}
return & p
}
func ( v * VMAF ) deallocatePicture ( p * C . VmafPicture ) bool {
if p == nil {
return false
}
2022-11-11 10:13:19 +00:00
return C . vmaf_picture_unref ( p ) == 0
2022-11-11 06:52:44 +00:00
}
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 )
2022-11-11 10:13:19 +00:00
defer C . free ( unsafe . Pointer ( name ) )
2022-11-11 06:52:44 +00:00
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
}
2022-11-11 10:13:19 +00:00
type PoolingMethod uint32
2022-11-11 06:52:44 +00:00
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
2022-11-11 10:13:19 +00:00
if ret := C . vmaf_score_pooled ( v . context , v . model , uint32 ( method ) , ( * C . double ) ( & score ) , C . uint ( indexLow ) , C . uint ( indexHigh ) ) ; ret != 0 {
2022-11-11 06:52:44 +00:00
return 0 , fmt . Errorf ( "vmaf_score_at_index error %d" , ret )
}
return score , nil
}
2022-11-11 07:15:44 +00:00
// 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 ) {
2022-11-11 06:52:44 +00:00
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 {
2022-11-11 07:15:44 +00:00
return index , err
2022-11-11 06:52:44 +00:00
}
2022-11-11 07:15:44 +00:00
return index , nil
2022-11-11 06:52:44 +00:00
} else if ! refOk {
2022-11-11 07:15:44 +00:00
return index , errors . New ( "reference ended before distorted" )
2022-11-11 06:52:44 +00:00
} else if ! distOk {
2022-11-11 07:15:44 +00:00
return index , errors . New ( "distorted ended before reference" )
2022-11-11 06:52:44 +00:00
} else {
if err = v . ReadPictures ( refFrame , distFrame , index ) ; err != nil {
2022-11-11 07:15:44 +00:00
return index , err
2022-11-11 06:52:44 +00:00
}
}
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 {
2022-11-11 07:15:44 +00:00
return fmt . Errorf ( "vmaf_read_pictures error %d" , ret )
2022-11-11 06:52:44 +00:00
} else {
return nil
}
}
ref := v . frameToPicture ( reference )
dist := v . frameToPicture ( distorted )
2022-11-12 10:04:17 +00:00
defer v . deallocatePicture ( ref )
defer v . deallocatePicture ( dist )
2022-11-11 06:52:44 +00:00
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 {
2022-11-11 07:15:44 +00:00
return fmt . Errorf ( "vmaf_read_pictures error %d" , ret )
2022-11-11 06:52:44 +00:00
} 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
}
}
2022-11-11 10:13:19 +00:00
type OutputFormat uint32
2022-11-11 06:52:44 +00:00
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 )
2022-11-11 10:13:19 +00:00
defer C . free ( unsafe . Pointer ( cPath ) )
if ret := C . vmaf_write_output ( v . context , cPath , uint32 ( format ) ) ; ret != 0 {
2022-11-11 07:15:44 +00:00
return fmt . Errorf ( "vmaf_write_output error %d" , ret )
2022-11-11 06:52:44 +00:00
}
return nil
}