Merge pull request #186 from cocoa-xu/cx-fix-outputBlob-in-list

Fixed outputBlob in a list
cocoa-xu · Mar 24, 2023 · 91fc18e · 91fc18e
2 parents 0c2550c + bc5fc64
commit 91fc18e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/lib/zoo/image_classification/mobilenet_v1.ex b/lib/zoo/image_classification/mobilenet_v1.ex
@@ -114,6 +114,13 @@ defmodule Evision.Zoo.ImageClassification.MobileNetV1 do
     inputBlob = preprocess(image)
     Evision.DNN.Net.setInput(self, inputBlob)
     outputBlob = Evision.DNN.Net.forward(self, outputName: "")
+    outputBlob =
+      if is_list(outputBlob) do
+        [outputBlob] = outputBlob
+        outputBlob
+      else
+        outputBlob
+      end
     # todo: use Evision.Backend when Nx.slice is implemented
     result = Nx.squeeze(Evision.Mat.to_nx(outputBlob, Nx.BinaryBackend))
     Nx.to_flat_list(Nx.argsort(result, direction: :desc)[[0..top_k-1]])

diff --git a/lib/zoo/text_recognition/crnn.ex b/lib/zoo/text_recognition/crnn.ex
@@ -136,6 +136,13 @@ defmodule Evision.Zoo.TextRecognition.CRNN do
     # Forward
     Evision.DNN.Net.setInput(self, inputBlob)
     outputBlob = Evision.DNN.Net.forward(self, outputName: "")
+    outputBlob =
+      if is_list(outputBlob) do
+        [outputBlob] = outputBlob
+        outputBlob
+      else
+        outputBlob
+      end
 
     # Postprocess
     charset = opts[:charset]