Multi-model image pipelines

/* Detect faces, crop each one without copying pixels, and compare the
 * first two with a face embedding model.
 *
 * usage: faces DETECTOR.slab EMBEDDER.slab IMAGE LICENSE_FILE
 */

#include <stdio.h>
#include <string.h>

#include "slab.h"

#define S(str) (uint8_t *)(str), strlen(str)

static int read_file(const char *path, char *buf, size_t cap, size_t *len) {
    FILE *f = fopen(path, "rb");
    if (!f) return 0;
    *len = fread(buf, 1, cap, f);
    fclose(f);
    while (*len && (buf[*len - 1] == '\n' || buf[*len - 1] == '\r')) (*len)--;
    return 1;
}

/* Crop face `i` from `frame` (a view, no copy), embed it into `out`. */
static int embed(uint8_t *embedder, uint8_t *frame, uint8_t *faces, int32_t i, uint8_t *out) {
    float x1 = slab_output_box_x1(faces, i), y1 = slab_output_box_y1(faces, i);
    float x2 = slab_output_box_x2(faces, i), y2 = slab_output_box_y2(faces, i);
    uint8_t *face = slab_frame_new();
    uint8_t *in = slab_input_new();
    int ok = slab_frame_crop_into(frame, face, (int32_t)x1, (int32_t)y1, (int32_t)(x2 - x1), (int32_t)(y2 - y1));
    if (ok) {
        slab_input_set_frame(in, face);
        ok = slab_session_run(embedder, in, out) && slab_output_normalize_l2(out, 0);
    }
    slab_input_free(in);
    slab_frame_free(face);
    return ok;
}

int main(int argc, char **argv) {
    if (argc < 5) { fprintf(stderr, "usage: %s DETECTOR.slab EMBEDDER.slab IMAGE LICENSE_FILE\n", argv[0]); return 2; }

    uint8_t *cfg = slab_config_new();
    char token[4096]; size_t token_len;
    if (!read_file(argv[4], token, sizeof token, &token_len) ||
        !slab_config_set_runtime_license(cfg, (uint8_t *)token, token_len)) {
        fprintf(stderr, "License rejected\n");
        return 1;
    }

    uint8_t *detector = slab_session_new(cfg);
    uint8_t *embedder = slab_session_new(cfg);
    if (!slab_session_open(detector, S(argv[1])) || !slab_session_open(embedder, S(argv[2]))) {
        fprintf(stderr, "Cannot open the models\n");
        return 1;
    }

    /* Preprocessing normally comes from the model file. Set it in code only
     * when a model file does not carry it, or to override it. */
    slab_session_set_resize_mode(detector, SLAB_RESIZE_LETTERBOX);
    slab_session_set_letterbox_pad(detector, 0);
    slab_session_set_color_bgr(detector, 1);
    slab_session_set_normalization(detector, 1.0f, 127.5f, 127.5f, 127.5f, 128.0f, 128.0f, 128.0f);
    slab_session_set_normalization(embedder, 1.0f, 127.5f, 127.5f, 127.5f, 127.5f, 127.5f, 127.5f);

    uint8_t *frame = slab_frame_new();
    if (!slab_frame_load_file(frame, S(argv[3]))) { fprintf(stderr, "Cannot load the image\n"); return 1; }

    uint8_t *in = slab_input_new();
    uint8_t *faces = slab_output_new();
    slab_input_set_frame(in, frame);
    if (!slab_session_run(detector, in, faces)) { fprintf(stderr, "Detection failed\n"); return 1; }

    /* Boxes and landmarks come back in the original image's pixels. */
    int32_t n = slab_output_decode_boxes(faces, 0.5f, 0.4f, 32);
    printf("%d faces\n", n);
    for (int32_t i = 0; i < n; i++) {
        printf("  [%.0f, %.0f] - [%.0f, %.0f] score %.3f, %d landmarks\n",
               slab_output_box_x1(faces, i), slab_output_box_y1(faces, i),
               slab_output_box_x2(faces, i), slab_output_box_y2(faces, i),
               slab_output_box_score(faces, i), slab_output_box_landmark_count(faces, i));
    }

    if (n >= 2) {
        uint8_t *a = slab_output_new(), *b = slab_output_new();
        if (embed(embedder, frame, faces, 0, a) && embed(embedder, frame, faces, 1, b))
            printf("Similarity of the first two faces: %.4f\n", slab_output_cosine_similarity(a, 0, b, 0));
        slab_output_free(a);
        slab_output_free(b);
    }

    slab_output_free(faces);
    slab_input_free(in);
    slab_frame_free(frame);
    slab_session_free(embedder);
    slab_session_free(detector);
    slab_config_free(cfg);
    return 0;
}