Python extensions for filter bank and stacker Signal Ops. (#2143)

Functionality to use the rest of the Signal Library OPs directly from python. Test with `bazel run python/tflite_micro/signal:stacker_op_test` and `bazel run python/tflite_micro/signal:filter_bank_ops_test` BUG=[287346710](http://b/287346710)
tensorflow · Jul 25, 2023 · e2cdefd · e2cdefd
1 parent bb3fda3
commit e2cdefd
Show file tree

Hide file tree

Showing 24 changed files with 4,136 additions and 1 deletion.
diff --git a/python/tflite_micro/python_ops_resolver.cc b/python/tflite_micro/python_ops_resolver.cc
@@ -53,6 +53,10 @@ PythonOpsResolver::PythonOpsResolver() {
   AddExpandDims();
   AddFftAutoScale();
   AddFill();
+  AddFilterBank();
+  AddFilterBankLog();
+  AddFilterBankSquareRoot();
+  AddFilterBankSpectralSubtraction();
   AddFloor();
   AddFloorDiv();
   AddFloorMod();

diff --git a/python/tflite_micro/signal/BUILD b/python/tflite_micro/signal/BUILD
@@ -18,8 +18,10 @@ cc_library(
         ":delay_op_cc",
         ":energy_op_cc",
         ":fft_ops_cc",
+        ":filter_bank_ops_cc",
         ":framer_op_cc",
         ":overlap_add_op_cc",
+        ":stacker_op_cc",
         ":window_op_cc",
     ],
 )
@@ -36,8 +38,10 @@ py_library(
         ":delay_op",
         ":energy_op",
         ":fft_ops",
+        ":filter_bank_ops",
         ":framer_op",
         ":overlap_add_op",
+        ":stacker_op",
         ":window_op",
     ],
 )
@@ -117,6 +121,40 @@ py_test(
     ],
 )
 
+py_tflm_signal_library(
+    name = "filter_bank_ops",
+    srcs = ["ops/filter_bank_ops.py"],
+    cc_op_defs = ["//signal/tensorflow_core/ops:filter_bank_ops"],
+    cc_op_kernels = [
+        "//signal/tensorflow_core/kernels:filter_bank_kernels",
+    ],
+    deps = [
+        "//python/tflite_micro/signal/utils:freq_to_mel",
+    ],
+)
+
+py_test(
+    name = "filter_bank_ops_test",
+    size = "small",
+    srcs = ["ops/filter_bank_ops_test.py"],
+    data = [
+        "//python/tflite_micro/signal/ops/testdata:filter_bank_accumulation_16k.txt",
+        "//python/tflite_micro/signal/ops/testdata:filter_bank_accumulation_44k.txt",
+        "//python/tflite_micro/signal/ops/testdata:filter_bank_accumulation_8k.txt",
+        "//python/tflite_micro/signal/ops/testdata:filter_bank_spectral_subtraction_test1.txt",
+        "//python/tflite_micro/signal/ops/testdata:filter_bank_square_root_test1.txt",
+        "//python/tflite_micro/signal/ops/testdata:filter_bank_test1.txt",
+    ],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    deps = [
+        ":filter_bank_ops",
+        "//python/tflite_micro/signal/utils:util",
+        requirement("numpy"),
+        requirement("tensorflow-cpu"),
+    ],
+)
+
 py_tflm_signal_library(
     name = "framer_op",
     srcs = ["ops/framer_op.py"],
@@ -167,6 +205,32 @@ py_test(
     ],
 )
 
+py_tflm_signal_library(
+    name = "stacker_op",
+    srcs = ["ops/stacker_op.py"],
+    cc_op_defs = ["//signal/tensorflow_core/ops:stacker_op"],
+    cc_op_kernels = [
+        "//signal/tensorflow_core/kernels:stacker_kernel",
+    ],
+)
+
+py_test(
+    name = "stacker_op_test",
+    size = "small",
+    srcs = ["ops/stacker_op_test.py"],
+    data = [
+        "//python/tflite_micro/signal/ops/testdata:stacker_test1.txt",
+    ],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    deps = [
+        ":stacker_op",
+        "//python/tflite_micro/signal/utils:util",
+        requirement("numpy"),
+        requirement("tensorflow-cpu"),
+    ],
+)
+
 py_tflm_signal_library(
     name = "window_op",
     srcs = ["ops/window_op.py"],