Add pickle support to MelBanksOptions.

This commit is contained in:
Fangjun Kuang 2021-11-04 10:40:58 +08:00
parent e5dd272963
commit 6eb7a3b243
3 changed files with 27 additions and 7 deletions

View File

@ -40,10 +40,11 @@ static void PybindFrameExtractionOptions(py::module &m) {
#endif
.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);
}));
.def(py::pickle(
[](const PyClass &self) -> py::dict { 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

@ -26,9 +26,15 @@ static void PybindMelBanksOptions(py::module &m) {
[](const PyClass &self) -> std::string { return self.ToString(); })
.def("as_dict",
[](const PyClass &self) -> py::dict { return AsDict(self); })
.def_static("from_dict", [](py::dict dict) -> PyClass {
return MelBanksOptionsFromDict(dict);
});
.def_static("from_dict",
[](py::dict dict) -> PyClass {
return MelBanksOptionsFromDict(dict);
})
.def(py::pickle(
[](const PyClass &self) -> py::dict { return AsDict(self); },
[](py::dict dict) -> PyClass {
return MelBanksOptionsFromDict(dict);
}));
}
void PybindMelComputations(py::module &m) { PybindMelBanksOptions(m); }

View File

@ -2,6 +2,8 @@
#
# Copyright (c) 2021 Xiaomi Corporation (authors: Fangjun Kuang)
import pickle
import kaldifeat
@ -82,12 +84,23 @@ def test_from_dict_full_and_as_dict():
assert opts3.htk_mode is True
def test_pickle():
opts = kaldifeat.MelBanksOptions()
opts.num_bins = 100
opts.low_freq = 22
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__":