Skip to content

feat(breath): support BreathLab v6 model with optional EP annotation #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Xiantaidu wants to merge 1 commit into from
Closed

feat(breath): support BreathLab v6 model with optional EP annotation #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions scripts/convert_breath_to_gguf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -368,7 +368,7 @@ def put(key: str, value: Any) -> None:

for key in ("threshold", "median_filter_frames", "min_dur_ms", "merge_gap_ms",
"sp_threshold", "sp_floor_percentile", "sp_floor_margin_db",
"sp_hf_guard_db", "sp_min_dur_ms"):
"sp_hf_guard_db", "sp_min_dur_ms", "ep_threshold"):
if key in pp:
put(f"breath.postprocess.{key}", pp[key])

Expand Down Expand Up @@ -448,15 +448,25 @@ def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument("--model-dir", type=pathlib.Path, required=True,
help="directory holding <name>.onnx + <name>.meta.json")
parser.add_argument("--name", required=True,
help="model stem, e.g. v5_24k")
parser.add_argument("--name", default=None,
help="model stem, e.g. model or v5_24k (default: autodetected from model-dir)")
parser.add_argument("-o", "--output", type=pathlib.Path, required=True,
help="output .gguf path")
parser.add_argument("--dtype", choices=sorted(_DTYPE_MAP), default="f32",
help="base dtype for the weight tensors (default: f32)")
parser.add_argument("-v", "--verbose", action="store_true")
args = parser.parse_args(argv)

if not args.name:
if (args.model_dir / "model.onnx").exists():
args.name = "model"
else:
candidates = list(args.model_dir.glob("*.onnx"))
if len(candidates) == 1:
args.name = candidates[0].stem
else:
parser.error("--name is required when model-dir contains multiple or no .onnx files")

logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO,
format="%(levelname)s %(message)s")
convert(args.model_dir, args.name, args.output, base_dtype=args.dtype)
Expand Down
72 changes: 64 additions & 8 deletions src/breath/breath.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ struct PostProcess {
float sp_floor_margin_db = 8.0f;
float sp_hf_guard_db = 6.0f;
float sp_min_dur_ms = 120.0f;
float ep_threshold = 0.5f;
};

// Python's round() is half-to-even; every constant in the shipped models lands
Expand Down Expand Up @@ -243,16 +244,19 @@ struct BreathModel::Impl {
int hop_frames = 600;
int head_ap = 0;
int head_sp = -1;
int head_ep = -1;

// features + windowed inference with the overlap average
void infer(const float * wav, std::size_t n, int sample_rate,
BreathFeatures & feats, std::vector<float> & ap_prob,
std::vector<float> & sp_prob) const;
std::vector<float> & sp_prob,
std::vector<float> * ep_prob = nullptr) const;
};

void BreathModel::Impl::infer(const float * wav, std::size_t n, int sample_rate,
BreathFeatures & feats, std::vector<float> & ap_prob,
std::vector<float> & sp_prob) const {
std::vector<float> & sp_prob,
std::vector<float> * ep_prob) const {
if (!wav || n == 0) throw InvalidArgument("breath: empty waveform");

// 1. resample to the model rate (the reference `load_wav` does the same)
Expand Down Expand Up @@ -282,6 +286,7 @@ void BreathModel::Impl::infer(const float * wav, std::size_t n, int sample_rate,
const int n_out = net.config().n_out;
std::vector<double> ap_acc(static_cast<std::size_t>(T), 0.0);
std::vector<double> sp_acc(static_cast<std::size_t>(T), 0.0);
std::vector<double> ep_acc(static_cast<std::size_t>(T), 0.0);
std::vector<double> wsum(static_cast<std::size_t>(T), 0.0);
for (int s : starts) {
const int w = std::min(win, T - s);
Expand All @@ -295,6 +300,9 @@ void BreathModel::Impl::infer(const float * wav, std::size_t n, int sample_rate,
if (head_sp >= 0 && head_sp < n_out) {
sp_acc[idx] += out[static_cast<std::size_t>(head_sp) + n_out * t];
}
if (head_ep >= 0 && head_ep < n_out) {
ep_acc[idx] += out[static_cast<std::size_t>(head_ep) + n_out * t];
}
wsum[idx] += 1.0;
}
}
Expand All @@ -313,6 +321,17 @@ void BreathModel::Impl::infer(const float * wav, std::size_t n, int sample_rate,
static_cast<float>(sp_acc[static_cast<std::size_t>(t)] / den);
}
}
if (ep_prob) {
ep_prob->clear();
if (head_ep >= 0 && head_ep < n_out) {
ep_prob->resize(static_cast<std::size_t>(T));
for (int t = 0; t < T; ++t) {
const double den = std::max(wsum[static_cast<std::size_t>(t)], 1e-8);
(*ep_prob)[static_cast<std::size_t>(t)] =
static_cast<float>(ep_acc[static_cast<std::size_t>(t)] / den);
}
}
}
}

