Improved channel transposition process

PINTO0309 · Oct 29, 2022 · 578cdd0 · 578cdd0
1 parent c53052a
commit 578cdd0
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -71,7 +71,7 @@ Video speed is adjusted approximately 50 times slower than actual speed.
 $ docker run --rm -it \
 -v `pwd`:/workdir \
 -w /workdir \
-ghcr.io/pinto0309/onnx2tf:1.0.30
+ghcr.io/pinto0309/onnx2tf:1.0.31
 
 or
 

diff --git a/onnx2tf/__init__.py b/onnx2tf/__init__.py
@@ -1,3 +1,3 @@
 from onnx2tf.onnx2tf import convert, main
 
-__version__ = '1.0.30'
+__version__ = '1.0.31'
diff --git a/onnx2tf/utils/common_functions.py b/onnx2tf/utils/common_functions.py
@@ -513,6 +513,16 @@ def channel_transpose(
         and list(const_or_var_2.shape).count(1) == len(const_or_var_2.shape)-1 \
         and const_or_var_1.shape[-1] == max(const_or_var_2.shape):
         const_or_var_2 = const_or_var_2.squeeze()
+    elif isinstance(const_or_var_2, tf.Tensor) \
+        and len(const_or_var_2.shape) > 1 \
+        and list(const_or_var_2.shape).count(1) == len(const_or_var_2.shape)-1 \
+        and const_or_var_1.shape[-1] == max(const_or_var_2.shape):
+        const_or_var_2 = tf.squeeze(const_or_var_2)
+    elif tf.keras.backend.is_keras_tensor(const_or_var_2) \
+        and len(const_or_var_2.shape) > 1 \
+        and list(const_or_var_2.shape).count(1) == len(const_or_var_2.shape)-1 \
+        and const_or_var_1.shape[-1] == max(const_or_var_2.shape):
+        const_or_var_2 = tf.squeeze(const_or_var_2)
     return const_or_var_2