test_add_marker.py

import logging
import time
from pathlib import Path
from typing import List

from riff_wave import AdtlLabel
from riff_wave.core import WaveContainer, riff_wave_logger, CueTiming, Cue, AdtlAnnotation

riff_wave_logger.setLevel(logging.DEBUG)

test_audio_path = 'samples/sil_2s_marker.wav'
save_dir = 'samples'
test_audio_stem = Path(test_audio_path).stem

# 1. read and parse
time_start = time.perf_counter()
with open(test_audio_path, 'rb') as f:
    wave = WaveContainer.from_stream(f)
time_end = time.perf_counter()
print(f'[parse] time cost: {time_end - time_start:.3f}s')
print(f'{wave.fmt = }')
print(f'{wave.duration = }')

# 2. get original adtl annotations
adtl_annotations = wave.get_adtl_annotations()
print(f'{adtl_annotations = }')

# 3. create new adtl annotations and replace the original ones
mark_times = [0.5, 1.0, 1.5]
mark_labels = ['A', 'B', 'C']

adtl_list: List[AdtlLabel] = []
cue_timings: List[CueTiming] = []
adtl_annotations: List[AdtlAnnotation] = []
for i, (marker_time, label) in enumerate(zip(mark_times, mark_labels)):
    cur_annotation = AdtlAnnotation(
        id=i+1,
        start_idx=int(marker_time * wave.sample_rate),
        end_idx=int(marker_time * wave.sample_rate),
        label=label,
        sample_rate=wave.sample_rate
    )
    adtl_annotations.append(cur_annotation)

flag_has_adtl = any(isinstance(x, Cue) for x in wave.sub_chunks)
if flag_has_adtl:
    riff_wave_logger.warning('Replacing existing ADTL annotations. Original ADTL annotations will be removed.')

wave.update_adtl_annotations(adtl_annotations=adtl_annotations)

# 4. rewrite
time_start = time.perf_counter()
with open(f'{save_dir}/{test_audio_stem}_new_marker.wav', 'wb') as f:
    f.write(wave.dump())
time_end = time.perf_counter()
print(f'[dump] time cost: {time_end - time_start:.3f}s')