examples/embed.py

"""Encode an image, or run a synthetic smoke check when no image is supplied."""

import argparse
from pathlib import Path

import coremltools as ct
import numpy as np
from PIL import Image, ImageOps


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("image", type=Path, nargs="?")
    parser.add_argument("--model", type=Path, default=Path("models/DINOv3ViTB16-FP32-448.mlpackage"))
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()
    model = ct.models.MLModel(str(args.model), compute_units=ct.ComputeUnit.ALL)
    image_type = model.get_spec().description.input[0].type.imageType
    size = (image_type.width, image_type.height)
    if args.image:
        with Image.open(args.image) as source:
            image = ImageOps.exif_transpose(source).convert("RGB").resize(size, Image.Resampling.BICUBIC)
    else:
        y, x = np.mgrid[0:size[1], 0:size[0]]
        pixels = np.stack((x * 255 // size[0], y * 255 // size[1], (x + y) * 255 // sum(size)), axis=-1).astype(np.uint8)
        image = Image.fromarray(pixels)
    prediction = model.predict({"image": image})
    vector = np.asarray(prediction["embedding"], dtype=np.float32).reshape(-1)
    patches = np.asarray(prediction["patch_embeddings"], dtype=np.float32).reshape(-1, 768)
    if vector.shape != (768,) or not np.isfinite(vector).all():
        raise SystemExit("Invalid embedding")
    if not np.isfinite(patches).all():
        raise SystemExit("Invalid patch embeddings")
    norm = float(np.linalg.norm(vector))
    if not np.isclose(norm, 1, atol=1e-4):
        raise SystemExit(f"Embedding is not normalized: {norm}")
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        np.save(args.output, vector)
    print(f"shape={vector.shape}, dtype={vector.dtype}, L2 norm={norm:.8f}")
    print(f"patches={patches.shape} (unnormalized)")


if __name__ == "__main__":
    main()