Don't transpose on the way back to Image

Author: kipcole9

lib/segmentation.ex

@@ -510,8 +510,8 @@ if ImageVision.ortex_configured?() do
         |> Nx.greater(0)
         |> Nx.multiply(255)
         |> Nx.as_type(:u8)
-        |> Nx.transpose()
         |> Nx.new_axis(2)
+        |> Nx.rename([:height, :width, :bands])
 
       binary
       |> Image.from_nx!()
@@ -606,8 +606,8 @@ if ImageVision.ortex_configured?() do
             |> Nx.greater(0.5)
             |> Nx.multiply(255)
             |> Nx.as_type(:u8)
-            |> Nx.transpose()
             |> Nx.new_axis(2)
+            |> Nx.rename([:height, :width, :bands])
             |> Image.from_nx!()
             |> Image.resize!(orig_w / mask_w, vertical_scale: orig_h / mask_h)
 

livebooks/image_vision.livemd

@@ -151,7 +151,7 @@ image =
 segments = Image.Segmentation.segment_panoptic(image)
 
 # Render a coloured overlay — one colour per segment label.
-{:ok, overlay} = Image.Segmentation.compose_overlay(image, segments)
+overlay = Image.Segmentation.compose_overlay(image, segments)
 
 Kino.Layout.grid(
   [