Multi-model image pipelines

// Detect faces, crop each one without copying pixels, and compare the
// first two with a face embedding model.
//
// One decoded Frame feeds both models; each session prepares the image
// the way its own model expects.

import "stdlib/io" as io;
import "stdlib/env" as env;
import "stdlib/fs" as fs;
import "stdlib/result" as result;
import "stdlib/option" as option;
import "ffi/api" as slab;

fn embed(ref embedder: slab::Session, ref frame: slab::Frame, ref faces: slab::Output, i: i32, ref out: slab::Output) -> bool {
    // A crop is a view into the original pixels, not a copy.
    var face = slab::Frame::new();
    let x = faces.box_x1(i) as i32;
    let y = faces.box_y1(i) as i32;
    let w = (faces.box_x2(i) - faces.box_x1(i)) as i32;
    let h = (faces.box_y2(i) - faces.box_y1(i)) as i32;
    if !frame.crop_into(face, x, y, w, h) { return false; }

    var inp = slab::Input::new();
    inp.set_frame(face);
    if !embedder.run(inp, out) { return false; }
    return out.normalize_l2(0);
}

fn main() -> i32 {
    var cfg = slab::Config::new();
    cfg.set_models_dir("models");
    cfg.set_cache_dir("/tmp/slab_cache");
    guard let option::Option::Some(path) = env::var("SLAB_RUNTIME_LICENSE_FILE") else { return 1; };
    guard let result::Result::Ok(token) = fs::read_to_string(path.view()) else { return 1; };
    if !cfg.set_runtime_license(token.view()) {
        io::eprintln("License rejected: ${cfg.runtime_license_error_message()}");
        return 1;
    }

    var detector = slab::Session::new(cfg);
    var embedder = slab::Session::new(cfg);
    if !detector.open("face_detector.slab") || !embedder.open("face_embedder.slab") {
        io::eprintln("Cannot open the models: ${detector.error_message()} ${embedder.error_message()}");
        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.
    detector.set_resize_mode(slab::RESIZE_LETTERBOX);
    detector.set_letterbox_pad(0);
    detector.set_color_bgr(true);
    detector.set_normalization(1.0, 127.5, 127.5, 127.5, 128.0, 128.0, 128.0);
    embedder.set_normalization(1.0, 127.5, 127.5, 127.5, 127.5, 127.5, 127.5);

    var frame = slab::Frame::new();
    if !frame.load_file("images/group.jpg") {
        io::eprintln("Cannot load the image: ${frame.error_message()}");
        return 1;
    }

    var inp = slab::Input::new();
    var faces = slab::Output::new();
    inp.set_frame(frame);
    if !detector.run(inp, faces) {
        io::eprintln("Detection failed: ${detector.error_message()}");
        return 1;
    }

    // Boxes and landmarks come back in the original image's coordinates.
    let n = faces.decode_boxes(0.5f32, 0.4f32, 32);
    io::println("${n} faces");
    for i in 0..n {
        io::println("  [${faces.box_x1(i)}, ${faces.box_y1(i)}] - [${faces.box_x2(i)}, ${faces.box_y2(i)}] score ${faces.box_score(i)}, ${faces.box_landmark_count(i)} landmarks");
    }
    if n < 2 { return 0; }

    var a = slab::Output::new();
    var b = slab::Output::new();
    if embed(embedder, frame, faces, 0, a) && embed(embedder, frame, faces, 1, b) {
        io::println("Similarity of the first two faces: ${a.cosine_similarity(0, b, 0)}");
    }
    return 0;
}