File size: 1,965 Bytes
1f71164 5ff8f6e 1f71164 5ff8f6e 1f71164 5ff8f6e 1f71164 5ff8f6e 1f71164 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | """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()
|