Add pickle support to FrameExtractionOptions.

This commit is contained in:
Fangjun Kuang 2021-11-04 10:35:52 +08:00
parent ec66d87fae
commit e5dd272963
3 changed files with 37 additions and 25 deletions

View File

@ -12,39 +12,38 @@
namespace kaldifeat {
static void PybindFrameExtractionOptions(py::module &m) {
py::class_<FrameExtractionOptions>(m, "FrameExtractionOptions")
using PyClass = FrameExtractionOptions;
py::class_<PyClass>(m, "FrameExtractionOptions")
.def(py::init<>())
.def_readwrite("samp_freq", &FrameExtractionOptions::samp_freq)
.def_readwrite("frame_shift_ms", &FrameExtractionOptions::frame_shift_ms)
.def_readwrite("frame_length_ms",
&FrameExtractionOptions::frame_length_ms)
.def_readwrite("dither", &FrameExtractionOptions::dither)
.def_readwrite("preemph_coeff", &FrameExtractionOptions::preemph_coeff)
.def_readwrite("remove_dc_offset",
&FrameExtractionOptions::remove_dc_offset)
.def_readwrite("window_type", &FrameExtractionOptions::window_type)
.def_readwrite("round_to_power_of_two",
&FrameExtractionOptions::round_to_power_of_two)
.def_readwrite("blackman_coeff", &FrameExtractionOptions::blackman_coeff)
.def_readwrite("snip_edges", &FrameExtractionOptions::snip_edges)
.def_readwrite("samp_freq", &PyClass::samp_freq)
.def_readwrite("frame_shift_ms", &PyClass::frame_shift_ms)
.def_readwrite("frame_length_ms", &PyClass::frame_length_ms)
.def_readwrite("dither", &PyClass::dither)
.def_readwrite("preemph_coeff", &PyClass::preemph_coeff)
.def_readwrite("remove_dc_offset", &PyClass::remove_dc_offset)
.def_readwrite("window_type", &PyClass::window_type)
.def_readwrite("round_to_power_of_two", &PyClass::round_to_power_of_two)
.def_readwrite("blackman_coeff", &PyClass::blackman_coeff)
.def_readwrite("snip_edges", &PyClass::snip_edges)
.def("as_dict",
[](const FrameExtractionOptions &self) -> py::dict {
return AsDict(self);
})
[](const PyClass &self) -> py::dict { return AsDict(self); })
.def_static("from_dict",
[](py::dict dict) -> FrameExtractionOptions {
[](py::dict dict) -> PyClass {
return FrameExtractionOptionsFromDict(dict);
})
#if 0
.def_readwrite("allow_downsample",
&FrameExtractionOptions::allow_downsample)
.def_readwrite("allow_upsample", &FrameExtractionOptions::allow_upsample)
&PyClass::allow_downsample)
.def_readwrite("allow_upsample", &PyClass::allow_upsample)
.def_readwrite("max_feature_vectors",
&FrameExtractionOptions::max_feature_vectors)
&PyClass::max_feature_vectors)
#endif
.def("__str__", [](const FrameExtractionOptions &self) -> std::string {
return self.ToString();
});
.def("__str__",
[](const PyClass &self) -> std::string { return self.ToString(); })
.def(py::pickle([](const PyClass &self) { return AsDict(self); },
[](py::dict dict) -> PyClass {
return FrameExtractionOptionsFromDict(dict);
}));
m.def("num_frames", &NumFrames, py::arg("num_samples"), py::arg("opts"),
py::arg("flush") = true);

View File

@ -15,7 +15,7 @@
/*
* This file contains code about `from_dict` and
* `to_dict` for various options in kaldifeat.
* `as_dict` for various options in kaldifeat.
*
* Regarding `from_dict`, users don't need to provide
* all the fields in the options. If some fields

View File

@ -2,6 +2,8 @@
#
# Copyright (c) 2021 Xiaomi Corporation (authors: Fangjun Kuang)
import pickle
import kaldifeat
@ -94,12 +96,23 @@ def test_from_dict_full_and_as_dict():
assert opts3.window_type == "hanning"
def test_pickle():
opts = kaldifeat.FrameExtractionOptions()
opts.samp_freq = 44100
opts.dither = 5.5
data = pickle.dumps(opts)
opts2 = pickle.loads(data)
assert str(opts) == str(opts2)
def main():
test_default()
test_set_get()
test_from_empty_dict()
test_from_dict_partial()
test_from_dict_full_and_as_dict()
test_pickle()
if __name__ == "__main__":