Skip to content

Commit 02fbaaf

Browse files
satraclaude
andauthored
Fix video I/O: context manager, no-audio guard, temp cleanup (#483)
* Fix video I/O: context manager, no-audio guard, temp cleanup Addresses PR #482 review comments: - Use `with av.open()` context manager (proper resource cleanup) - Guard against videos with no audio track (was IndexError) - Clean up temp directory with shutil.rmtree in finally block - Skip videos without audio instead of crashing Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Use TemporaryDirectory context manager instead of mkdtemp Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Fix fragile path calculation and empty audio list crash - Use os.path.relpath instead of str.replace for safe relative paths (leading slash caused writes to root directory) - Return empty dict when no audio tracks found instead of crashing read_files_from_disk with empty list Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent d503a06 commit 02fbaaf

1 file changed

Lines changed: 33 additions & 25 deletions

File tree

src/senselab/video/tasks/input_output.py

Lines changed: 33 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
import numpy as np
1010
import soundfile as sf
1111

12-
from senselab.utils.data_structures import from_strings_to_files, get_common_directory
12+
from senselab.utils.data_structures import from_strings_to_files, get_common_directory, logger
1313
from senselab.utils.tasks.input_output import read_files_from_disk
1414

1515

@@ -33,22 +33,29 @@ def extract_audios_from_local_videos(
3333
A dataset dict of the extracted audio files.
3434
"""
3535

36-
def _extract_audio_from_local_video(video_path: Path, output_audio_path: str, fmt: str, codec: str) -> None:
37-
"""Extract audio from a video file using PyAV."""
38-
container = av.open(str(video_path))
39-
audio_stream = container.streams.audio[0]
40-
sample_rate = audio_stream.rate or 16000
36+
def _extract_audio_from_local_video(video_path: Path, output_audio_path: str, fmt: str, codec: str) -> bool:
37+
"""Extract audio from a video file using PyAV.
4138
42-
frames = []
43-
for frame in container.decode(audio=0):
44-
arr = frame.to_ndarray()
45-
if arr.ndim > 1:
46-
arr = arr.mean(axis=0) # downmix to mono
47-
frames.append(arr)
48-
container.close()
39+
Returns:
40+
True if audio was extracted, False if the video has no audio track.
41+
"""
42+
with av.open(str(video_path)) as container:
43+
if not container.streams.audio:
44+
logger.warning("No audio track found in %s, skipping.", video_path)
45+
return False
46+
47+
audio_stream = container.streams.audio[0]
48+
sample_rate = audio_stream.rate or 16000
49+
50+
frames = []
51+
for frame in container.decode(audio=0):
52+
arr = frame.to_ndarray()
53+
if arr.ndim > 1:
54+
arr = arr.mean(axis=0) # downmix to mono
55+
frames.append(arr)
4956

5057
if not frames:
51-
return
58+
return False
5259

5360
audio_data = np.concatenate(frames)
5461

@@ -61,20 +68,21 @@ def _extract_audio_from_local_video(video_path: Path, output_audio_path: str, fm
6168
)
6269

6370
sf.write(output_audio_path, audio_data, sample_rate, format=fmt.upper())
71+
return True
6472

6573
if isinstance(files, str):
6674
files = [files]
6775
formatted_files = from_strings_to_files(files)
6876
common_path = get_common_directory(files)
6977

70-
temp_dir = tempfile.mkdtemp()
71-
72-
audio_files_paths = []
73-
for file in formatted_files:
74-
base_file_name = os.path.splitext(str(file.filepath).replace(common_path, ""))[0]
75-
output_audio_path = os.path.join(temp_dir, f"{base_file_name}.{audio_format}")
76-
os.makedirs(os.path.dirname(output_audio_path), exist_ok=True)
77-
_extract_audio_from_local_video(file.filepath, output_audio_path, fmt=audio_format, codec=acodec)
78-
audio_files_paths.append(output_audio_path)
79-
80-
return read_files_from_disk(audio_files_paths)
78+
with tempfile.TemporaryDirectory(prefix="senselab-video-io-") as temp_dir:
79+
audio_files_paths = []
80+
for file in formatted_files:
81+
rel_path = os.path.relpath(file.filepath, common_path)
82+
base_file_name = os.path.splitext(rel_path)[0]
83+
output_audio_path = os.path.join(temp_dir, f"{base_file_name}.{audio_format}")
84+
os.makedirs(os.path.dirname(output_audio_path), exist_ok=True)
85+
if _extract_audio_from_local_video(file.filepath, output_audio_path, fmt=audio_format, codec=acodec):
86+
audio_files_paths.append(output_audio_path)
87+
88+
return read_files_from_disk(audio_files_paths) if audio_files_paths else {}

0 commit comments

Comments
 (0)