[0.1.32] macOS phase 5: Mask2Former segmentation on the GPU through native Core ML

README.md

@@ -28,11 +28,12 @@ see [Screen capture](#screen-capture).
    place (`C:\Program Files\NVIDIA\CUDNN`) is found without changing PATH. The OCR window says
    what, if anything, is missing.
 
-   **On macOS**, labelling and OCR run on the GPU through Core ML, from Core ML copies of the models
+   **On macOS**, labelling, OCR and segmentation run on the GPU through Core ML, from Core ML copies
-   made once with `tools/export_siglip_coreml.py` (SigLIP 2's image model, about ten times as fast as
+   of the models made once with `tools/export_siglip_coreml.py` (SigLIP 2's image model, about ten
-   the CPU) and `tools/export_ppocr_coreml.py` (PP-OCRv5's detector and recognizer, about twice as
+   times as fast as the CPU), `tools/export_ppocr_coreml.py` (PP-OCRv5's detector and recognizer,
-   fast; run it with `--model mobile` and `--model server`; their first lines say how to set up the
+   about twice as fast; run it with `--model mobile` and `--model server`) and
-   Python they need). Without them, each runs on the
+   `tools/export_mask2former_coreml.py` (Mask2Former, about four times as fast; after
+   `tools/export_mask2former.py`). Their first lines say how to set up the Python they need. Without them, each runs on the
    CPU. The first scan after making them is slower: Core ML prepares each image size the first time it
    meets it, and macOS keeps that for later runs.
 

check_targets.jai

@@ -49,7 +49,7 @@ APP_NAME        :: "mcap";
 APP_FOLDER_NAME :: "mcap";
 GIT_BRANCH    :: "check";
 GIT_HASH      :: "-";
-VERSION       :: "0.1.31";
+VERSION       :: "0.1.32";
 BUILD_DATE    :: "-";
 DEBUG         :: true;
 COMPILE_TESTS :: true;

first.jai

@@ -8,7 +8,7 @@ APP_NAME              :: "mcap";
 APP_FOLDER_NAME       :: "mcap";
 APP_BUNDLE_IDENTIFIER :: "com.example.mcap";
 
-VERSION :: "0.1.31";
+VERSION :: "0.1.32";
 // Where everything goes: bin/, or bin-headless/ for builds made with no_crash_dialog, so automated
 // builds never write over the ones you run by hand.
 build_directory := "bin";
@@ -407,6 +407,7 @@ copy_assets :: () {
     copy_if_changed("modules/Segment/models.config", tprint("%/segment-models/models.config", build_directory));
     if file_exists("modules/Segment/models") {
         for file_list("modules/Segment/models")  copy_if_changed(it, tprint("%/segment-models/%", build_directory, path_filename(it)));
+        #if OS == .MACOS  clone_packages("modules/Segment/models", tprint("%/segment-models", build_directory));
     } else {
         log("Segmentation: the model has not been made; see modules/Segment/README.md.");
     }

modules/Core_ML/module.jai

@@ -85,9 +85,18 @@ Coreml_Output :: struct {
   copy: *float;       // when data is a copy: the default allocator's
 }
 
-// Runs the model on input, a float32 tensor of shape. ok is false, with why, when it fails.
+// Runs the model on input, a float32 tensor of shape: its output (model.output_name). ok is false,
+// with why, when it fails.
 coreml_predict :: (model: *Coreml_Model, input: [] float, shape: [] s64) -> Coreml_Output, ok: bool, error: string {
-  result: Coreml_Output;
+  outs: [1] Coreml_Output;
+  ok, error := coreml_predict_outputs(model, input, shape, .[model.output_name], outs);
+  return outs[0], ok, error;
+}
+
+// As coreml_predict, for a model with several outputs: those named, into outs (one each, in order),
+// each released with coreml_release_output. On failure none is set.
+coreml_predict_outputs :: (model: *Coreml_Model, input: [] float, shape: [] s64, names: [] string, outs: [] Coreml_Output) -> ok: bool, error: string {
+  assert(names.count == outs.count);
   init_objective_c();
   pool := objc_new(NSAutoreleasePool);
   defer NSAutoreleasePool.drain(pool);
@@ -96,13 +105,13 @@ coreml_predict :: (model: *Coreml_Model, input: [] float, shape: [] s64) -> Core
   strides := NewArray(shape.count, s64,, temp);
   step: s64 = 1;
   for < shape { strides[it_index] = step; step *= it; }
-  if step != input.count  return result, false, tprint("the input has % values, not % for its shape %", input.count, step, shape);
+  if step != input.count  return false, tprint("the input has % values, not % for its shape %", input.count, step, shape);
 
   error: *NSError;
   func_array_init :: #type (*void, Selector, *void, *NSObject, s64, *NSObject, *void, **NSError) -> *NSObject #c_call;
   array := (cast(func_array_init) objc_msgSend)(objc_alloc(objc_getClass("MLMultiArray")), selector("initWithDataPointer:shape:dataType:strides:deallocator:error:"),
                                                 input.data, numbers(shape), MULTIARRAY_FLOAT32, numbers(strides), null, *error);
-  if !array  return result, false, tprint("Core ML could not take the input: %", error_text(error));
+  if !array  return false, tprint("Core ML could not take the input: %", error_text(error));
   autorelease(array);
 
   func_value :: #type (*void, Selector, *NSObject) -> *NSObject #c_call;
@@ -112,18 +121,51 @@ coreml_predict :: (model: *Coreml_Model, input: [] float, shape: [] s64) -> Core
   dictionary := (cast(func_dictionary) objc_msgSend)(objc_getClass("NSDictionary"), selector("dictionaryWithObject:forKey:"), value, key);
   func_provider :: #type (*void, Selector, *NSObject, **NSError) -> *NSObject #c_call;
   features := (cast(func_provider) objc_msgSend)(objc_alloc(objc_getClass("MLDictionaryFeatureProvider")), selector("initWithDictionary:error:"), dictionary, *error);
-  if !features  return result, false, tprint("Core ML could not take the input: %", error_text(error));
+  if !features  return false, tprint("Core ML could not take the input: %", error_text(error));
   autorelease(features);
 
   func_predict :: #type (*void, Selector, *NSObject, **NSError) -> *NSObject #c_call;
   lock(*model.mutex);
   prediction := (cast(func_predict) objc_msgSend)(model.model, selector("predictionFromFeatures:error:"), features, *error);
   unlock(*model.mutex);
-  if !prediction  return result, false, tprint("Core ML could not run the model: %", error_text(error));
+  if !prediction  return false, tprint("Core ML could not run the model: %", error_text(error));
 
-  out_value := (cast(func_value) objc_msgSend)(prediction, selector("featureValueForName:"), NSString.getTempString(model.output_name));
+  for names {
+    out, ok, read_error := read_output(prediction, it);
+    if !ok {
+      for 0..it_index - 1  coreml_release_output(*outs[it]);
+      return false, read_error;
+    }
+    outs[it_index] = out;
+  }
+  return true, "";
+}
+
+coreml_release_output :: (out: *Coreml_Output) {
+  if out.array  release(out.array);
+  if out.copy   free(out.copy,, context.default_allocator);
+  out.* = .{};
+}
+
+// coreml_predict, its output copied into output (which must hold exactly as many values).
+coreml_run :: (model: *Coreml_Model, input: [] float, shape: [] s64, output: [] float) -> ok: bool, error: string {
+  out, ok, error := coreml_predict(model, input, shape);
+  if !ok  return false, error;
+  defer coreml_release_output(*out);
+  if out.count != output.count  return false, tprint("the model gives % values, not %", out.count, output.count);
+  memcpy(output.data, out.data, out.count * size_of(float));
+  return true, "";
+}
+
+#scope_file
+
+// The output named name of a prediction (an MLFeatureProvider), as float32 values in C order.
+read_output :: (prediction: *NSObject, name: string) -> Coreml_Output, ok: bool, error: string {
+  result: Coreml_Output;
+  func_value :: #type (*void, Selector, *NSObject) -> *NSObject #c_call;
+  out_value := (cast(func_value) objc_msgSend)(prediction, selector("featureValueForName:"), NSString.getTempString(name));
   out_array := ifx out_value then cast(*NSObject) objc_msgSend_typed(out_value, selector("multiArrayValue")) else null;
-  if !out_array  return result, false, tprint("the model gives no %", model.output_name);
+  if !out_array  return result, false, tprint("the model gives no %", name);
 
   out_shape   := integers(cast(*NSObject) objc_msgSend_typed(out_array, selector("shape")));
   out_strides := integers(cast(*NSObject) objc_msgSend_typed(out_array, selector("strides")));
@@ -159,23 +201,6 @@ coreml_predict :: (model: *Coreml_Model, input: [] float, shape: [] s64) -> Core
   return result, true, "";
 }
 
-coreml_release_output :: (out: *Coreml_Output) {
-  if out.array  release(out.array);
-  if out.copy   free(out.copy,, context.default_allocator);
-  out.* = .{};
-}
-
-// coreml_predict, its output copied into output (which must hold exactly as many values).
-coreml_run :: (model: *Coreml_Model, input: [] float, shape: [] s64, output: [] float) -> ok: bool, error: string {
-  out, ok, error := coreml_predict(model, input, shape);
-  if !ok  return false, error;
-  defer coreml_release_output(*out);
-  if out.count != output.count  return false, tprint("the model gives % values, not %", out.count, output.count);
-  memcpy(output.data, out.data, out.count * size_of(float));
-  return true, "";
-}
-
-#scope_file
 
 MULTIARRAY_FLOAT32 :: 0x10000 | 32;   // MLMultiArrayDataTypeFloat32
 MULTIARRAY_FLOAT16 :: 0x10000 | 16;   // MLMultiArrayDataTypeFloat16

modules/Segment/README.md

@@ -18,7 +18,10 @@ conda run -n mcap-seg python tools/export_mask2former.py --check FOLDER_OF_IMAGE
 ```
 
 This writes `models/mask2former-swin-base-coco-panoptic.onnx` (462 MB) and its class list, then checks
-the .onnx against PyTorch on the images. Building the app copies `models.config` and the models into
+the .onnx against PyTorch on the images. On macOS, `tools/export_mask2former_coreml.py --check
+FOLDER_OF_IMAGES` then makes `models/mask2former-swin-base-coco-panoptic.mlpackage`, which runs on the
+GPU through Core ML (`modules/Core_ML`), one image a run, about four times as fast as the CPU with the
+same segments. Building the app copies `models.config` and the models into
 `bin/segment-models`. The models folder is not in git.
 
 ## Using it

modules/Segment/models.config

@@ -6,6 +6,7 @@
 [model mask2former-swin-base-coco-panoptic]
 title:             "Mask2Former Swin-B, COCO panoptic (133 classes)"
 model:             "mask2former-swin-base-coco-panoptic.onnx"
+model_coreml:      "mask2former-swin-base-coco-panoptic.mlpackage"   # macOS: on the GPU (tools/export_mask2former_coreml.py makes it)
 classes:           "mask2former-swin-base-coco-panoptic.classes.txt"
 image_size:        384             # squashed to 384 x 384, not keeping its shape, as the checkpoint's processor does
 mean:              0.485 0.456 0.406

modules/Segment/module.jai

@@ -11,7 +11,7 @@
 // transformers' post_process_panoptic_segmentation does (tools/export_mask2former.py checks that
 // they match).
 
-Segment_Provider :: enum { CPU; CUDA; }
+Segment_Provider :: enum { CPU; CUDA; COREML; }
 
 // Tightly packed or not: stride is the bytes from one row to the next.
 Segment_Image :: struct {
@@ -23,6 +23,7 @@ Segment_Model_Info :: struct {
   name: string;
   title: string;
   model: string;                       // the .onnx, in the models folder
+  model_coreml: string;                // macOS: the model as a Core ML package (tools/export_mask2former_coreml.py), run on the GPU; "" for none
   classes: string;                     // the class list, in the models folder
   image_size: s32 = 384;               // the input's width and height: images are resized to it, not keeping their shape
   mean := float.[0.485, 0.456, 0.406]; // ImageNet's, per channel, after scaling to 0..1
@@ -50,6 +51,7 @@ Segment_Engine :: struct {
   env: *OrtEnv;
   memory_info: *OrtMemoryInfo;
   session: Ort_Session;
+  #if OS == .MACOS  coreml: *Coreml_Model;   // the model, when Core ML runs it (one image a run)
   cancel: *bool;                       // when set, and true, segment_images stops between batches: the rest fail as "cancelled"
   stats: Segment_Stats;                // since it was made
 }
@@ -114,6 +116,7 @@ segment_list_models :: (models_folder: string) -> [] Segment_Model_Info, error:
       case "title";             model.title   = copy_string(value);
       case "model";             model.model   = copy_string(value);
       case "classes";           model.classes = copy_string(value);
+      case "model_coreml";      model.model_coreml = copy_string(value);
       case "image_size";        if number_ok  model.image_size = cast(s32) number;
       case "mean";              segment_parse_three(value, *model.mean, path, line_number);
       case "std";               segment_parse_three(value, *model.std, path, line_number);
@@ -155,11 +158,29 @@ segment_create :: (models_folder: string, model_name: string, prefer_gpu := true
   HINT :: "make it with tools/export_mask2former.py (see modules/Segment/README.md) and rebuild";
   path := tprint("%/%", models_folder, model.model);
   use_gpu := false;
-  if prefer_gpu {
+  #if OS == .MACOS  if prefer_gpu {
+    // Core ML on the GPU; on any failure, onnxruntime on the CPU below, with the reason.
+    coreml_path := tprint("%/%", models_folder, model.model_coreml);
+    if !model.model_coreml || !file_exists(coreml_path) {
+      engine.gpu_explanation = sprint("Segmentation runs on the CPU: % is not there; tools/export_mask2former_coreml.py makes it.", ifx model.model_coreml then model.model_coreml else "its Core ML model");
+    } else {
+      engine.coreml, ok, error = coreml_load(coreml_path, "pixel_values", "class_queries_logits");
+      if ok {
+        engine.provider = .COREML;
+        ok, error = segment_warm_up(engine);
+        if ok  return engine, true, "";
+        coreml_free(engine.coreml);
+        engine.coreml = null;
+        engine.provider = .CPU;
+      }
+      engine.gpu_explanation = sprint("Segmentation runs on the CPU: Core ML could not run % on the GPU:\n%", model.model_coreml, error);
+    }
+  }
+  if prefer_gpu && OS != .MACOS {
     available, explanation := ort_cuda_check("Segmentation");
     use_gpu = available;
     if !available  engine.gpu_explanation = copy_string(explanation);
-  } else {
+  } else if !prefer_gpu {
     engine.gpu_explanation = copy_string("Segmentation runs on the CPU, as it was asked to.");
   }
   if use_gpu {
@@ -190,6 +211,7 @@ segment_destroy :: (engine: *Segment_Engine) {
   if engine.provider == .CUDA  ort_gpu_lock();
   ort_release_session(*engine.session);
   if engine.provider == .CUDA  ort_gpu_unlock();
+  #if OS == .MACOS  coreml_free(engine.coreml);
   if engine.memory_info  ort.api.ReleaseMemoryInfo(engine.memory_info);
   if engine.env          ort.api.ReleaseEnv(engine.env);
   for engine.classes { free(it.name); free(it.model_name); }
@@ -198,6 +220,24 @@ segment_destroy :: (engine: *Segment_Engine) {
   free(engine);
 }
 
+// Images a run: Core ML's model takes one.
+segment_batch :: (engine: *Segment_Engine) -> s32 {
+  if engine.provider == {
+    case .CUDA;   return engine.model.batch_gpu;
+    case .COREML; return 1;
+  }
+  return engine.model.batch_cpu;
+}
+
+// What runs the model, for people: "CUDA", "Core ML" or "CPU".
+segment_provider_name :: (engine: *Segment_Engine) -> string {
+  if engine.provider == {
+    case .CUDA;   return "CUDA";
+    case .COREML; return "Core ML";
+  }
+  return "CPU";
+}
+
 // The map size for an image of width x height: its shape, the longest side at most max_side.
 segment_map_size :: (width: s32, height: s32, max_side: s32) -> s32, s32 {
   scale := min(cast(float64) 1, cast(float64) max_side / max(width, height));
@@ -210,7 +250,7 @@ segment_images :: (engine: *Segment_Engine, images: [] Segment_Image, allocator
   results := NewArray(images.count, Segment_Result,, allocator);
   size := engine.model.image_size;
   plane := cast(s64) size * size;
-  batch := cast(s64) max(ifx engine.provider == .CUDA then engine.model.batch_gpu else engine.model.batch_cpu, 1);
+  batch := cast(s64) max(segment_batch(engine), 1);
   input := NewArray(plane * 3 * batch, float,, context.default_allocator);
   defer array_free(input,, context.default_allocator);
 
@@ -224,10 +264,10 @@ segment_images :: (engine: *Segment_Engine, images: [] Segment_Image, allocator
     start := current_time_monotonic();
     for 0..count - 1  segment_prepare_image(engine.model, images[first + it], input.data + plane * 3 * it);
     shape := s64.[count, 3, size, size];
-    outs: [2] Ort_Output;
+    outs: [2] Segment_Output;
     ran := current_time_monotonic();
     engine.stats.prepare_us += to_microseconds(ran - start);
-    ok, error := ort_run_float_outputs(*engine.session, engine.memory_info, array_view(input, 0, plane * 3 * count), shape, .["class_queries_logits", "masks_queries_logits"], outs);
+    ok, error := segment_run(engine, array_view(input, 0, plane * 3 * count), shape, outs);
     posted := current_time_monotonic();
     engine.stats.model_us += to_microseconds(posted - ran);
     engine.stats.images += count;
@@ -249,8 +289,7 @@ segment_images :: (engine: *Segment_Engine, images: [] Segment_Image, allocator
                                                 cast(s32) queries, cast(s32) labels, cast(s32) mask_w, cast(s32) mask_h, image.width, image.height, allocator);
         }
       }
-      ort_release_output(*outs[0]);
+      segment_release_outputs(engine, outs);
-      ort_release_output(*outs[1]);
     }
     first += count;
   }
@@ -524,14 +563,48 @@ segment_warm_up :: (engine: *Segment_Engine) -> ok: bool, error: string {
   size := engine.model.image_size;
   input := NewArray(3 * size * size, float, initialized = true,, context.default_allocator);
   defer array_free(input,, context.default_allocator);
-  outs: [2] Ort_Output;
+  outs: [2] Segment_Output;
-  ok, error := ort_run_float_outputs(*engine.session, engine.memory_info, input, s64.[1, 3, size, size], .["class_queries_logits", "masks_queries_logits"], outs);
+  ok, error := segment_run(engine, input, s64.[1, 3, size, size], outs);
   if !ok  return false, tprint("The first run on the GPU failed: %", error);
-  ort_release_output(*outs[0]);
+  segment_release_outputs(engine, outs);
-  ort_release_output(*outs[1]);
   return true, "";
 }
 
+// The two outputs, class and mask logits, as onnxruntime or Core ML gives them.
+Segment_Output :: struct {
+  data: *float;
+  dims: [4] s64;
+  dim_count: s64;
+  ort: Ort_Output;
+  #if OS == .MACOS  coreml: Coreml_Output;
+}
+
+segment_run :: (engine: *Segment_Engine, input: [] float, shape: [] s64, outs: [] Segment_Output) -> ok: bool, error: string {
+  NAMES :: string.["class_queries_logits", "masks_queries_logits"];
+  #if OS == .MACOS  if engine.coreml {
+    coreml_outs: [2] Coreml_Output;
+    ok, error := coreml_predict_outputs(engine.coreml, input, shape, NAMES, coreml_outs);
+    if !ok  return false, error;
+    for * outs {
+      it.* = .{data = coreml_outs[it_index].data, dim_count = min(coreml_outs[it_index].dim_count, 4), coreml = coreml_outs[it_index]};
+      for d: 0..it.dim_count - 1  it.dims[d] = coreml_outs[it_index].dims[d];
+    }
+    return true, "";
+  }
+  ort_outs: [2] Ort_Output;
+  ok, error := ort_run_float_outputs(*engine.session, engine.memory_info, input, shape, NAMES, ort_outs);
+  if !ok  return false, error;
+  for * outs  it.* = .{data = ort_outs[it_index].data, dims = ort_outs[it_index].dims, dim_count = ort_outs[it_index].dim_count, ort = ort_outs[it_index]};
+  return true, "";
+}
+
+segment_release_outputs :: (engine: *Segment_Engine, outs: [] Segment_Output) {
+  for * outs {
+    ort_release_output(*it.ort);
+    #if OS == .MACOS  coreml_release_output(*it.coreml);
+  }
+}
+
 // The class list: `class: INDEX thing|stuff "model's name" "name"`, a line each, in index order.
 segment_load_classes :: (engine: *Segment_Engine, path: string) -> ok: bool, error: string {
   text, ok := read_entire_file(path, log_errors = false,, temp);
@@ -597,9 +670,11 @@ segment_strip_comment :: (line: string) -> string {
 #import "Basic";
 #import "String";
 #import "File";
+#import "File_Utilities";
 #import "Math";
 #import "System";
 #import "Hash";
 #import "Sort";
 #import "Capture_File";   // segment_mcap_chunks
 #import "onnxruntime";
+#if OS == .MACOS  #import "Core_ML";

src/app/segment_ui.jai

@@ -198,7 +198,7 @@ segment_ui_scan_task :: () {
       return;
     }
     segment_ui.engine = engine;
-    log("[Segment] % on % in % ms.%", SEGMENT_MODEL, engine.provider, to_milliseconds(current_time_monotonic() - start), ifx engine.gpu_explanation then tprint(" %", engine.gpu_explanation) else "");
+    log("[Segment] % on % in % ms.%", SEGMENT_MODEL, segment_provider_name(engine), to_milliseconds(current_time_monotonic() - start), ifx engine.gpu_explanation then tprint(" %", engine.gpu_explanation) else "");
   }
   engine := segment_ui.engine;
   // Cancelling (or closing the app) stops it between batches of the engine's, rather than waiting the
@@ -335,7 +335,7 @@ segment_window_contents :: () { using ImGui;
   // The model.
   model_file := tprint("%/%.onnx", segment_models_folder(), SEGMENT_MODEL);
   if !busy && segment_ui.engine {
-    Text("% on %", SEGMENT_MODEL, ifx segment_ui.engine.provider == .CUDA then "the GPU" else "the CPU");
+    Text("% on %", SEGMENT_MODEL, ifx segment_ui.engine.provider != .CPU then tprint("the GPU (%)", segment_provider_name(segment_ui.engine)) else "the CPU");
     if segment_ui.engine.gpu_explanation  TextColored(app_theme.palette.dim, "%", segment_ui.engine.gpu_explanation);
   } else if !busy && segment_ui.engine_error {
     TextColored(app_theme.palette.warning, "% could not start: %", SEGMENT_MODEL, segment_ui.engine_error);

tools/export_mask2former_coreml.py

@@ -0,0 +1,173 @@
+# Converts Mask2Former (facebook/mask2former-swin-base-coco-panoptic) to a Core ML model for
+# modules/Segment on macOS, then checks it against the ONNX model on a folder of images, segment by
+# segment. Run once, on a Mac, after tools/export_mask2former.py; mcap only uses the .mlpackage it writes.
+#
+#   python3 -m venv ~/mcap-coreml
+#   ~/mcap-coreml/bin/pip install torch==2.7.0 transformers coremltools==9.0 onnxruntime pillow scipy
+#   ~/mcap-coreml/bin/pip install --no-deps timm
+#   ~/mcap-coreml/bin/python tools/export_mask2former_coreml.py --check FOLDER_OF_IMAGES
+#
+# The model as tools/export_mask2former.py exports it (pixel_values [1, 3, 384, 384] to
+# class_queries_logits [1, 100, 134] and masks_queries_logits [1, 100, 96, 96]), as an ML Program in
+# float32, one image a run (a batch of 4 did not convert). Three things coremltools 9.0 cannot take
+# are worked around for the conversion only, none changing what the model computes:
+#   - it casts a one-element tensor with int(array), which NumPy 2 refuses;
+#   - its meshgrid refuses the inputs Mask2Former's traced linspaces come out as;
+#   - Core ML takes tensors of at most 5 dimensions, and the multi-scale deformable attention makes 6
+#     (batch, queries, heads, levels, points, xy): it is run a level at a time instead, which gives
+#     the same output in PyTorch (checked below) with at most 5.
+# Measured on an M5 Pro: 42 ms an image on the GPU against 166 ms for onnxruntime on the CPU, with
+# the same segments (IoU 1.0) on every test image.
+
+import argparse, glob, os, sys, time
+import numpy as np
+import torch
+from torch import nn
+from PIL import Image
+import coremltools as ct
+from coremltools.converters.mil import Builder as mb
+from coremltools.converters.mil.frontend.torch import ops as ct_ops
+from coremltools.converters.mil.frontend.torch.torch_op_registry import register_torch_op
+from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
+from transformers.utils import logging
+import transformers.models.mask2former.modeling_mask2former as modeling
+
+logging.set_verbosity_error()
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from export_mask2former import Logits
+
+CHECKPOINT = "facebook/mask2former-swin-base-coco-panoptic"
+MODELS = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "modules", "Segment", "models")
+ONNX = os.path.join(MODELS, "mask2former-swin-base-coco-panoptic.onnx")
+OUTPUT = os.path.join(MODELS, "mask2former-swin-base-coco-panoptic.mlpackage")
+SIZE = 384
+
+
+_original_cast = ct_ops._cast
+def _cast_one_element(context, node, dtype, dtype_name):
+    x = ct_ops._get_inputs(context, node, expected=1)[0]
+    if x.can_be_folded_to_const() and np.ndim(x.val) > 0 and np.size(x.val) == 1:
+        context.add(ct_ops.mb.const(val=dtype(np.asarray(x.val).item()), name=node.name), node.name)
+        return
+    _original_cast(context, node, dtype, dtype_name)
+ct_ops._cast = _cast_one_element
+
+
+@register_torch_op(override=True)
+def meshgrid(context, node):
+    """torch.meshgrid, from its inputs flattened to 1-d first."""
+    inputs = ct_ops._get_inputs(context, node, expected=[1, 2])
+    indexing = inputs[1].val if len(inputs) > 1 else "ij"
+    indexing = ct_ops._get_kwinputs(context, node, "indexing", default=[indexing])[0]
+    if hasattr(indexing, "val"):
+        indexing = indexing.val
+    tensors = [mb.reshape(x=t, shape=[-1], name=f"{node.name}_flat_{i}") for i, t in enumerate(inputs[0])]
+    sizes = [t.shape[0] for t in tensors]
+    grids = []
+    for i, t in enumerate(tensors):
+        view = [1] * len(tensors); view[i] = -1
+        reps = list(sizes); reps[i] = 1
+        grid = mb.tile(x=mb.reshape(x=t, shape=view, name=f"{node.name}_view_{i}"), reps=reps, name=f"{node.name}_grid_{i}")
+        if indexing == "xy":
+            grid = mb.transpose(x=grid, perm=[1, 0] + list(range(2, len(tensors))), name=f"{node.name}_xy_{i}")
+        grids.append(grid)
+    context.add(tuple(grids), node.name)
+
+
+def deformable_attention_rank5(self, hidden_states, attention_mask=None, encoder_hidden_states=None, encoder_attention_mask=None,
+                               position_embeddings=None, reference_points=None, spatial_shapes_list=None, level_start_index=None,
+                               output_attentions=False):
+    """Mask2FormerPixelDecoderEncoderMultiscaleDeformableAttention.forward, a level at a time."""
+    if position_embeddings is not None:
+        hidden_states = self.with_pos_embed(hidden_states, position_embeddings)
+    if reference_points.shape[-1] != 2:
+        raise ValueError("only 2-d reference points are rewritten")
+    B, Q, _ = hidden_states.shape
+    _, S, _ = encoder_hidden_states.shape
+    H, L, P, D = self.n_heads, self.n_levels, self.n_points, self.d_model // self.n_heads
+    value = self.value_proj(encoder_hidden_states)
+    if attention_mask is not None:
+        value = value.masked_fill(attention_mask[..., None], float(0))
+    value = value.view(B, S, H, D)
+    offsets = self.sampling_offsets(hidden_states).view(B, Q, H, L, P * 2)
+    weights = nn.functional.softmax(self.attention_weights(hidden_states).view(B, Q, H, L * P), -1)
+    weights = weights.transpose(1, 2).reshape(B * H, 1, Q, L * P)
+    values = value.split([h * w for h, w in spatial_shapes_list], dim=1)
+    sampled = []
+    for level, (h, w) in enumerate(spatial_shapes_list):
+        value_l = values[level].flatten(2).transpose(1, 2).reshape(B * H, D, h, w)
+        offset_l = offsets[:, :, :, level].reshape(B, Q, H, P, 2)
+        location = reference_points[:, :, None, level, None, :] + offset_l / torch.tensor([w, h], dtype=offset_l.dtype)
+        grid = (2 * location - 1).transpose(1, 2).flatten(0, 1)
+        sampled.append(nn.functional.grid_sample(value_l, grid, mode="bilinear", padding_mode="zeros", align_corners=False))
+    output = (torch.stack(sampled, dim=-2).flatten(-2) * weights).sum(-1).view(B, H * D, Q).transpose(1, 2).contiguous()
+    return self.output_proj(output), None
+
+
+def segments_of(processor, class_logits, mask_logits):
+    holder = type("Outputs", (), {"class_queries_logits": torch.from_numpy(class_logits), "masks_queries_logits": torch.from_numpy(mask_logits)})()
+    return processor.post_process_panoptic_segmentation(holder, target_sizes=[(SIZE, SIZE)])[0]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--check", metavar="FOLDER", help="images to check the Core ML model against the ONNX one on")
+    args = parser.parse_args()
+
+    processor = AutoImageProcessor.from_pretrained(CHECKPOINT)
+    model = Mask2FormerForUniversalSegmentation.from_pretrained(CHECKPOINT).eval()
+    example = torch.zeros(1, 3, SIZE, SIZE)
+    with torch.no_grad():
+        original = Logits(model)(torch.randn(1, 3, SIZE, SIZE, generator=torch.Generator().manual_seed(0)))
+        modeling.Mask2FormerPixelDecoderEncoderMultiscaleDeformableAttention.forward = deformable_attention_rank5
+        rewritten = Logits(model)(torch.randn(1, 3, SIZE, SIZE, generator=torch.Generator().manual_seed(0)))
+        difference = max((a - b).abs().max().item() for a, b in zip(original, rewritten))
+        print(f"The attention a level at a time, against the original in PyTorch: max difference {difference:.2e}")
+        if difference > 1e-4:
+            sys.exit("The rewritten attention does not give the same output.")
+        traced = torch.jit.trace(Logits(model).eval(), example)
+    start = time.perf_counter()
+    ml = ct.convert(traced, inputs=[ct.TensorType(name="pixel_values", shape=(1, 3, SIZE, SIZE))],
+                    outputs=[ct.TensorType(name="class_queries_logits"), ct.TensorType(name="masks_queries_logits")],
+                    convert_to="mlprogram", compute_precision=ct.precision.FLOAT32, minimum_deployment_target=ct.target.macOS14)
+    ml.short_description = f"{CHECKPOINT}, class and mask logits, float32, for mcap"
+    ml.save(OUTPUT)
+    print(f"Wrote {OUTPUT} in {time.perf_counter() - start:.0f} s.")
+
+    if not args.check:
+        return
+    import onnxruntime as ort
+    paths = sorted(p for p in glob.glob(os.path.join(args.check, "**", "*"), recursive=True)
+                   if p.lower().endswith((".png", ".jpg", ".jpeg", ".webp")))
+    if not paths:
+        sys.exit(f"No images in {args.check}.")
+    session = ort.InferenceSession(ONNX, providers=["CPUExecutionProvider"])
+    coreml = ct.models.MLModel(OUTPUT, compute_units=ct.ComputeUnit.CPU_AND_GPU)
+    total, worst_iou, unmatched, times = 0, 1.0, 0, []
+    for path in paths:
+        pixels = processor(images=Image.open(path).convert("RGB"), return_tensors="np")["pixel_values"].astype(np.float32)
+        ref = session.run(["class_queries_logits", "masks_queries_logits"], {"pixel_values": pixels})
+        start = time.perf_counter()
+        got = coreml.predict({"pixel_values": pixels})
+        times.append(time.perf_counter() - start)
+        a = segments_of(processor, ref[0], ref[1])
+        b = segments_of(processor, got["class_queries_logits"].astype(np.float32), got["masks_queries_logits"].astype(np.float32))
+        for sa in a["segments_info"]:
+            total += 1
+            best = 0.0
+            for sb in b["segments_info"]:
+                if sb["label_id"] == sa["label_id"]:
+                    ma, mbk = (a["segmentation"] == sa["id"]).numpy(), (b["segmentation"] == sb["id"]).numpy()
+                    best = max(best, (ma & mbk).sum() / max((ma | mbk).sum(), 1))
+            if best == 0:
+                unmatched += 1
+            else:
+                worst_iou = min(worst_iou, best)
+    print(f"{len(paths)} images, {total} segments: worst IoU against the ONNX model {worst_iou:.4f}, unmatched {unmatched}; "
+          f"{np.median(times[1:] or times) * 1000:.0f} ms an image on the GPU")
+    if unmatched or worst_iou < 0.98:
+        sys.exit("The Core ML model segments differently from the ONNX one.")
+
+
+if __name__ == "__main__":
+    main()

tools/segment_probe.jai

@@ -33,7 +33,7 @@ main :: () {
   engine, ok, error := segment_create(models_folder, model_name, prefer_gpu);
   if !ok { print("Cannot start %: %\n", model_name, error); exit(1); }
   defer segment_destroy(engine);
-  print("% on %, % classes (made in % ms)%\n", model_name, engine.provider, engine.classes.count, ms(start),
+  print("% on %, % classes (made in % ms)%\n", model_name, segment_provider_name(engine), engine.classes.count, ms(start),
         ifx engine.gpu_explanation then tprint(": %", engine.gpu_explanation) else "");
 
   failures := 0;
@@ -129,7 +129,7 @@ main :: () {
 
   print("\n% images: decoding % ms each, segmenting % ms each (model and post-processing, batch %)\n", done,
         formatFloat(decode_ms / max(done, 1), trailing_width = 0), formatFloat(model_ms / max(done, 1), trailing_width = 0),
-        ifx engine.provider == .CUDA then engine.model.batch_gpu else engine.model.batch_cpu);
+        segment_batch(engine));
   per := max(engine.stats.images, 1) * 1000.0;
   print("  of which: preparing % ms, the model % ms, making segments % ms, per image\n", formatFloat(engine.stats.prepare_us / per, trailing_width = 1),
         formatFloat(engine.stats.model_us / per, trailing_width = 1), formatFloat(engine.stats.post_us / per, trailing_width = 1));