feature_extractor.h

Go to the documentation of this file.

 
#ifndef __VMAF_FEATURE_EXTRACTOR_H__
#define __VMAF_FEATURE_EXTRACTOR_H__
 
#include <stdatomic.h>
#include <stdint.h>
#include <stdlib.h>
 
#include "dict.h"
#include "framesync.h"
#include "feature_collector.h"
#include "opt.h"
 
#include "libvmaf/picture.h"
 
#ifdef HAVE_CUDA
#include "cuda/common.h"
#endif
 
enum VmafFeatureExtractorFlags {
    VMAF_FEATURE_EXTRACTOR_TEMPORAL = 1 << 0,
    VMAF_FEATURE_EXTRACTOR_CUDA = 1 << 1,
    VMAF_FEATURE_FRAME_SYNC = 1 << 2,
    VMAF_FEATURE_EXTRACTOR_PREV_REF = 1 << 3,
};
 
typedef struct VmafFeatureExtractor {
    const char *name; 
    int (*init)(struct VmafFeatureExtractor *fex, enum VmafPixelFormat pix_fmt,
                unsigned bpc, unsigned w, unsigned h);
    int (*extract)(struct VmafFeatureExtractor *fex,
                   VmafPicture *ref_pic, VmafPicture *ref_pic_90,
                   VmafPicture *dist_pic, VmafPicture *dist_pic_90,
                   unsigned index, VmafFeatureCollector *feature_collector);
    int (*flush)(struct VmafFeatureExtractor *fex,
                 VmafFeatureCollector *feature_collector);
    int (*close)(struct VmafFeatureExtractor *fex);
    const VmafOption *options; 
    void *priv; 
    size_t priv_size; 
    uint64_t flags; 
    const char **provided_features; 
 
    #ifdef HAVE_CUDA
    VmafCudaState *cu_state; 
    #endif
 
    VmafFrameSyncContext *framesync;
    VmafPicture prev_ref; 
    VmafPicture prev_prev_ref; 
 
} VmafFeatureExtractor;
 
VmafFeatureExtractor *vmaf_get_feature_extractor_by_name(const char *name);
VmafFeatureExtractor *vmaf_get_feature_extractor_by_feature_name(const char *name,
                                                                 unsigned flags);
 
enum VmafFeatureExtractorContextFlags {
    VMAF_FEATURE_EXTRACTOR_CONTEXT_DO_NOT_OVERWRITE = 1 << 0,
};
 
typedef struct VmafFeatureExtractorContext {
    bool is_initialized, is_closed;
    VmafDictionary *opts_dict;
    VmafFeatureExtractor *fex;
} VmafFeatureExtractorContext;
 
int vmaf_feature_extractor_context_create(VmafFeatureExtractorContext **fex_ctx,
                                          VmafFeatureExtractor *fex,
                                          VmafDictionary *opts_dict);
 
int vmaf_feature_extractor_context_init(VmafFeatureExtractorContext *fex_ctx,
                                        enum VmafPixelFormat pix_fmt,
                                        unsigned bpc, unsigned w, unsigned h);
 
int vmaf_feature_extractor_context_extract(VmafFeatureExtractorContext *fex_ctx,
                                           VmafPicture *ref, VmafPicture *ref_90,
                                           VmafPicture *dist, VmafPicture *dist_90,
                                           unsigned pic_index,
                                           VmafFeatureCollector *vfc);
 
int vmaf_feature_extractor_context_flush(VmafFeatureExtractorContext *fex_ctx,
                                         VmafFeatureCollector *vfc);
 
int vmaf_feature_extractor_context_close(VmafFeatureExtractorContext *fex_ctx);
 
int vmaf_feature_extractor_context_delete(VmafFeatureExtractorContext *fex_ctx);
 
int vmaf_feature_extractor_context_destroy(VmafFeatureExtractorContext *fex_ctx);
 
typedef struct VmafFeatureExtractorContextPool {
    struct fex_list_entry {
        VmafFeatureExtractor *fex;
        VmafDictionary *opts_dict;
        struct {
            VmafFeatureExtractorContext *fex_ctx;
            bool in_use;
        } *ctx_list;
        atomic_int capacity, in_use;
        pthread_cond_t full;
    } *fex_list;
    unsigned cnt, capacity;
    pthread_mutex_t lock;
    unsigned n_threads;
} VmafFeatureExtractorContextPool;
 
int vmaf_fex_ctx_pool_create(VmafFeatureExtractorContextPool **pool,
                             unsigned n_threads);
 
int vmaf_fex_ctx_pool_aquire(VmafFeatureExtractorContextPool *pool,
                             VmafFeatureExtractor *fex,
                             VmafDictionary *opts_dict,
                             VmafFeatureExtractorContext **fex_ctx);
 
int vmaf_fex_ctx_pool_release(VmafFeatureExtractorContextPool *pool,
                             VmafFeatureExtractorContext *fex_ctx);
 
int vmaf_fex_ctx_pool_flush(VmafFeatureExtractorContextPool *pool,
                            VmafFeatureCollector *feature_collector);
 
int vmaf_fex_ctx_pool_destroy(VmafFeatureExtractorContextPool *pool);
 
#endif /* __VMAF_FEATURE_EXTRACTOR_H__ */