Tests are green

Author: kipcole9

.gitignore

@@ -41,3 +41,7 @@ test/perf/*.jpeg
 # ML model
 *.pt
 *.onnx
+
+# Tools
+mise.toml
+.tool-versions

config/test.exs

@@ -1,8 +1,6 @@
 import Config
 
-config :nx,
-  default_backend: EXLA.Backend
-
 config :logger,
   level: :warning
 
+config :nx, :default_defn_options, [compiler: EXLA]

lib/classification.ex

@@ -157,7 +157,6 @@ if ImageVision.bumblebee_configured?() do
     @spec classifier(configuration :: Keyword.t()) ::
             {Nx.Serving, Keyword.t()} | {:error, Image.error()}
     def classifier(classifier \\ Application.get_env(:image_vision, :classifier, [])) do
-      Application.ensure_all_started(:exla)
       classifier = Keyword.merge(@default_classifier, classifier)
 
       model = Keyword.fetch!(classifier, :model)
@@ -224,7 +223,6 @@ if ImageVision.bumblebee_configured?() do
     @spec embedder(configuration :: Keyword.t()) ::
             {Nx.Serving, Keyword.t()} | {:error, Image.error()}
     def embedder(embedder \\ Application.get_env(:image_vision, :embedder, [])) do
-      Application.ensure_all_started(:exla)
       embedder = Keyword.merge(@default_embedder, embedder)
 
       model = Keyword.fetch!(embedder, :model)
@@ -256,7 +254,7 @@ if ImageVision.bumblebee_configured?() do
            {:ok, featurizer} = Bumblebee.load_featurizer(featurizer, featurizer_options) do
         Bumblebee.Vision.image_classification(model_info, featurizer,
           compile: [batch_size: batch_size],
-          defn_options: [compiler: EXLA]
+          defn_options: defn_options()
         )
       end
     end
@@ -267,11 +265,23 @@ if ImageVision.bumblebee_configured?() do
            {:ok, featurizer} = Bumblebee.load_featurizer(featurizer, featurizer_options) do
         Bumblebee.Vision.image_embedding(model_info, featurizer,
           compile: [batch_size: batch_size],
-          defn_options: [compiler: EXLA]
+          defn_options: defn_options()
         )
       end
     end
 
+    # Use EXLA as the Nx compiler when it is properly loaded and implements
+    # the current Nx.Defn.Compiler protocol. Falls back to the default
+    # evaluator when EXLA is absent or version-mismatched (e.g. EXLA 0.10
+    # paired with Nx 0.11 does not export __compile__/4).
+    defp defn_options do
+      if Code.ensure_loaded?(EXLA) and function_exported?(EXLA, :__compile__, 4) do
+        [compiler: EXLA]
+      else
+        []
+      end
+    end
+
     @doc """
     Classifies an image and returns the full prediction map.
 

lib/detection.ex

@@ -266,15 +266,11 @@ if ImageVision.ortex_configured?() do
 
       padded = Image.embed!(resized, @input_size, @input_size, x: 0, y: 0)
 
-      mean = Nx.tensor([0.485, 0.456, 0.406])
-      std = Nx.tensor([0.229, 0.224, 0.225])
-
       tensor =
         padded
-        |> Image.to_nx!()
+        |> Image.to_nx!(backend: Nx.BinaryBackend)
         |> Nx.as_type(:f32)
         |> Nx.divide(255.0)
-        |> NxImage.normalize(mean, std)
         |> Nx.transpose(axes: [2, 0, 1])
         |> Nx.new_axis(0)
 
@@ -297,8 +293,12 @@ if ImageVision.ortex_configured?() do
       original_height = Keyword.fetch!(opts, :original_height)
       min_score = Keyword.fetch!(opts, :min_score)
 
-      scores = Nx.sigmoid(logits[0])
-      boxes = pred_boxes[0]
+      scores =
+        logits[0]
+        |> Nx.backend_transfer(Nx.BinaryBackend)
+        |> Nx.sigmoid()
+
+      boxes = pred_boxes[0] |> Nx.backend_transfer(Nx.BinaryBackend)
 
       best_class = Nx.argmax(scores, axis: 1)
       best_score = Nx.reduce_max(scores, axes: [1])

lib/segmentation.ex

@@ -180,7 +180,7 @@ if ImageVision.ortex_configured?() do
 
       orig_w = Image.width(image)
       orig_h = Image.height(image)
-      input_scale = @sam_input_size / max(orig_w, orig_h)
+      input_scale = max(orig_w, orig_h) / @sam_input_size
 
       {point_coords, point_labels} =
         encode_sam_prompt(prompt, orig_w, orig_h, input_scale)
@@ -265,13 +265,13 @@ if ImageVision.ortex_configured?() do
       orig_w = Image.width(image)
       orig_h = Image.height(image)
 
-      {batch, input_h, input_w} = detr_preprocess(image)
+      {batch, _input_h, _input_w} = detr_preprocess(image)
 
       {logits, _pred_boxes, pred_masks} = Ortex.run(model, batch)
+      logits = Nx.backend_transfer(logits, Nx.BinaryBackend)
+      pred_masks = Nx.backend_transfer(pred_masks, Nx.BinaryBackend)
 
       detr_postprocess(logits, pred_masks, id2label,
-        input_h: input_h,
-        input_w: input_w,
         orig_w: orig_w,
         orig_h: orig_h,
         min_score: min_score
@@ -413,7 +413,7 @@ if ImageVision.ortex_configured?() do
 
       tensor =
         padded
-        |> Image.to_nx!()
+        |> Image.to_nx!(backend: Nx.BinaryBackend)
         |> Nx.as_type(:f32)
         |> Nx.divide(255.0)
         |> NxImage.normalize(Nx.tensor(@imagenet_mean), Nx.tensor(@imagenet_std))
@@ -424,16 +424,17 @@ if ImageVision.ortex_configured?() do
     end
 
     defp sam_encode(encoder, tensor) do
-      {image_embed, high_res_feats_0, high_res_feats_1} = Ortex.run(encoder, tensor)
+      # Model outputs in order: high_res_feats_0, high_res_feats_1, image_embed
+      {high_res_feats_0, high_res_feats_1, image_embed} = Ortex.run(encoder, tensor)
       {image_embed, high_res_feats_0, high_res_feats_1}
     end
 
     # Converts a user prompt into SAM point_coords + point_labels tensors.
     # Coords are in the 1024×1024 input space.
-    defp encode_sam_prompt(:auto, orig_w, orig_h, _input_scale) do
+    defp encode_sam_prompt(:auto, orig_w, orig_h, input_scale) do
       cx = orig_w / 2
       cy = orig_h / 2
-      encode_sam_prompt({:point, cx, cy}, orig_w, orig_h, orig_w / @sam_input_size)
+      encode_sam_prompt({:point, cx, cy}, orig_w, orig_h, input_scale)
     end
 
     defp encode_sam_prompt({:point, x, y}, _orig_w, _orig_h, scale) do
@@ -469,7 +470,7 @@ if ImageVision.ortex_configured?() do
       mask_input = Nx.broadcast(Nx.tensor(0, type: :f32), {1, 1, 256, 256})
       has_mask = Nx.tensor([0], type: :f32)
 
-      {masks, iou_predictions, _low_res} =
+      {masks, iou_predictions} =
         Ortex.run(decoder, {
           image_embed,
           high_res_feats_0,
@@ -480,29 +481,33 @@ if ImageVision.ortex_configured?() do
           has_mask
         })
 
-      {masks, iou_predictions}
+      {
+        Nx.backend_transfer(masks, Nx.BinaryBackend),
+        Nx.backend_transfer(iou_predictions, Nx.BinaryBackend)
+      }
     end
 
-    # Converts a SAM output mask tensor [256, 256] (logits) to a
+    # Converts a SAM output mask tensor {256, 256} (logits) to a
     # single-band Vimage at original image dimensions.
     defp sam_mask_to_image(mask_tensor, orig_w, orig_h, resized_w, resized_h) do
-      # Up to the padded 1024×1024 input, then crop to the valid region,
-      # then resize to original dimensions.
-      upscaled =
-        mask_tensor
-        |> Nx.reshape({1024, 1024, 1})
-        |> Nx.slice([0, 0, 0], [resized_h, resized_w, 1])
+      # The 256×256 mask covers the full 1024×1024 padded input space.
+      # Crop to the valid region, which is resized_w×resized_h in input
+      # space, scaled to mask space by the factor 256/1024 = 1/4.
+      valid_h = round(resized_h * 256 / @sam_input_size)
+      valid_w = round(resized_w * 256 / @sam_input_size)
 
-      # Threshold at 0 (logits > 0 → object)
       binary =
-        upscaled
+        mask_tensor
+        |> Nx.slice([0, 0], [valid_h, valid_w])
         |> Nx.greater(0)
         |> Nx.multiply(255)
         |> Nx.as_type(:u8)
+        |> Nx.transpose()
+        |> Nx.new_axis(2)
 
       binary
       |> Image.from_nx!()
-      |> Image.resize!(orig_w / resized_w, vertical_scale: orig_h / resized_h)
+      |> Image.resize!(orig_w / valid_w, vertical_scale: orig_h / valid_h)
     end
 
     # --- Private: DETR-panoptic pre/post --------------------------------
@@ -525,14 +530,15 @@ if ImageVision.ortex_configured?() do
 
       tensor =
         resized
-        |> Image.to_nx!()
+        |> Image.to_nx!(backend: Nx.BinaryBackend)
         |> Nx.as_type(:f32)
         |> Nx.divide(255.0)
         |> NxImage.normalize(Nx.tensor(@imagenet_mean), Nx.tensor(@imagenet_std))
         |> Nx.transpose(axes: [2, 0, 1])
         |> Nx.new_axis(0)
 
-      {tensor, input_h, input_w}
+      pixel_mask = Nx.broadcast(Nx.tensor(1, type: :s64), {1, 64, 64})
+      {{tensor, pixel_mask}, input_h, input_w}
     end
 
     # Loads id2label from config.json; cached in :persistent_term.
@@ -556,8 +562,6 @@ if ImageVision.ortex_configured?() do
     # Converts raw DETR-panoptic outputs into a list of segments.
     defp detr_postprocess(logits, pred_masks, id2label, opts) do
       min_score = Keyword.fetch!(opts, :min_score)
-      input_h = Keyword.fetch!(opts, :input_h)
-      input_w = Keyword.fetch!(opts, :input_w)
       orig_w = Keyword.fetch!(opts, :orig_w)
       orig_h = Keyword.fetch!(opts, :orig_h)
 
@@ -578,9 +582,8 @@ if ImageVision.ortex_configured?() do
       class_list = Nx.to_flat_list(best_class)
       score_list = Nx.to_flat_list(best_score)
 
-      # pred_masks: [1, 100, H/4, W/4]
-      mask_h = div(input_h, 4)
-      mask_w = div(input_w, 4)
+      # pred_masks: [1, queries, H/4, W/4]
+      {1, _num_queries, mask_h, mask_w} = Nx.shape(pred_masks)
 
       Enum.zip(class_list, score_list)
       |> Enum.with_index()
@@ -591,11 +594,12 @@ if ImageVision.ortex_configured?() do
 
           mask =
             mask_tensor
-            |> Nx.reshape({mask_h, mask_w, 1})
             |> Nx.sigmoid()
             |> Nx.greater(0.5)
             |> Nx.multiply(255)
             |> Nx.as_type(:u8)
+            |> Nx.transpose()
+            |> Nx.new_axis(2)
             |> Image.from_nx!()
             |> Image.resize!(orig_w / mask_w, vertical_scale: orig_h / mask_h)
 

mix.exs

@@ -75,6 +75,7 @@ defmodule ImageVision.MixProject do
       {:nx, "~> 0.11", optional: true, override: true},
       {:nx_image, "~> 0.1", optional: true},
       {:bumblebee, "~> 0.6", optional: true},
+      {:exla, "~> 0.11", optional: true},
 
       # --- Tooling ---
       {:ex_doc, "~> 0.18", only: [:release, :dev, :docs]},

mix.lock

@@ -12,9 +12,9 @@
   "elixir_make": {:hex, :elixir_make, "0.9.0", "6484b3cd8c0cee58f09f05ecaf1a140a8c97670671a6a0e7ab4dc326c3109726", [:mix], [], "hexpm", "db23d4fd8b757462ad02f8aa73431a426fe6671c80b200d9710caf3d1dd0ffdb"},
   "erlex": {:hex, :erlex, "0.2.8", "cd8116f20f3c0afe376d1e8d1f0ae2452337729f68be016ea544a72f767d9c12", [:mix], [], "hexpm", "9d66ff9fedf69e49dc3fd12831e12a8a37b76f8651dd21cd45fcf5561a8a7590"},
   "ex_doc": {:hex, :ex_doc, "0.40.1", "67542e4b6dde74811cfd580e2c0149b78010fd13001fda7cfeb2b2c2ffb1344d", [:mix], [{:earmark_parser, "~> 1.4.44", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "bcef0e2d360d93ac19f01a85d58f91752d930c0a30e2681145feea6bd3516e00"},
-  "exla": {:hex, :exla, "0.10.0", "93e7d75a774fbc06ce05b96de20c4b01bda413b315238cb3c727c09a05d2bc3a", [:make, :mix], [{:elixir_make, "~> 0.6", [hex: :elixir_make, repo: "hexpm", optional: false]}, {:fine, "~> 0.1.0", [hex: :fine, repo: "hexpm", optional: false]}, {:nimble_pool, "~> 1.0", [hex: :nimble_pool, repo: "hexpm", optional: false]}, {:nx, "~> 0.10.0", [hex: :nx, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}, {:xla, "~> 0.9.0", [hex: :xla, repo: "hexpm", optional: false]}], "hexpm", "16fffdb64667d7f0a3bc683fdcd2792b143a9b345e4b1f1d5cd50330c63d8119"},
+  "exla": {:hex, :exla, "0.11.0", "1428de9edcb297480a64611d3a72fcefe13c93c115bba6d38e910583c37e38c8", [:make, :mix], [{:elixir_make, "~> 0.6", [hex: :elixir_make, repo: "hexpm", optional: false]}, {:fine, "~> 0.1", [hex: :fine, repo: "hexpm", optional: false]}, {:nimble_pool, "~> 1.0", [hex: :nimble_pool, repo: "hexpm", optional: false]}, {:nx, "~> 0.11.0", [hex: :nx, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}, {:xla, "~> 0.10.0", [hex: :xla, repo: "hexpm", optional: false]}], "hexpm", "1067207c802bd6f28cded6a2664979ee2e25dddce95cb84be3f0a3ebfbab2c74"},
   "finch": {:hex, :finch, "0.21.0", "b1c3b2d48af02d0c66d2a9ebfb5622be5c5ecd62937cf79a88a7f98d48a8290c", [:mix], [{:mime, "~> 1.0 or ~> 2.0", [hex: :mime, repo: "hexpm", optional: false]}, {:mint, "~> 1.6.2 or ~> 1.7", [hex: :mint, repo: "hexpm", optional: false]}, {:nimble_options, "~> 0.4 or ~> 1.0", [hex: :nimble_options, repo: "hexpm", optional: false]}, {:nimble_pool, "~> 1.1", [hex: :nimble_pool, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "87dc6e169794cb2570f75841a19da99cfde834249568f2a5b121b809588a4377"},
-  "fine": {:hex, :fine, "0.1.5", "54880d13ab2c57884a105502d4fdce041f5852fe19bcfdcbcd7327225b1a6d5a", [:mix], [], "hexpm", "39f8f3f48a21e053c483f362b7b6a3bb38fdd987b31debc4d4e7a77fe8564335"},
+  "fine": {:hex, :fine, "0.1.6", "4bf7151493443c454aac9f2fa2f34f5fefd0346a83fb5586a016c4a135c63247", [:mix], [], "hexpm", "5638eb4495488e885ebec167fa57973e5c35e1a50c344eb7666c90ec1c4e3b12"},
   "hpax": {:hex, :hpax, "1.0.3", "ed67ef51ad4df91e75cc6a1494f851850c0bd98ebc0be6e81b026e765ee535aa", [:mix], [], "hexpm", "8eab6e1cfa8d5918c2ce4ba43588e894af35dbd8e91e6e55c817bca5847df34a"},
   "image": {:hex, :image, "0.65.0", "44908233a1a0dcdbb6ae873ec09fd9ae533d1840d300d8b0b1b186d586b935e6", [:mix], [{:color, "~> 0.4", [hex: :color, repo: "hexpm", optional: false]}, {:evision, "~> 0.1.33 or ~> 0.2", [hex: :evision, repo: "hexpm", optional: true]}, {:exla, "0.11.0", [hex: :exla, repo: "hexpm", optional: true]}, {:jason, "~> 1.4", [hex: :jason, repo: "hexpm", optional: true]}, {:kino, "~> 0.13", [hex: :kino, repo: "hexpm", optional: true]}, {:nx, "~> 0.11.0", [hex: :nx, repo: "hexpm", optional: true]}, {:nx_image, "~> 0.1", [hex: :nx_image, repo: "hexpm", optional: true]}, {:phoenix_html, "~> 2.1 or ~> 3.2 or ~> 4.0", [hex: :phoenix_html, repo: "hexpm", optional: false]}, {:plug, "~> 1.13", [hex: :plug, repo: "hexpm", optional: true]}, {:req, "~> 0.4", [hex: :req, repo: "hexpm", optional: true]}, {:rustler, "> 0.0.0", [hex: :rustler, repo: "hexpm", optional: true]}, {:scholar, "~> 0.3", [hex: :scholar, repo: "hexpm", optional: true]}, {:sweet_xml, "~> 0.7", [hex: :sweet_xml, repo: "hexpm", optional: false]}, {:vix, "~> 0.33", [hex: :vix, repo: "hexpm", optional: false]}, {:xav, "~> 0.10", [hex: :xav, repo: "hexpm", optional: true]}], "hexpm", "d2060e08d0f42564f49de1ea97a82a5d237f9ac91edb141dece51f1238dd8b4a"},
   "jason": {:hex, :jason, "1.4.4", "b9226785a9aa77b6857ca22832cffa5d5011a667207eb2a0ad56adb5db443b8a", [:mix], [{:decimal, "~> 1.0 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: true]}], "hexpm", "c5eb0cab91f094599f94d55bc63409236a8ec69a21a67814529e8d5f6cc90b3b"},
@@ -47,5 +47,5 @@
   "unpickler": {:hex, :unpickler, "0.1.0", "c2262c0819e6985b761e7107546cef96a485f401816be5304a65fdd200d5bd6a", [:mix], [], "hexpm", "e2b3f61e62406187ac52afead8a63bfb4e49394028993f3c4c42712743cab79e"},
   "unzip": {:hex, :unzip, "0.12.0", "beed92238724732418b41eba77dcb7f51e235b707406c05b1732a3052d1c0f36", [:mix], [], "hexpm", "95655b72db368e5a84951f0bed586ac053b55ee3815fd96062fce10ce4fc998d"},
   "vix": {:hex, :vix, "0.38.0", "77529ee4f6ced339c3d5f90a9eacf306f5b7109d3d1b5e3ef391a984ad404f75", [:make, :mix], [{:cc_precompiler, "~> 0.1.4 or ~> 0.2", [hex: :cc_precompiler, repo: "hexpm", optional: false]}, {:elixir_make, "~> 0.7.3 or ~> 0.8", [hex: :elixir_make, repo: "hexpm", optional: false]}, {:kino, "~> 0.7", [hex: :kino, repo: "hexpm", optional: true]}], "hexpm", "dca58f654922fa678d5df8e028317483d9c0f8acb2e2714076a8468695687aa7"},
-  "xla": {:hex, :xla, "0.9.1", "cca0040ff94902764007a118871bfc667f1a0085d4a5074533a47d6b58bec61e", [:make, :mix], [{:elixir_make, "~> 0.4", [hex: :elixir_make, repo: "hexpm", optional: false]}], "hexpm", "eb5e443ae5391b1953f253e051f2307bea183b59acee138053a9300779930daf"},
+  "xla": {:hex, :xla, "0.10.0", "41121e9f011456242d3a79b9289910ce43419be0b0e7ebe67cc1292c6b3f232f", [:make, :mix], [{:elixir_make, "~> 0.4", [hex: :elixir_make, repo: "hexpm", optional: false]}], "hexpm", "f57d91aea6e661b52bf12239316c598679e9170628122bbd941235f040122bc6"},
 }

test/classification_test.exs

@@ -2,13 +2,15 @@ defmodule Image.ClassificationTest do
   use ExUnit.Case, async: false
 
   @moduletag :ml
+  @moduletag :classification
 
   @images Path.join(__DIR__, "support/images")
 
   # Start the classifier serving once for the whole suite. Loading
   # ConvNeXt-tiny-224 takes several seconds, so we do it here rather
   # than per-test.
   setup_all do
+    Application.ensure_all_started(:exla)
     spec = Image.Classification.classifier()
     start_supervised!(spec)
     :ok
@@ -29,7 +31,7 @@ defmodule Image.ClassificationTest do
   describe "labels/2" do
     test "classifies a Cavalier King Charles Spaniel as a Blenheim spaniel" do
       image = Image.open!(Path.join(@images, "puppy.webp"))
-      labels = Image.Classification.labels(image)
+      labels = Image.Classification.labels(image, min_score: 0.1)
       assert "Blenheim spaniel" in labels
     end
 

test/detection_test.exs

@@ -2,6 +2,7 @@ defmodule Image.DetectionTest do
   use ExUnit.Case, async: false
 
   @moduletag :ml
+  @moduletag :ortex
 
   @corpus_dir Path.join(__DIR__, "support/images/segmentation")
   @corpus Path.join(@corpus_dir, "corpus.json")
@@ -41,7 +42,7 @@ defmodule Image.DetectionTest do
       for entry <- @corpus do
         image = open_image(entry)
         expected = entry["coco_class"]
-        detections = Image.Detection.detect(image)
+        detections = Image.Detection.detect(image, min_score: 0.1)
         labels = Enum.map(detections, & &1.label)
 
         assert Enum.any?(labels, &(&1 == expected)),
@@ -55,7 +56,7 @@ defmodule Image.DetectionTest do
         expected = entry["coco_class"]
         [gt_x, gt_y, gt_w, gt_h] = entry["prompt_box"]
 
-        detections = Image.Detection.detect(image)
+        detections = Image.Detection.detect(image, min_score: 0.1)
 
         best_iou =
           detections