// ---------------------------------------------------------------------------
Expand Down Expand Up @@ -369,6 +388,7 @@ BreathModel BreathModel::load(const std::string & gguf_path) {
pp.sp_floor_margin_db = get_float(gguf, "breath.postprocess.sp_floor_margin_db", 8.0f);
pp.sp_hf_guard_db = get_float(gguf, "breath.postprocess.sp_hf_guard_db", 6.0f);
pp.sp_min_dur_ms = get_float(gguf, "breath.postprocess.sp_min_dur_ms", 120.0f);
pp.ep_threshold = get_float(gguf, "breath.postprocess.ep_threshold", 0.5f);

impl.win_frames = static_cast<int>(std::lround(
get_float(gguf, "breath.eval.infer_window_sec", 12.0f) * impl.fps));
Expand All @@ -381,26 +401,35 @@ BreathModel BreathModel::load(const std::string & gguf_path) {
// energy + high-band rule in full_segmentation
const auto sp = gguf.get_int_opt("breath.head.sp");
impl.head_sp = sp ? static_cast<int>(*sp) : -1;
const auto ep = gguf.get_int_opt("breath.head.ep");
impl.head_ep = ep ? static_cast<int>(*ep) : -1;

impl.net = BreathNet::load(gguf_path);
return m;
}

bool BreathModel::has_ep() const noexcept { return impl_->head_ep >= 0; }
float BreathModel::ep_threshold() const noexcept { return impl_->pp.ep_threshold; }

void BreathModel::probabilities(const float * wav, std::size_t n, int sample_rate,
std::vector<float> & ap_prob,
std::vector<float> & sp_prob) const {
std::vector<float> & sp_prob,
std::vector<float> * ep_prob) const {
BreathFeatures feats;
impl_->infer(wav, n, sample_rate, feats, ap_prob, sp_prob);
impl_->infer(wav, n, sample_rate, feats, ap_prob, sp_prob, ep_prob);
}

void BreathModel::run(const float * wav, std::size_t n, int sample_rate,
std::vector<BreathEvent> & ap_events,
std::vector<BreathSegment> & segments) const {
std::vector<BreathSegment> & segments,
std::vector<BreathEvent> * ep_events,
bool include_ep) const {
const auto & impl = *impl_;

BreathFeatures feats;
std::vector<float> ap_prob, sp_prob;
impl.infer(wav, n, sample_rate, feats, ap_prob, sp_prob);
std::vector<float> ap_prob, sp_prob, ep_prob;
const bool want_ep = (ep_events != nullptr || include_ep) && (impl.head_ep >= 0);
impl.infer(wav, n, sample_rate, feats, ap_prob, sp_prob, want_ep ? &ep_prob : nullptr);

std::vector<char> cls;
const auto runs = full_segmentation(ap_prob, sp_prob, feats.energy_db, feats.hf_db,
Expand All @@ -414,14 +443,41 @@ void BreathModel::run(const float * wav, std::size_t n, int sample_rate,
static_cast<float>(r.second * inv_fps)});
}

std::vector<std::pair<int, int>> ep_runs;
if (want_ep && !ep_prob.empty()) {
ep_runs = prob_to_events(ep_prob, impl.fps, impl.pp.ep_threshold, impl.pp);
if (ep_events) {
ep_events->clear();
ep_events->reserve(ep_runs.size());
for (const auto & r : ep_runs) {
ep_events->push_back(BreathEvent{static_cast<float>(r.first * inv_fps),
static_cast<float>(r.second * inv_fps)});
}
}
}

if (include_ep && !ep_runs.empty()) {
const int nf = static_cast<int>(cls.size());
for (const auto & r : ep_runs) {
const int s = std::max(0, r.first);
const int e = std::min(nf, r.second);
for (int t = s; t < e; ++t) {
if (cls[static_cast<std::size_t>(t)] != 1) {
cls[static_cast<std::size_t>(t)] = 3;
}
}
}
}

segments.clear();
const int nf = static_cast<int>(cls.size());
int i = 0;
while (i < nf) {
int j = i;
while (j < nf && cls[static_cast<std::size_t>(j)] == cls[static_cast<std::size_t>(i)]) ++j;
const char * name = cls[static_cast<std::size_t>(i)] == 1 ? "AP"
: cls[static_cast<std::size_t>(i)] == 2 ? "SP" : "V";
: cls[static_cast<std::size_t>(i)] == 2 ? "SP"
: cls[static_cast<std::size_t>(i)] == 3 ? "EP" : "V";
segments.push_back(BreathSegment{name, static_cast<float>(i * inv_fps),
static_cast<float>(j * inv_fps)});
i = j;
Expand Down
19 changes: 13 additions & 6 deletions src/breath/breath.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,21 +48,28 @@ class BreathModel {
BreathModel & operator=(const BreathModel &) = delete;

// AP events + the full AP/SP/V timeline for one mono waveform (any rate).
// When include_ep is true or ep_events is non-null, optional EP events
// are extracted for models with an EP head (BreathLab v6).
void run(const float * wav, std::size_t n, int sample_rate,
std::vector<BreathEvent> & ap_events,
std::vector<BreathSegment> & segments) const;
std::vector<BreathSegment> & segments,
std::vector<BreathEvent> * ep_events = nullptr,
bool include_ep = false) const;

// ---- diagnostics -----------------------------------------------------

// Window-averaged per-frame probabilities ([T] each; `sp` is empty when
// the model has no SP head).
// the model has no SP head, `ep_prob` populated when non-null and model has EP).
void probabilities(const float * wav, std::size_t n, int sample_rate,
std::vector<float> & ap_prob,
std::vector<float> & sp_prob) const;
std::vector<float> & sp_prob,
std::vector<float> * ep_prob = nullptr) const;

int sample_rate() const noexcept; // rate the model runs at
int fps() const noexcept; // frames per second
float threshold() const noexcept; // AP threshold from the GGUF
int sample_rate() const noexcept; // rate the model runs at
int fps() const noexcept; // frames per second
float threshold() const noexcept; // AP threshold from the GGUF
bool has_ep() const noexcept; // whether model carries an EP head
float ep_threshold() const noexcept; // EP threshold from the GGUF
const char * backend_name() const noexcept;

private:
Expand Down
51 changes: 44 additions & 7 deletions src/cli/main.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ Breathe options:
-> align --textgrid)
--phones-tier NAME tier to merge into (default: phones)
--min-insert-ms MS shortest AP/SP inserted (default: 50)
--ep detect and annotate EP (exhale) events (BreathLab v6)
-q, --quiet only report errors

