Skip to content

Commit

Permalink
fix(tools): load audio close file after reading
Browse files Browse the repository at this point in the history
  • Loading branch information
fumiama committed Nov 28, 2024
1 parent c3948c8 commit 00c56ee
Showing 1 changed file with 33 additions and 11 deletions.
44 changes: 33 additions & 11 deletions tools/audio/av.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,36 +42,53 @@ def wav2(i: BytesIO, o: BufferedWriter, format: str):

def load_audio(
file: Union[str, BytesIO, Path],
sr: Optional[int]=None,
format: Optional[str]=None,
mono=True
sr: Optional[int] = None,
format: Optional[str] = None,
mono=True,
) -> Union[np.ndarray, Tuple[np.ndarray, int]]:
"""
https://github.com/fumiama/Retrieval-based-Voice-Conversion-WebUI/blob/412a9950a1e371a018c381d1bfb8579c4b0de329/infer/lib/audio.py#L39
"""
if (isinstance(file, str) and not Path(file).exists()) or (isinstance(file, Path) and not file.exists()):
if (isinstance(file, str) and not Path(file).exists()) or (
isinstance(file, Path) and not file.exists()
):
raise FileNotFoundError(f"File not found: {file}")
rate = 0

container = av.open(file, format=format)
audio_stream = next(s for s in container.streams if s.type == "audio")
channels = 1 if audio_stream.layout == "mono" else 2
container.seek(0)
resampler = AudioResampler(format="fltp", layout=audio_stream.layout, rate=sr) if sr is not None else None
resampler = (
AudioResampler(format="fltp", layout=audio_stream.layout, rate=sr)
if sr is not None
else None
)

# Estimated maximum total number of samples to pre-allocate the array
# AV stores length in microseconds by default
estimated_total_samples = int(container.duration * sr // 1_000_000) if sr is not None else 48000
decoded_audio = np.zeros(estimated_total_samples + 1 if channels == 1 else (channels, estimated_total_samples + 1), dtype=np.float32)
estimated_total_samples = (
int(container.duration * sr // 1_000_000) if sr is not None else 48000
)
decoded_audio = np.zeros(
(
estimated_total_samples + 1
if channels == 1
else (channels, estimated_total_samples + 1)
),
dtype=np.float32,
)

offset = 0

def process_packet(packet: List[AudioFrame]):
frames_data = []
rate = 0
for frame in packet:
frame.pts = None # 清除时间戳,避免重新采样问题
resampled_frames = resampler.resample(frame) if resampler is not None else [frame]
# frame.pts = None # 清除时间戳,避免重新采样问题
resampled_frames = (
resampler.resample(frame) if resampler is not None else [frame]
)
for resampled_frame in resampled_frames:
frame_data = resampled_frame.to_ndarray()
rate = resampled_frame.rate
Expand All @@ -83,16 +100,21 @@ def frame_iter(container):
yield p.decode()

for r, frames_data in map(process_packet, frame_iter(container)):
if not rate: rate = r
if not rate:
rate = r
for frame_data in frames_data:
end_index = offset + len(frame_data[0])

# 检查 decoded_audio 是否有足够的空间,并在必要时调整大小
if end_index > decoded_audio.shape[1]:
decoded_audio = np.resize(decoded_audio, (decoded_audio.shape[0], end_index*4))
decoded_audio = np.resize(
decoded_audio, (decoded_audio.shape[0], end_index * 4)
)

np.copyto(decoded_audio[..., offset:end_index], frame_data)
offset += len(frame_data[0])

container.close()

# Truncate the array to the actual size
decoded_audio = decoded_audio[..., :offset]
Expand Down

0 comments on commit 00c56ee

Please sign in to comment.