Multi-model image pipelines

"""Detect faces, crop each one without copying pixels, and compare the
first two with a face embedding model.

usage: python faces.py DETECTOR.slab EMBEDDER.slab IMAGE LICENSE_FILE
"""

import sys

import slab

detector_path, embedder_path, image, license_file = sys.argv[1:5]

cfg = slab.Config(license=open(license_file).read())
detector = slab.Session(cfg, detector_path)
embedder = slab.Session(cfg, embedder_path)

# 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(slab.Resize.LETTERBOX, letterbox_pad=0)
detector.set_color_bgr(True)
detector.set_normalization(1.0, mean=(127.5,) * 3, std=(128.0,) * 3)
embedder.set_normalization(1.0, mean=(127.5,) * 3, std=(127.5,) * 3)

frame = slab.Frame.open(image)
faces = detector.run(slab.Input(frame=frame)).boxes(score_threshold=0.5, iou_threshold=0.4, max_boxes=32)
print(f"{len(faces)} faces")
for f in faces:
    print(f"  [{f.x1:.0f}, {f.y1:.0f}, {f.width:.0f}x{f.height:.0f}] score {f.score:.3f}, {len(f.landmarks)} landmarks")


def embed(face):
    out = embedder.run(slab.Input(frame=frame.crop(face.x1, face.y1, face.width, face.height)))
    out.normalize_l2()
    return out


if len(faces) >= 2:
    print(f"Similarity of the first two faces: {embed(faces[0]).cosine_similarity(embed(faces[1])):.4f}")