Align options:
Expand Down Expand Up @@ -540,6 +541,7 @@ struct BreatheOptions {
std::string merge; // alignment TextGrid file or directory
std::string phones_tier = "phones";
double min_insert_ms = 50.0;
bool ep = false;
bool quiet = false;
};

Expand Down Expand Up @@ -664,9 +666,11 @@ int cmd_breathe(const std::string & input, const BreatheOptions & opt) {

const auto t0 = std::chrono::steady_clock::now();
std::vector<BreathEvent> ap_events;
std::vector<BreathEvent> ep_events;
std::vector<BreathSegment> segments;
model.run(audio_buf.samples.data(), audio_buf.samples.size(),
audio_buf.sample_rate, ap_events, segments);
audio_buf.sample_rate, ap_events, segments,
opt.ep ? &ep_events : nullptr, opt.ep);
const auto t1 = std::chrono::steady_clock::now();
const double ms = std::chrono::duration<double, std::milli>(t1 - t0).count();

Expand All @@ -685,11 +689,26 @@ int cmd_breathe(const std::string & input, const BreatheOptions & opt) {
}

if (opt.merge.empty()) {
std::vector<TextGridTier> out_tiers;
out_tiers.push_back(std::move(breath_tier));
if (opt.ep && !ep_events.empty()) {
TextGridTier ep_tier;
ep_tier.name = "EP";
ep_tier.intervals.reserve(ep_events.size());
for (const BreathEvent & ev : ep_events) {
TextGridInterval interval;
interval.xmin = ev.start;
interval.xmax = ev.end;
interval.text = "EP";
ep_tier.intervals.push_back(std::move(interval));
}
out_tiers.push_back(std::move(ep_tier));
}
write_textgrid_file(
(out_dir / (identifier + ".breath.TextGrid")).string(),
{ std::move(breath_tier) }, seconds);
std::move(out_tiers), seconds);
} else {
// 2PASS step 2: fold the AP/SP segments into the phones
// 2PASS step 2: fold the AP/SP/EP segments into the phones
// tier of the first-pass alignment and write the enriched
// TextGrid the second `align --textgrid` pass consumes.
fs::path alignment = opt.merge;
Expand Down Expand Up @@ -717,8 +736,8 @@ int cmd_breathe(const std::string & input, const BreatheOptions & opt) {
(out_dir / (identifier + ".TextGrid")).string(),
{ std::move(merged_tier), std::move(breath_tier) }, seconds);
if (!opt.quiet) {
std::fprintf(stderr, "breath: merged %zu AP/SP into %zu phones\n",
inserted, phones.size());
std::fprintf(stderr, "breath: merged %zu %s into %zu phones\n",
inserted, opt.ep ? "AP/SP/EP" : "AP/SP", phones.size());
}
}
}
Expand All @@ -740,13 +759,22 @@ int cmd_breathe(const std::string & input, const BreatheOptions & opt) {
};

