Clean up code

software-mansion-labs · Jul 3, 2024 · 418def9 · 418def9
1 parent c1a32ee
commit 418def9
Show file tree

Hide file tree

Showing 7 changed files with 13 additions and 17 deletions.
diff --git a/lib/ex_vision/classification/efficientnet_v2_l.ex b/lib/ex_vision/classification/efficientnet_v2_l.ex
@@ -15,8 +15,8 @@ defmodule ExVision.Classification.EfficientNet_V2_L do
     image
     |> ExVision.Utils.resize({480, 480})
     |> NxImage.normalize(
-      Nx.tensor([0.5, 0.5, 0.5]),
-      Nx.tensor([0.5, 0.5, 0.5]),
+      Nx.f32([0.5, 0.5, 0.5]),
+      Nx.f32([0.5, 0.5, 0.5]),
       channels: :first
     )
   end

diff --git a/lib/ex_vision/classification/efficientnet_v2_m.ex b/lib/ex_vision/classification/efficientnet_v2_m.ex
@@ -15,8 +15,8 @@ defmodule ExVision.Classification.EfficientNet_V2_M do
     image
     |> ExVision.Utils.resize({480, 480})
     |> NxImage.normalize(
-      Nx.tensor([0.485, 0.456, 0.406]),
-      Nx.tensor([0.229, 0.224, 0.225]),
+      Nx.f32([0.485, 0.456, 0.406]),
+      Nx.f32([0.229, 0.224, 0.225]),
       channels: :first
     )
   end

diff --git a/lib/ex_vision/classification/efficientnet_v2_s.ex b/lib/ex_vision/classification/efficientnet_v2_s.ex
@@ -15,8 +15,8 @@ defmodule ExVision.Classification.EfficientNet_V2_S do
     image
     |> ExVision.Utils.resize({384, 384})
     |> NxImage.normalize(
-      Nx.tensor([0.485, 0.456, 0.406]),
-      Nx.tensor([0.229, 0.224, 0.225]),
+      Nx.f32([0.485, 0.456, 0.406]),
+      Nx.f32([0.229, 0.224, 0.225]),
       channels: :first
     )
   end

diff --git a/lib/ex_vision/classification/generic_classifier.ex b/lib/ex_vision/classification/generic_classifier.ex
@@ -15,7 +15,7 @@ defmodule ExVision.Classification.GenericClassifier do
     scores
     |> Nx.backend_transfer()
     |> Nx.flatten()
-    |> Utils.softmax()
+    |> Axon.Activations.softmax(axis: [0])
     |> Nx.to_flat_list()
     |> then(&Enum.zip(categories, &1))
     |> Map.new()

diff --git a/lib/ex_vision/classification/mobilenet_v3_small.ex b/lib/ex_vision/classification/mobilenet_v3_small.ex
@@ -15,8 +15,8 @@ defmodule ExVision.Classification.MobileNetV3Small do
     image
     |> ExVision.Utils.resize({224, 224})
     |> NxImage.normalize(
-      Nx.tensor([0.485, 0.456, 0.406]),
-      Nx.tensor([0.229, 0.224, 0.225]),
+      Nx.f32([0.485, 0.456, 0.406]),
+      Nx.f32([0.229, 0.224, 0.225]),
       channels: :first
     )
   end

diff --git a/lib/ex_vision/classification/squeezenet1_1.ex b/lib/ex_vision/classification/squeezenet1_1.ex
@@ -15,8 +15,8 @@ defmodule ExVision.Classification.SqueezeNet1_1 do
     image
     |> ExVision.Utils.resize({224, 224})
     |> NxImage.normalize(
-      Nx.tensor([0.485, 0.456, 0.406]),
-      Nx.tensor([0.229, 0.224, 0.225]),
+      Nx.f32([0.485, 0.456, 0.406]),
+      Nx.f32([0.229, 0.224, 0.225]),
       channels: :first
     )
   end

diff --git a/lib/ex_vision/utils.ex b/lib/ex_vision/utils.ex
@@ -86,8 +86,8 @@ defmodule ExVision.Utils do
 
   defp ensure_grad_3(tensor) do
     tensor
-    |> Nx.shape()
-    |> tuple_size()
+
+    Nx.rank()
     |> case do
       3 -> [tensor]
       4 -> tensor |> Nx.to_batched(1) |> Stream.map(&Nx.squeeze(&1, axes: [0])) |> Enum.to_list()
@@ -149,10 +149,6 @@ defmodule ExVision.Utils do
     Enum.map(outputs, fn {name, _type, _shape} -> name end)
   end
 
-  defn softmax(x) do
-    Nx.divide(Nx.exp(x), Nx.sum(Nx.exp(x)))
-  end
-
   @spec batched_run(atom(), ExVision.Model.input_t()) :: ExVision.Model.output_t()
   def batched_run(process_name, input) when is_list(input) do
     Nx.Serving.batched_run(process_name, input)