Coverage for tts_utils.py: 72%
355 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Utility functions for TTS processing, including silence detection and audio chunking.
3"""
4import os
5import re
6import io
7import base64
8import math
9import wave
10import tempfile
12import numpy as np
14from typing import List
15from typing import Tuple
16from typing import Optional
18from scipy.io import wavfile
21SILENCE_THRESHOLDS_MS = [200, 100, 50, 10] # thresholds for silence duration in milliseconds
22# AMPLITUDE_THRESHOLD = 500 # amplitude threshold for silence
23AMPLITUDE_THRESHOLD = 200 # amplitude threshold for silence
24MIN_SILENCE_DURATION = 0.2 # seconds
27def detect_silences(
28 data: np.ndarray,
29 rate: int,
30 amp_silence_threshold: float = AMPLITUDE_THRESHOLD,
31 min_silence_duration_seconds: float = MIN_SILENCE_DURATION
32) -> List[Tuple[float, float]]:
33 """
34 Detect silences by identifying continuous regions.
35 """
36 is_silent = np.abs(data) < amp_silence_threshold
37 min_silence_samples = int(rate * min_silence_duration_seconds)
39 silences = []
40 start = None
41 for i, silent in enumerate(is_silent):
42 if silent and start is None:
43 start = i
44 elif not silent and start is not None:
45 if i - start >= min_silence_samples:
46 start_seconds = start / rate
47 end_seconds = i / rate
48 silences.append((
49 start_seconds,
50 end_seconds))
51 start = None
52 if start is not None and len(data) - start >= min_silence_samples:
53 start_seconds = start / rate
54 end_seconds = len(data) / rate
55 silences.append((start_seconds, end_seconds))
56 return silences
59def get_time_in_period(
60 start: float,
61 end: float,
62 method: str = "middle" # 'start', 'end', or 'middle'
63) -> float:
64 """Get a time point in the given period based on the method."""
65 if start >= end:
66 raise ValueError("Start time must be less than end time.")
67 if method == "start":
68 return start
69 elif method == "end":
70 return end
71 elif method == "middle":
72 return (start + end) / 2.0
73 return (start + end) / 2.0
76def is_audio_silence(
77 audio_data: np.ndarray,
78 amp_silence_threshold: float = AMPLITUDE_THRESHOLD
79) -> bool:
80 """Check if the audio data is silence based on amplitude threshold."""
81 if not isinstance(audio_data, np.ndarray):
82 raise TypeError(f"Expected np.ndarray for audio_data, got {type(audio_data)}")
83 return bool(np.all(np.abs(audio_data) < amp_silence_threshold))
86def merge_trailing_chunks(
87 chunks: List[Tuple[float, float]],
88 rate: int = 16000,
89 data: Optional[np.ndarray] = None
90) -> List[Tuple[float, float]]:
91 """Merge chunks if it is into the previous chunk."""
92 if len(chunks) < 2:
93 return chunks
94 if data is None:
95 return chunks
96 last_start, last_end = chunks[-1]
97 last_start_idx = int(last_start * rate)
98 last_end_idx = int(last_end * rate)
99 last_chunk = data[last_start_idx:last_end_idx]
100 if is_audio_silence(last_chunk):
101 chunks[-2] = (chunks[-2][0], last_end)
102 chunks.pop()
103 return chunks
106def merge_chunks(
107 chunks: List[Tuple[float, float]],
108 max_duration_seconds: float = 5.0,
109) -> List[Tuple[float, float]]:
110 """Merge chunks if the combined duration is under max_duration_seconds."""
111 if not chunks:
112 return chunks
113 if max_duration_seconds <= 0:
114 return chunks
115 if len(chunks) < 2:
116 return chunks
118 merged_chunks = []
119 current_start, current_end = chunks[0]
120 for start, end in chunks[1:]:
121 if (end - current_start) <= max_duration_seconds:
122 current_end = end # Extend the current chunk
123 else:
124 merged_chunks.append((current_start, current_end))
125 current_start, current_end = start, end
126 merged_chunks.append((current_start, current_end)) # Add the last chunk
127 return merged_chunks
130def align_chunks(
131 chunks: List[Tuple[float, float]],
132 chunk_alignment_seconds: float = 1 / 30.0, # Align to 30 FPS
133) -> List[Tuple[float, float]]:
134 """Align chunk boundaries to the next frame (ceil) based on FPS."""
135 if chunk_alignment_seconds <= 0 or not chunks:
136 return chunks
138 aligned_chunks = []
139 for start, end in chunks:
140 aligned_start = math.ceil(start / chunk_alignment_seconds) * chunk_alignment_seconds
141 aligned_end = math.ceil(end / chunk_alignment_seconds) * chunk_alignment_seconds
142 if aligned_start >= aligned_end:
143 aligned_end = aligned_start + chunk_alignment_seconds
144 aligned_chunks.append((aligned_start, min(aligned_end, chunks[-1][1])))
145 return aligned_chunks
148def get_audio_chunks_by_silences_greedy_new(
149 audio_path: str,
150 max_duration_seconds: float = 5.0, # 5 seconds is what fantasy talking allows
151 chunk_alignment_seconds: float = 1 / 30.0, # Align to 30 FPS
152 min_chunk_duration_seconds: float = 0.5,
153 method: str = "start", # 'start', 'end', or 'middle'
154) -> List[Tuple[float, float]]:
155 """
156 Chunk audio using hierarchical greedy silence-based splitting:
157 - Prefer longer silences first
158 - Chunks <= max_duration_seconds
159 - Merge tiny intermediate and trailing chunks
160 - Boundaries aligned to chunk_alignment_seconds
161 """
162 if not isinstance(audio_path, str):
163 raise TypeError(f"Expected str for audio_path, got {type(audio_path)}")
164 if not os.path.exists(audio_path):
165 raise FileNotFoundError(f"WAV file does not exist: {audio_path}")
167 rate, data = wavfile.read(audio_path)
168 if data.ndim > 1:
169 data = data.mean(axis=1).astype(data.dtype)
171 total_duration = len(data) / rate
172 chunk_start = 0.0
173 chunks: List[Tuple[float, float]] = []
175 sorted_thresholds_ms = sorted(SILENCE_THRESHOLDS_MS, reverse=True)
177 # Precompute silences for all thresholds
178 silence_map = {
179 threshold_ms: detect_silences(
180 data,
181 rate,
182 min_silence_duration_seconds=threshold_ms / 1000.0)
183 for threshold_ms in sorted_thresholds_ms
184 }
186 while chunk_start < total_duration:
187 chunk_deadline = min(chunk_start + max_duration_seconds, total_duration)
188 chunk_end = None
190 # Try longer silences first
191 for threshold_ms in sorted_thresholds_ms:
192 candidate_silences = [
193 (s, e) for s, e in silence_map[threshold_ms]
194 if chunk_start < s <= chunk_deadline
195 ]
196 if candidate_silences:
197 # Pick the last silence in period
198 last_silence = max(candidate_silences, key=lambda se: get_time_in_period(se[0], se[1], method))
199 chunk_end = get_time_in_period(last_silence[0], last_silence[1], method)
200 break
202 # No silence found → use max_duration
203 if chunk_end is None:
204 chunk_end = chunk_deadline
206 # Merge tiny chunks (intermediate or trailing)
207 if chunks and (chunk_end - chunk_start) < min_chunk_duration_seconds:
208 prev_start, prev_end = chunks.pop()
209 chunk_start = prev_start
210 chunk_end = max(prev_end, chunk_end)
212 # Align boundaries
213 aligned_start = math.floor(chunk_start / chunk_alignment_seconds) * chunk_alignment_seconds
214 aligned_end = math.ceil(chunk_end / chunk_alignment_seconds) * chunk_alignment_seconds
215 aligned_end = min(aligned_end, total_duration)
217 chunks.append((aligned_start, aligned_end))
218 chunk_start = chunk_end
220 return chunks
223def get_audio_chunks_by_silences_greedy(
224 audio_path: str,
225 max_duration_seconds: float = 5.0, # 5 seconds is what fantasy talking allows
226 chunk_alignment_seconds: float = 1 / 30.0, # Align to 30 FPS
227 min_chunk_duration_seconds: float = 0.5,
228 method: str = "start", # 'start', 'end', or 'middle'
229) -> List[Tuple[float, float]]:
230 """
231 Chunk audio using greedy silence-based splitting with contiguous segments.
232 Each chunk ends at the last silence under max_duration.
233 Prevents creating multiple small chunks unnecessarily.
234 """
235 if not isinstance(audio_path, str):
236 raise TypeError(f"Expected str for audio_path, got {type(audio_path)}")
237 if not os.path.exists(audio_path):
238 raise FileNotFoundError(f"WAV file does not exist: {audio_path}")
239 rate, data = wavfile.read(audio_path)
241 if data.ndim > 1: # Convert stereo to mono if needed
242 data = data.mean(axis=1).astype(data.dtype)
244 total_duration = len(data) / rate
245 chunk_start_seconds = 0.0
246 chunks = []
248 # Precompute silences at all thresholds
249 silence_map = {
250 duration_ms: detect_silences(
251 data, rate,
252 min_silence_duration_seconds=duration_ms / 1000.0)
253 for duration_ms in SILENCE_THRESHOLDS_MS
254 }
256 # Find within the silences
257 while chunk_start_seconds < total_duration:
258 chunk_deadline = min(chunk_start_seconds + max_duration_seconds, total_duration)
260 # Find the last silence before the chunk deadline
261 best_silence_end = None
262 for duration_ms in SILENCE_THRESHOLDS_MS:
263 silences = silence_map[duration_ms]
264 candidate_silences = [
265 (start, end) for start, end in silences
266 if chunk_start_seconds < start <= chunk_deadline
267 ]
268 if candidate_silences:
269 # Take the last silence under deadline
270 candidate_time = max(
271 get_time_in_period(start, end, method)
272 for start, end in candidate_silences
273 )
274 if candidate_time > (chunk_deadline - 0.5): # prefer silence close to end
275 best_silence_end = candidate_time
276 break
277 elif not best_silence_end:
278 # fallback: keep the last available silence if no better found
279 best_silence_end = candidate_time
281 # If no silence or too early, just go to max duration
282 if best_silence_end is None or best_silence_end < chunk_start_seconds + 0.5:
283 chunk_end_seconds = chunk_deadline
284 else:
285 chunk_end_seconds = best_silence_end
287 chunks.append((chunk_start_seconds, chunk_end_seconds))
288 chunk_start_seconds = chunk_end_seconds
290 merged_chunks = merge_trailing_chunks(chunks, rate, data)
291 merged_chunks = merge_chunks(merged_chunks, max_duration_seconds)
292 aligned_chunks = align_chunks(merged_chunks, chunk_alignment_seconds)
294 # Final checks
295 # If not enough chunks
296 min_num_chunks = math.ceil(total_duration / max_duration_seconds)
297 if len(aligned_chunks) < min_num_chunks:
298 return get_audio_chunks_hard_cutoff(
299 total_duration,
300 max_duration_seconds)
301 # If any chunk is longer than max_duration_seconds
302 if any((end - start) > max_duration_seconds for start, end in aligned_chunks):
303 return get_audio_chunks_hard_cutoff(
304 total_duration,
305 max_duration_seconds)
307 return aligned_chunks
310def get_audio_chunks_hard_cutoff(
311 total_duration: float,
312 max_duration_seconds: float = 5.0,
313) -> List[Tuple[float, float]]:
314 """Chunk audio using hard cutoff at max_duration_seconds."""
315 chunks = []
316 num_chunks = math.ceil(total_duration / max_duration_seconds)
317 chunk_duration = total_duration / num_chunks
318 for i in range(num_chunks):
319 start = i * chunk_duration
320 end = min((i + 1) * chunk_duration, total_duration)
321 chunks.append((start, end))
322 return chunks
325def get_audio_chunks_by_silences_binary(
326 audio_path: str,
327 max_duration_seconds: float = 5.0,
328 chunk_alignment_seconds: float = 1 / 30.0, # Align to 30 FPS
329 method: str = "middle", # 'start', 'end', or 'middle'
330) -> List[Tuple[float, float]]:
331 """
332 Chunk audio using recursive binary search.
333 If a segment is longer than max_duration_seconds, split it at the longest
334 silence inside (according to `method`). Falls back to hard cutoff.
335 Returns a list of (start_time, end_time) in seconds.
336 """
337 if not isinstance(audio_path, str):
338 raise TypeError(f"Expected str for audio_path, got {type(audio_path)}")
340 rate, data = wavfile.read(audio_path)
342 # Convert stereo to mono if needed
343 if data.ndim > 1:
344 data = data.mean(axis=1).astype(data.dtype)
346 total_duration = len(data) / rate
348 # Detect all silences once and sort by length (longest first)
349 silences = detect_silences(data, rate)
350 silences = sorted(silences, key=lambda s: (s[1] - s[0]), reverse=True)
352 def split_segment(start: float, end: float) -> List[Tuple[float, float]]:
353 duration = end - start
354 if duration <= max_duration_seconds:
355 return [(start, end)]
357 # Find longest silence inside this segment
358 candidate = None
359 for s_start, s_end in silences:
360 if start < s_start and s_end < end:
361 candidate = (s_start, s_end)
362 break # longest first due to sorting
364 if candidate:
365 split_point = get_time_in_period(candidate[0], candidate[1], method)
366 return split_segment(start, split_point) + split_segment(split_point, end)
368 # No silence inside -> hard cutoff
369 cutoff = min(start + max_duration_seconds, end)
370 left = [(start, cutoff)]
371 right = split_segment(cutoff, end) if cutoff < end else []
372 return left + right
374 chunks = split_segment(0.0, total_duration)
376 merged_chunks = merge_trailing_chunks(chunks, rate, data)
377 merged_chunks = merge_chunks(merged_chunks, max_duration_seconds)
378 aligned_chunks = align_chunks(merged_chunks, chunk_alignment_seconds)
380 return aligned_chunks
383def get_audio_chunks_by_silences(
384 audio_path: str,
385 max_duration_seconds: float = 5.0, # 5 seconds is what fantasy talking allows
386 chunk_alignment_seconds: float = 1 / 30.0, # Align to 30 FPS
387 min_chunk_duration_seconds: float = 0.5,
388 method: str = "start", # 'start', 'end', or 'middle'
389) -> List[Tuple[float, float]]:
390 # get_audio_chunks_by_silences_binary
391 # get_audio_chunks_by_silences_greedy
392 return get_audio_chunks_by_silences_greedy(
393 audio_path,
394 max_duration_seconds,
395 chunk_alignment_seconds,
396 min_chunk_duration_seconds,
397 method)
400def is_audio_path_silence(audio_path: str) -> bool:
401 """
402 Check if the audio is silence based on amplitude threshold.
403 Returns True if the audio is considered silence, False otherwise.
404 """
405 if not isinstance(audio_path, str):
406 raise TypeError(f"Expected str for audio_path, got {type(audio_path)}")
407 if not os.path.exists(audio_path):
408 raise FileNotFoundError(f"WAV file does not exist: {audio_path}")
409 _, data = wavfile.read(audio_path)
410 if data.ndim > 1:
411 data = data.mean(axis=1).astype(data.dtype)
412 return is_audio_silence(data)
415def is_audio_base64_silence(audio_base64: str) -> bool:
416 """
417 Check if the base64-encoded audio is silence based on amplitude threshold.
418 """
419 if not isinstance(audio_base64, str):
420 raise TypeError(f"Expected str for audio_base64, got {type(audio_base64)}")
421 audio_bytes = base64.b64decode(audio_base64)
422 buffer = io.BytesIO(audio_bytes)
423 _, data = wavfile.read(buffer)
424 if data.ndim > 1:
425 data = data.mean(axis=1).astype(data.dtype)
426 return is_audio_silence(data)
429def strip_audio_file_silence(
430 input_path: str,
431 strip_start: bool = False,
432 strip_end: bool = True,
433 output_path: Optional[str] = None,
434 amp_silence_threshold: float = AMPLITUDE_THRESHOLD
435) -> str:
436 """
437 Strip silence from the start and/or end of a WAV audio file.
438 """
439 if not isinstance(input_path, str):
440 raise TypeError(f"Expected str for input_path, got {type(input_path)}")
441 if not os.path.exists(input_path):
442 raise FileNotFoundError(f"Input WAV file does not exist: {input_path}")
443 if not input_path.lower().endswith(".wav"):
444 raise ValueError(f"Input file must be a WAV file: {input_path}")
446 with wave.open(input_path, "rb") as wf:
447 params = wf.getparams()
448 n_channels = wf.getnchannels()
449 sampwidth = wf.getsampwidth()
450 nframes = wf.getnframes()
451 frames = wf.readframes(nframes)
453 dtype = {1: np.int8, 2: np.int16, 4: np.int32}[sampwidth]
454 audio = np.frombuffer(frames, dtype=dtype)
456 # Handle multi-channel audio: reshape so axis=1 is channels
457 if n_channels > 1:
458 audio = audio.reshape(-1, n_channels)
459 amplitude = np.max(np.abs(audio), axis=1) # collapse channels
460 else:
461 amplitude = np.abs(audio)
463 # Identify non-silence indices
464 non_silent_indices = np.where(amplitude >= amp_silence_threshold)[0]
466 if non_silent_indices.size == 0:
467 # Entire file is silent: return original or empty file
468 if not output_path:
469 output_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
470 with wave.open(output_path, "wb") as wf:
471 wf.setparams(params._replace(nframes=0))
472 wf.writeframes(b"")
473 return output_path
475 start_idx, end_idx = 0, len(audio)
477 if strip_start:
478 start_idx = non_silent_indices[0]
479 if strip_end:
480 end_idx = non_silent_indices[-1] + 1
482 trimmed = audio[start_idx:end_idx]
484 if not output_path:
485 output_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
487 with wave.open(output_path, "wb") as wf:
488 wf.setparams(params._replace(nframes=len(trimmed)))
489 wf.writeframes(trimmed.tobytes())
491 return output_path
494def get_sentences(text: str) -> list[str]:
495 """
496 Extract sentences from a text including the punctuation.
497 """
498 if not text:
499 return []
500 # Simple regex to split sentences while keeping the punctuation
501 sentences = re.split(r'([.!?]+)', text)
502 sentences = [s.strip() for s in sentences if s.strip()] # Remove empty strings
503 # Combine punctuation with the sentence
504 return [''.join(sentences[i:i + 2]).strip() for i in range(0, len(sentences), 2)]
507def estimate_num_words_from_audio_duration(duration_seconds: float, speed: float = 1.0) -> int:
508 """
509 Estimate number of words based on audio duration.
510 0.4 seconds per word on average (16 * num_words / 40.0).
511 Data in "audio_duration.csv".
512 """
513 num_words = int(math.ceil((duration_seconds * speed) / 0.4)) # seconds -> words
514 return max(1, num_words)
517def estimate_audio_duration_from_words(num_words: int, speed: float = 1.0) -> float:
518 """
519 Estimate audio duration based on the number of words.
520 0.4 seconds per word on average (16 * num_words / 40.0).
521 Data in "audio_duration.csv".
522 """
523 duration_seconds = 0.4 * num_words # words -> seconds
524 return duration_seconds / speed
527def estimate_audio_duration_from_chars(num_chars: int, speed: float = 1.0) -> float:
528 """
529 Estimate audio duration based on the number of characters.
530 0.064 seconds per char on average (13 * num_chars / 200.0).
531 Data in "audio_duration.csv".
532 """
533 duration_seconds = 0.065 * num_chars # chars -> seconds
534 return duration_seconds / speed
537def estimate_audio_duration(text: str, speed: float = 1.0) -> float:
538 """
539 Estimate audio duration based on the number of words and characters.
540 """
541 num_chars = len(text)
542 num_words = len(text.strip().split())
543 # sub_sentences = get_sentences(text)
544 # num_sentences = len(sub_sentences)
546 duration_seconds = max(
547 estimate_audio_duration_from_chars(num_chars, speed),
548 estimate_audio_duration_from_words(num_words, speed),
549 ) # seconds
550 return duration_seconds
553def split_into_sentences_max_duration(
554 text: str,
555 max_duration: float = 5.0,
556) -> list[str]:
557 """
558 Split the text into sub-sentences with a maximum estimated audio duration.
559 If a sentence exceeds the limit, it is further split by words.
560 """
561 if not text:
562 return []
564 if estimate_audio_duration(text) <= max_duration:
565 return [text]
567 sentences = get_sentences(text)
568 if not sentences:
569 return [text]
571 chunks = []
572 current_chunk = ""
573 current_duration = 0.0 # seconds
575 for sentence in sentences:
576 sentence_duration = estimate_audio_duration(sentence)
578 if sentence_duration <= max_duration:
579 # Sentence fits within max duration
580 if current_duration + sentence_duration > max_duration:
581 # Commit current chunk and start a new one
582 current_chunk = current_chunk.strip()
583 if current_chunk:
584 chunks.append(current_chunk)
585 current_chunk = sentence
586 current_duration = sentence_duration
587 else:
588 # Append sentence to current chunk
589 if current_chunk:
590 current_chunk += " " + sentence
591 else:
592 current_chunk = sentence
593 current_duration += sentence_duration
595 else:
596 # Sentence too long -> split by words
597 words = sentence.split()
598 word_chunk = ""
599 word_duration = 0.0
601 for word in words:
602 duration = estimate_audio_duration(word + " ")
603 if word_duration + duration > max_duration:
604 # Commit the chunk of words
605 word_chunk = word_chunk.strip()
606 if word_chunk:
607 chunks.append(word_chunk)
608 word_chunk = word
609 word_duration = duration
610 else:
611 word_chunk += " " + word if word_chunk else word
612 word_duration += duration
614 word_chunk = word_chunk.strip()
615 if word_chunk:
616 chunks.append(word_chunk)
618 # Reset main chunk tracking after forced split
619 current_chunk = ""
620 current_duration = 0.0
622 current_chunk = current_chunk.strip()
623 if current_chunk:
624 chunks.append(current_chunk)
626 return chunks
629def generate_waveform_plt(wav_file_name: str) -> str:
630 """Generate a waveform PNG image from a WAV audio file using matplotlib."""
631 import matplotlib.pyplot as plt
633 # Open the file
634 # Read the WAV file to generate the waveform
635 rate, data = wavfile.read(wav_file_name)
636 if data.ndim > 1: # Stereo or multi-channel audio to mono
637 data = data.mean(axis=1).astype(data.dtype)
638 data_plt = data / np.max(np.abs(data)) if np.max(np.abs(data)) != 0 else data
640 # Generate and add silences
641 silences = {}
642 for silence_duration_ms in [200, 100, 50, 10]:
643 silences[silence_duration_ms] = detect_silences(
644 data, rate,
645 min_silence_duration_seconds=silence_duration_ms / 1000.0)
647 waveform_path = f"{wav_file_name}-waveform.png"
648 times = np.arange(len(data)) / rate
649 plt.figure(figsize=(14, 4))
650 plt.plot(times, data_plt, label="Waveform", color="steelblue")
652 color_map = {
653 200: "red",
654 100: "orange",
655 50: "yellow",
656 10: "green"
657 }
658 for silence_duration, start_end in reversed(silences.items()):
659 for silence_start, silence_end in start_end:
660 label = None
661 if silence_start == start_end[0][0]:
662 label = f"Silence {silence_duration} ms"
663 plt.axvspan(
664 silence_start, silence_end,
665 color=color_map[silence_duration],
666 alpha=0.5,
667 label=label)
669 duration_seconds = len(data) / rate
670 plt.xlim(0, duration_seconds)
671 plt.ylim(-1, 1)
672 plt.title("Waveform with silences highlighted")
673 plt.xlabel("Time (seconds)")
674 plt.ylabel("Amplitude (normalized)")
675 plt.legend(loc="lower left")
676 plt.savefig(waveform_path, bbox_inches="tight", dpi=300)
677 plt.close()
679 return waveform_path