Value ap_list = arr();
Value ep_list = arr();
Value seg_list = arr();
for (const BreathEvent & event : ap_events) {
Value item = obj();
item.obj.push_back({ "start", num(event.start) });
item.obj.push_back({ "end", num(event.end) });
ap_list.arr.push_back(std::move(item));
}
if (opt.ep) {
for (const BreathEvent & event : ep_events) {
Value item = obj();
item.obj.push_back({ "start", num(event.start) });
item.obj.push_back({ "end", num(event.end) });
ep_list.arr.push_back(std::move(item));
}
}
for (const BreathSegment & segment : segments) {
Value item = obj();
item.obj.push_back({ "label", text(segment.label) });
Expand All @@ -758,15 +786,23 @@ int cmd_breathe(const std::string & input, const BreatheOptions & opt) {
doc.obj.push_back({ "sample_rate", num(audio_buf.sample_rate) });
doc.obj.push_back({ "duration", num(seconds) });
doc.obj.push_back({ "ap", std::move(ap_list) });
if (opt.ep) {
doc.obj.push_back({ "ep", std::move(ep_list) });
}
doc.obj.push_back({ "segments", std::move(seg_list) });
std::ofstream out(out_dir / (identifier + ".breath.json"), std::ios::binary);
out << bren_json::dump(doc, 2) << "\n";
}

++ok;
if (!opt.quiet) {
std::printf("%-40s %6.1f s %3zu AP events %4zu segments %7.1f ms\n",
identifier.c_str(), seconds, ap_events.size(), segments.size(), ms);
if (opt.ep) {
std::printf("%-40s %6.1f s %3zu AP %3zu EP %4zu segments %7.1f ms\n",
identifier.c_str(), seconds, ap_events.size(), ep_events.size(), segments.size(), ms);
} else {
std::printf("%-40s %6.1f s %3zu AP events %4zu segments %7.1f ms\n",
identifier.c_str(), seconds, ap_events.size(), segments.size(), ms);
}
}
} catch (const std::exception & e) {
++skipped;
Expand Down Expand Up @@ -848,6 +884,7 @@ int main(int argc, char ** argv) {
else if (a == "--merge") opt.merge = value();
else if (a == "--phones-tier") opt.phones_tier = value();
else if (a == "--min-insert-ms") opt.min_insert_ms = std::strtod(value().c_str(), nullptr);
else if (a == "--ep") opt.ep = true;
else if (a == "-q" || a == "--quiet") opt.quiet = true;
else throw InvalidArgument("unknown option: " + a);
}
Expand Down
4 changes: 2 additions & 2 deletions src/cli/phone_input.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,7 @@ std::vector<TimedInterval> merge_breath_into_phones(
struct Insert { double xmin, xmax; const std::string * label; };
std::vector<Insert> inserts;
for (const tifa_ggml::BreathSegment & s : segments) {
if ((s.label == "AP" || s.label == "SP") && s.end - s.start >= min_insert) {
if ((s.label == "AP" || s.label == "SP" || s.label == "EP") && s.end - s.start >= min_insert) {
inserts.push_back({s.start, s.end, &s.label});
}
}
Expand Down Expand Up @@ -253,7 +253,7 @@ std::vector<TimedInterval> merge_breath_into_phones(
// adjacent phones of any other label is not ours to decide).
constexpr double kBreathMergeGap = 0.1;
const auto is_breath_label = [](const std::string & text) {
return text == "AP" || text == "SP" || text == "br" || text == "sil"
return text == "AP" || text == "SP" || text == "EP" || text == "br" || text == "sil"
|| text == "pau";
};
std::vector<TimedInterval> deduped;
Expand Down
17 changes: 17 additions & 0 deletions tests/test_breathe_merge.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,3 +107,20 @@ TEST(BreatheMerge, VoiceSegmentsAreNeverInserted) {
ASSERT_EQ(render(out), render(std::vector<TimedInterval>{iv(0.0, 1.0, "a")}));
EXPECT_EQ(inserted, 0u);
}

TEST(BreatheMerge, InsertsExhaleSegmentsWhenPresent) {
// phone 'a' followed by silence, with an EP detected between them
const std::vector<TimedInterval> phones{
iv(0.0, 1.0, "a"), iv(1.0, 2.0, "SP")};
std::size_t inserted = 0;
const auto out = merge_breath_into_phones(
phones, {seg("EP", 0.9, 1.2)}, 0.05, &inserted);
// straddling phone boundary [0.9, 1.2] vs [0.0, 1.0]: not inside phone, but if inside SP or phone:
// let's test EP inside the trailing SP [1.0, 1.3]
const auto out2 = merge_breath_into_phones(
phones, {seg("EP", 1.0, 1.3)}, 0.05, &inserted);
ASSERT_EQ(render(out2), render(std::vector<TimedInterval>{
iv(0.0, 1.0, "a"), iv(1.0, 1.3, "EP"), iv(1.3, 2.0, "SP")}));
EXPECT_EQ(inserted, 1u);
}

Loading