324 lines
10 KiB
Python
324 lines
10 KiB
Python
import os
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import List, Optional
|
|
|
|
import librosa
|
|
import numpy as np
|
|
from soundfile import write
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VocalDetectionConfig:
|
|
hop_ms: int = 20
|
|
smooth_ms: int = 200
|
|
start_ms: int = 120
|
|
end_ms: int = 200
|
|
prepad_ms: int = 80
|
|
postpad_ms: int = 120
|
|
min_len_ms: int = 1000
|
|
max_len_ms: int = 20000
|
|
short_seg_merge_gap_ms: int = 8000
|
|
small_gap_ms: int = 500
|
|
lookback_ms: int = 200
|
|
lookahead_ms: int = 100
|
|
|
|
|
|
def _moving_average(x: np.ndarray, win: int) -> np.ndarray:
|
|
if win <= 1:
|
|
return x
|
|
kernel = np.ones(win, dtype=np.float32) / float(win)
|
|
return np.convolve(x, kernel, mode="same")
|
|
|
|
|
|
def _merge_short_segments(
|
|
segments_ms: List[List[int]],
|
|
*,
|
|
min_len_ms: int,
|
|
max_len_ms: int,
|
|
short_seg_merge_gap_ms: int,
|
|
small_gap_ms: int,
|
|
) -> List[List[int]]:
|
|
if not segments_ms:
|
|
return []
|
|
|
|
merged: List[List[int]] = []
|
|
cur_start, cur_end = segments_ms[0]
|
|
|
|
for next_start, next_end in segments_ms[1:]:
|
|
cur_len = cur_end - cur_start
|
|
gap_ms = next_start - cur_end
|
|
merged_len = next_end - cur_start
|
|
|
|
should_merge = (
|
|
(cur_len < min_len_ms and gap_ms < short_seg_merge_gap_ms)
|
|
or (gap_ms < small_gap_ms and merged_len < max_len_ms)
|
|
)
|
|
|
|
if should_merge:
|
|
cur_end = next_end
|
|
continue
|
|
|
|
if (cur_end - cur_start) >= min_len_ms:
|
|
merged.append([cur_start, cur_end])
|
|
|
|
cur_start, cur_end = next_start, next_end
|
|
|
|
if (cur_end - cur_start) >= min_len_ms:
|
|
merged.append([cur_start, cur_end])
|
|
|
|
if not merged:
|
|
return segments_ms
|
|
|
|
return merged
|
|
|
|
|
|
def _voiced_to_segments(
|
|
voiced: np.ndarray,
|
|
*,
|
|
hop_ms: int,
|
|
smooth_ms: int,
|
|
start_ms: int,
|
|
end_ms: int,
|
|
prepad_ms: int,
|
|
postpad_ms: int,
|
|
max_len_ms: int,
|
|
) -> List[List[int]]:
|
|
smooth_frames = max(1, int(round(smooth_ms / hop_ms)))
|
|
smooth_voiced = _moving_average(voiced.astype(np.float32), smooth_frames)
|
|
active = smooth_voiced >= 0.5
|
|
|
|
segments: List[List[int]] = []
|
|
start_idx = None
|
|
start_frames = max(1, int(round(start_ms / hop_ms)))
|
|
end_frames = max(1, int(round(end_ms / hop_ms)))
|
|
prepad_frames = max(0, int(round(prepad_ms / hop_ms)))
|
|
postpad_frames = max(0, int(round(postpad_ms / hop_ms)))
|
|
active_count = 0
|
|
inactive_count = 0
|
|
|
|
for i, flag in enumerate(active):
|
|
if flag:
|
|
active_count += 1
|
|
inactive_count = 0
|
|
else:
|
|
inactive_count += 1
|
|
active_count = 0
|
|
|
|
if start_idx is None:
|
|
if active_count >= start_frames:
|
|
start_idx = max(0, i - start_frames + 1 - prepad_frames)
|
|
else:
|
|
if inactive_count >= end_frames:
|
|
end_idx = min(len(active) - 1, i - end_frames + 1 + postpad_frames)
|
|
start_ms_val = start_idx * hop_ms
|
|
end_ms_val = end_idx * hop_ms + hop_ms
|
|
if end_ms_val > start_ms_val:
|
|
segments.append([int(start_ms_val), int(end_ms_val)])
|
|
start_idx = None
|
|
|
|
if start_idx is not None:
|
|
start_ms_val = start_idx * hop_ms
|
|
end_idx = min(len(active) - 1, len(active) - 1 + postpad_frames)
|
|
end_ms_val = end_idx * hop_ms + hop_ms
|
|
if end_ms_val > start_ms_val:
|
|
segments.append([int(start_ms_val), int(end_ms_val)])
|
|
|
|
def _split_segment(seg: List[int]) -> List[List[int]]:
|
|
start_ms_val, end_ms_val = seg
|
|
start_frame = int(start_ms_val // hop_ms)
|
|
end_frame = int((end_ms_val - 1) // hop_ms)
|
|
end_frame = max(start_frame, min(end_frame, len(active) - 1))
|
|
|
|
best_start = None
|
|
best_len = 0
|
|
cur_start = None
|
|
cur_len = 0
|
|
for idx in range(start_frame, end_frame + 1):
|
|
if not active[idx]:
|
|
if cur_start is None:
|
|
cur_start = idx
|
|
cur_len = 1
|
|
else:
|
|
cur_len += 1
|
|
else:
|
|
if cur_start is not None and cur_len > best_len:
|
|
best_start, best_len = cur_start, cur_len
|
|
cur_start = None
|
|
cur_len = 0
|
|
if cur_start is not None and cur_len > best_len:
|
|
best_start, best_len = cur_start, cur_len
|
|
|
|
if best_start is None:
|
|
split_frame = (start_frame + end_frame) // 2
|
|
else:
|
|
split_frame = best_start + best_len // 2
|
|
|
|
split_ms = split_frame * hop_ms
|
|
if split_ms <= start_ms_val:
|
|
split_ms = start_ms_val + hop_ms
|
|
if split_ms >= end_ms_val:
|
|
split_ms = end_ms_val - hop_ms
|
|
|
|
if split_ms <= start_ms_val or split_ms >= end_ms_val:
|
|
return [seg]
|
|
|
|
return [[start_ms_val, int(split_ms)], [int(split_ms), end_ms_val]]
|
|
|
|
queue = segments[:]
|
|
segments = []
|
|
while queue:
|
|
seg = queue.pop(0)
|
|
if (seg[1] - seg[0]) <= max_len_ms:
|
|
segments.append(seg)
|
|
continue
|
|
parts = _split_segment(seg)
|
|
if len(parts) == 1:
|
|
segments.append(seg)
|
|
else:
|
|
queue = parts + queue
|
|
|
|
return segments
|
|
|
|
|
|
class VocalDetector:
|
|
"""Detect vocal segments based on f0 voiced decisions.
|
|
|
|
This component consumes a precomputed ``*_f0.npy`` track and
|
|
produces vocal segments (and cuts wav files) for downstream
|
|
transcription or singing voice tasks.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
cut_wavs_output_dir: str = "cut_wavs",
|
|
config: VocalDetectionConfig | None = None,
|
|
*,
|
|
verbose: bool = True,
|
|
):
|
|
"""Initialize the vocal detector.
|
|
|
|
Args:
|
|
cut_wavs_output_dir: Directory to save cut wav segments.
|
|
config: Detection configuration; uses :class:`VocalDetectionConfig` by default.
|
|
verbose: Whether to print verbose logs.
|
|
"""
|
|
self.cut_wavs_output_dir = cut_wavs_output_dir
|
|
self.config = config or VocalDetectionConfig()
|
|
self.verbose = verbose
|
|
|
|
if self.verbose:
|
|
print(
|
|
"[vocal detection] init success:",
|
|
f"cut_wavs_output_dir={self.cut_wavs_output_dir}",
|
|
f"hop_ms={self.config.hop_ms}",
|
|
)
|
|
|
|
def process(self, audio_path: str, f0: np.ndarray, *, verbose: Optional[bool] = None) -> List[dict]:
|
|
"""Run vocal detection on a single wav.
|
|
|
|
Args:
|
|
audio_path: Path to the input wav file.
|
|
f0: The f0 contour to use for vocal detection.
|
|
verbose: Override instance-level verbose flag for this call.
|
|
|
|
Returns:
|
|
A list of segment metadata dicts with fields like
|
|
``item_name``, ``wav_fn``, ``start_time_ms``, ``end_time_ms``.
|
|
"""
|
|
verbose = self.verbose if verbose is None else verbose
|
|
if verbose:
|
|
print(f"[vocal detection] process: start: {audio_path}")
|
|
t0 = time.time()
|
|
|
|
os.makedirs(self.cut_wavs_output_dir, exist_ok=True)
|
|
|
|
base_name = os.path.basename(audio_path)
|
|
base_name_no_ext = os.path.splitext(base_name)[0]
|
|
|
|
voiced = f0 > 0
|
|
|
|
segments_ms = _voiced_to_segments(
|
|
voiced,
|
|
hop_ms=self.config.hop_ms,
|
|
smooth_ms=self.config.smooth_ms,
|
|
start_ms=self.config.start_ms,
|
|
end_ms=self.config.end_ms,
|
|
prepad_ms=self.config.prepad_ms,
|
|
postpad_ms=self.config.postpad_ms,
|
|
max_len_ms=self.config.max_len_ms,
|
|
)
|
|
|
|
if verbose:
|
|
print(f"[vocal detection] segments(before_merge)={len(segments_ms)}")
|
|
|
|
segments_ms = _merge_short_segments(
|
|
segments_ms,
|
|
min_len_ms=self.config.min_len_ms,
|
|
max_len_ms=self.config.max_len_ms,
|
|
short_seg_merge_gap_ms=self.config.short_seg_merge_gap_ms,
|
|
small_gap_ms=self.config.small_gap_ms,
|
|
)
|
|
|
|
if verbose:
|
|
print(f"[vocal detection] segments(after_merge)={len(segments_ms)}")
|
|
|
|
y, sr = librosa.load(audio_path, sr=None, mono=True)
|
|
|
|
# Apply global lookback/lookahead in milliseconds
|
|
lookback_ms = self.config.lookback_ms
|
|
lookahead_ms = self.config.lookahead_ms
|
|
|
|
adjusted_segments: List[List[int]] = []
|
|
prev_end = 0
|
|
for start_ms, end_ms in segments_ms:
|
|
start_ms = max(0, start_ms - lookback_ms)
|
|
end_ms = min(end_ms + lookahead_ms, int(y.shape[0] / sr * 1000))
|
|
|
|
# Enforce non-overlap with previous segment, move backward the previous one.
|
|
if start_ms < prev_end and len(adjusted_segments) > 0:
|
|
adjusted_segments[-1][1] = start_ms
|
|
|
|
adjusted_segments.append([start_ms, end_ms])
|
|
prev_end = end_ms
|
|
|
|
segment_infos = []
|
|
for idx, (start_ms, end_ms) in enumerate(adjusted_segments):
|
|
if end_ms - start_ms > self.config.max_len_ms:
|
|
start_ms = end_ms - self.config.max_len_ms
|
|
|
|
key = f"{base_name_no_ext}_{idx}"
|
|
start_sample = librosa.time_to_samples(start_ms / 1000, sr=sr)
|
|
end_sample = librosa.time_to_samples(end_ms / 1000, sr=sr)
|
|
# Use Python int to avoid 32-bit overflow on Windows
|
|
start_sample = int(start_sample)
|
|
end_sample = int(end_sample)
|
|
segment = y[start_sample:end_sample]
|
|
|
|
write(f"{self.cut_wavs_output_dir}/{key}.wav", segment, sr)
|
|
segment_infos.append(
|
|
{
|
|
"item_name": key,
|
|
"wav_fn": f"{self.cut_wavs_output_dir}/{key}.wav",
|
|
"start_time_ms": int(start_sample * 1000 / sr),
|
|
"end_time_ms": int(end_sample * 1000 / sr),
|
|
"origin_wav_fn": audio_path,
|
|
"duration": int((end_sample - start_sample) * 1000 / sr),
|
|
}
|
|
)
|
|
|
|
if verbose:
|
|
dt = time.time() - t0
|
|
print(
|
|
"[vocal detection] process: done:",
|
|
f"n_segments={len(segment_infos)}",
|
|
f"time={dt:.3f}s",
|
|
)
|
|
|
|
return segment_infos
|
|
|
|
|
|
if __name__ == "__main__":
|
|
m = VocalDetector(cut_wavs_output_dir="outputs/transcription/cut_wavs")
|
|
segment_infos = m.process("example/audio/zh_prompt.mp3", np.load("example/audio/zh_prompt_f0.npy"))
|
|
print(segment_infos)
|