Style transfer

lib/ex_vision/model/definition.ex

@@ -26,14 +26,14 @@ defmodule ExVision.Model.Definition do
 
     options =
       Keyword.validate!(options, [
-        :categories,
+        categories: nil,
         name: module_to_name(__CALLER__.module)
       ])
 
     quote do
-      unless is_nil(unquote(options[:categories])) do
-        use ExVision.Model.Definition.Parts.WithCategories, unquote(options)
-      end
+      # conditional defintion based on whether `categories` option is present has to be moved inside __using__ macro
+      # here is explenation https://cocoa-research.works/2022/10/conditional-compliation-with-if-and-use-in-elixir/
+      use ExVision.Model.Definition.Parts.WithCategories, unquote(options)
 
       @behaviour ExVision.Model.Definition
 

lib/ex_vision/model/definition/parts/with_categories.ex

@@ -5,22 +5,24 @@ defmodule ExVision.Model.Definition.Parts.WithCategories do
 
   defmacro __using__(options) do
     options = Keyword.validate!(options, [:name, :categories])
-    categories = options |> Keyword.fetch!(:categories) |> Utils.load_categories()
-    spec = categories |> Enum.uniq() |> Bunch.Typespec.enum_to_alternative()
+    unless is_nil(options |> Keyword.fetch!(:categories)) do
+      categories = options |> Keyword.fetch!(:categories) |> Utils.load_categories()
+      spec = categories |> Enum.uniq() |> Bunch.Typespec.enum_to_alternative()
 
-    quote do
-      require Bunch.Typespec
+      quote do
+        require Bunch.Typespec
 
-      @typedoc """
-      Type describing all categories recognised by #{unquote(options[:name])}
-      """
-      @type category_t() :: unquote(spec)
+        @typedoc """
+        Type describing all categories recognised by #{unquote(options[:name])}
+        """
+        @type category_t() :: unquote(spec)
 
-      @doc """
-      Returns a list of all categories recognised by #{unquote(options[:name])}
-      """
-      @spec categories() :: [category_t()]
-      def categories(), do: unquote(categories)
+        @doc """
+        Returns a list of all categories recognised by #{unquote(options[:name])}
+        """
+        @spec categories() :: [category_t()]
+        def categories(), do: unquote(categories)
+      end
     end
   end
 end

lib/ex_vision/style_transfer/style_transfer.ex

@@ -0,0 +1,56 @@
+defmodule Configuration do
+  @low_resolution {400,300}
+  @high_resolution {640,480}
+  def configuration do
+    %{
+      ExVision.StyleTransfer.Candy  => [model: "candy.onnx", resolution: @high_resolution],
+      ExVision.StyleTransfer.CandyFast  => [model: "candy_fast.onnx", resolution: @low_resolution],
+      ExVision.StyleTransfer.Princess  => [model: "princess.onnx", resolution: @high_resolution],
+      ExVision.StyleTransfer.PrincessFast  => [model: "princess_fast.onnx", resolution: @low_resolution],
+      ExVision.StyleTransfer.Udnie  => [model: "udnie.onnx", resolution: @high_resolution],
+      ExVision.StyleTransfer.UdnieFast  => [model: "udnie_fast.onnx", resolution: @low_resolution],
+      ExVision.StyleTransfer.Mosaic  => [model: "mosaic.onnx", resolution: @high_resolution],
+      ExVision.StyleTransfer.MosaicFast  => [model: "mosaic_fast.onnx", resolution: @low_resolution],
+    }
+  end
+end
+
+for {module, opts} <- Configuration.configuration() do
+  defmodule module do
+    @moduledoc """
+    #{module} is a custom style transfer model optimised for devices with low computational capabilities and CPU inference.
+    """
+    require Logger
+    @type output_t() :: [Nx.Tensor.t()]
+
+    use ExVision.Model.Definition.Ortex, model: unquote(opts[:model])
+
+    @impl true
+    def load(options \\ []) do
+      if Keyword.has_key?(options, :batch_size) do
+        Logger.warning(
+          "`:max_batch_size` was given, but this model can only process batch of size 1. Overriding"
+        )
+      end
+
+      options
+      |> Keyword.put(:batch_size, 1)
+      |> default_model_load()
+    end
+
+    @impl true
+    def preprocessing(img, _metdata) do
+      ExVision.Utils.resize(img, unquote(opts[:resolution])) |> Nx.divide(255.0)
+    end
+
+    @impl true
+    def postprocessing(
+          stylized_frame,
+          metadata
+        ) do
+
+      stylized_frame = stylized_frame["55"]
+      NxImage.resize(stylized_frame, metadata.original_size, channels: :first)
+    end
+  end
+end

mix.exs

@@ -99,6 +99,14 @@ defmodule ExVision.Mixfile do
           ExVision.Classification.EfficientNet_V2_L,
           ExVision.Classification.SqueezeNet1_1,
           ExVision.SemanticSegmentation.DeepLabV3_MobileNetV3,
+          ExVision.StyleTransfer.Candy,
+          ExVision.StyleTransfer.CandyFast,
+          ExVision.StyleTransfer.Udnie,
+          ExVision.StyleTransfer.UdnieFast,
+          ExVision.StyleTransfer.Mosaic,
+          ExVision.StyleTransfer.MosaicFast,
+          ExVision.StyleTransfer.Princess,
+          ExVision.StyleTransfer.PrincessFast,
           ExVision.InstanceSegmentation.MaskRCNN_ResNet50_FPN_V2,
           ExVision.ObjectDetection.Ssdlite320_MobileNetv3,
           ExVision.ObjectDetection.FasterRCNN_ResNet50_FPN,
@@ -123,6 +131,7 @@ defmodule ExVision.Mixfile do
         ExVision.Types,
         ExVision.Classification,
         ExVision.SemanticSegmentation,
+        ExVision.StyleTransfer,
         ExVision.InstanceSegmentation,
         ExVision.ObjectDetection,
         ExVision.KeypointDetection