Tiny-AI extractor template¶

/**
 *  Copyright 2026 <author>
 *  SPDX-License-Identifier: BSD-3-Clause-Plus-Patent
 */

#include <assert.h>
#include <errno.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include "libvmaf/dnn.h"
#include "libvmaf/picture.h"

#include "config.h"
#include "feature_collector.h"
#include "feature_extractor.h"
#include "log.h"
#include "mem.h"
#include "opt.h"

#include "dnn/tensor_io.h"
#include "dnn/tiny_extractor_template.h"

typedef struct MyExtractorState {
    char *model_path; /* feature option, owned by opt.c */
    VmafDnnSession *sess;
    unsigned w, h;
    uint8_t *rgb8[3];   /* per-channel uint8 RGB scratch */
    float *tensor_in;   /* 3 * w * h floats, NCHW */
    float *out_buf;     /* model output buffer */
} MyExtractorState;

static void my_release(MyExtractorState *s)
{
    if (!s)
        return;
    for (int i = 0; i < 3; ++i) {
        if (s->rgb8[i]) {
            aligned_free(s->rgb8[i]);
            s->rgb8[i] = NULL;
        }
    }
    if (s->tensor_in) {
        aligned_free(s->tensor_in);
        s->tensor_in = NULL;
    }
    if (s->out_buf) {
        aligned_free(s->out_buf);
        s->out_buf = NULL;
    }
    if (s->sess) {
        (void)vmaf_dnn_session_close(s->sess);
        s->sess = NULL;
    }
}

static int my_init(VmafFeatureExtractor *fex, enum VmafPixelFormat pix_fmt, unsigned bpc,
                   unsigned w, unsigned h)
{
    MyExtractorState *s = fex->priv;
    if (pix_fmt == VMAF_PIX_FMT_YUV400P || bpc != 8)
        return -ENOTSUP;

    const char *path =
        vmaf_tiny_ai_resolve_model_path("my_extractor", s->model_path, "VMAF_MY_MODEL_PATH");
    if (!path)
        return -EINVAL;

    int rc = vmaf_tiny_ai_open_session("my_extractor", path, &s->sess);
    if (rc < 0)
        return rc;

    s->w = w;
    s->h = h;
    const size_t plane = (size_t)w * (size_t)h;
    for (int i = 0; i < 3; ++i) {
        s->rgb8[i] = (uint8_t *)aligned_malloc(plane, 32);
        if (!s->rgb8[i])
            goto oom;
    }
    s->tensor_in = (float *)aligned_malloc(3u * plane * sizeof(float), 32);
    s->out_buf = (float *)aligned_malloc(plane * sizeof(float), 32);
    if (!s->tensor_in || !s->out_buf)
        goto oom;
    return 0;
oom:
    my_release(s);
    return -ENOMEM;
}

static int my_extract(VmafFeatureExtractor *fex, VmafPicture *ref, VmafPicture *ref90,
                      VmafPicture *dist, VmafPicture *dist90, unsigned index,
                      VmafFeatureCollector *fc)
{
    (void)ref90;
    (void)dist90;
    MyExtractorState *s = fex->priv;
    int rc = vmaf_tiny_ai_yuv_to_rgb8_planes(dist, s->rgb8[0], s->rgb8[1], s->rgb8[2]);
    if (rc < 0)
        return rc;
    rc = vmaf_tensor_from_rgb_imagenet(s->rgb8[0], s->w, s->rgb8[1], s->w, s->rgb8[2], s->w,
                                       (int)s->w, (int)s->h, s->tensor_in);
    if (rc < 0)
        return rc;
    const int64_t shape[4] = {1, 3, (int64_t)s->h, (int64_t)s->w};
    const size_t plane = (size_t)s->w * (size_t)s->h;
    VmafDnnInput inputs[1] = {
        {.name = "input", .data = s->tensor_in, .shape = shape, .rank = 4u},
    };
    VmafDnnOutput outputs[1] = {
        {.name = "score", .data = s->out_buf, .capacity = plane, .written = 0u},
    };
    rc = vmaf_dnn_session_run(s->sess, inputs, 1u, outputs, 1u);
    if (rc < 0)
        return rc;
    /* derive scalar feature from outputs[0].data + outputs[0].written ... */
    return vmaf_feature_collector_append(fc, "my_score", /* derived */ 0.0, index);
}

static int my_close(VmafFeatureExtractor *fex)
{
    MyExtractorState *s = fex->priv;
    if (!s)
        return 0;
    my_release(s);
    memset(s, 0, sizeof(*s));
    return 0;
}

static const VmafOption my_options[] = {
    VMAF_TINY_AI_MODEL_PATH_OPTION(MyExtractorState,
        "Filesystem path to the my-extractor ONNX model. "
        "Overrides the VMAF_MY_MODEL_PATH env var."),
    {NULL},
};

static const char *my_provided[] = {"my_score", NULL};

VmafFeatureExtractor vmaf_fex_my = {
    .name = "my_extractor",
    .init = my_init,
    .extract = my_extract,
    .close = my_close,
    .options = my_options,
    .priv_size = sizeof(MyExtractorState),
    .provided_features = my_provided,
    .chars = {0},
};