flaxdiff 0.1.38.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

flaxdiff/data/sources/audio_utils.py

@@ -0,0 +1,142 @@
+"""
+Audio utilities for efficiently loading audio data from video files.
+This module provides alternatives to decord's AudioReader/AVReader (which have memory leaks).
+"""
+
+import os
+import tempfile
+import subprocess
+import numpy as np
+from typing import Tuple, Optional, Union
+
+
+def read_audio_ffmpeg(
+    video_path: str, 
+    start_time: Optional[float] = None,
+    duration: Optional[float] = None,
+    target_sr: int = 16000
+) -> Tuple[np.ndarray, int]:
+    """
+    Extract audio from video file using ffmpeg subprocess calls.
+    
+    Args:
+        video_path: Path to the video file.
+        start_time: Start time in seconds (optional).
+        duration: Duration to extract in seconds (optional).
+        target_sr: Target sample rate for the audio.
+        
+    Returns:
+        Tuple of (audio_data, sample_rate) where audio_data is a numpy array.
+    """
+    # Create a temporary file for the audio
+    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp_file:
+        tmp_path = tmp_file.name
+    
+    try:
+        # Build the ffmpeg command
+        cmd = ['ffmpeg', '-y', '-i', video_path]
+        
+        # Add time parameters if specified
+        if start_time is not None:
+            cmd.extend(['-ss', str(start_time)])
+            
+        if duration is not None:
+            cmd.extend(['-t', str(duration)])
+            
+        # Set output parameters (mono, target sample rate)
+        cmd.extend([
+            '-ac', '1',  # mono
+            '-ar', str(target_sr),  # sample rate
+            '-vn',  # no video
+            '-f', 'wav',  # wav format
+            tmp_path
+        ])
+        
+        # Execute the command
+        subprocess.run(cmd, check=True, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
+        
+        # Read the audio file using numpy
+        audio_data = np.fromfile(tmp_path, np.int16).astype(np.float32) / 32768.0  # Convert to float in [-1, 1]
+        
+        return audio_data, target_sr
+        
+    finally:
+        # Always clean up the temporary file
+        try:
+            os.unlink(tmp_path)
+        except:
+            pass
+
+
+def read_audio_moviepy(
+    video_path: str,
+    start_time: Optional[float] = None,
+    duration: Optional[float] = None,
+    target_sr: int = 16000
+) -> Tuple[np.ndarray, int]:
+    """
+    Extract audio from video file using moviepy.
+    Requires the moviepy package: pip install moviepy
+    
+    Args:
+        video_path: Path to the video file.
+        start_time: Start time in seconds (optional).
+        duration: Duration to extract in seconds (optional).
+        target_sr: Target sample rate for the audio.
+        
+    Returns:
+        Tuple of (audio_data, sample_rate) where audio_data is a numpy array.
+    """
+    try:
+        from moviepy import VideoFileClip
+    except ImportError:
+        raise ImportError("moviepy is not installed. Install it with 'pip install moviepy'")
+    
+    # Load video file
+    if start_time is not None or duration is not None:
+        start_t = start_time if start_time is not None else 0
+        end_t = start_t + duration if duration is not None else None
+        video = VideoFileClip(video_path).subclipped(start_t, end_t)
+    else:
+        video = VideoFileClip(video_path)
+    # Extract audio
+    audio = video.audio.with_fps(target_sr)
+    
+    # Get audio data
+    audio_data = audio.to_soundarray()
+    
+    # Convert to mono if stereo
+    if audio_data.ndim > 1 and audio_data.shape[1] > 1:
+        audio_data = np.mean(audio_data, axis=1)
+    
+    # Clean up
+    video.close()
+    
+    return audio_data, target_sr
+
+
+# Helper function to choose the best available method
+def read_audio(
+    video_path: str,
+    start_time: Optional[float] = None,
+    duration: Optional[float] = None,
+    target_sr: int = 16000,
+    method: str = 'ffmpeg'
+) -> Tuple[np.ndarray, int]:
+    """
+    Extract audio from video file using the specified method.
+    
+    Args:
+        video_path: Path to the video file.
+        start_time: Start time in seconds (optional).
+        duration: Duration to extract in seconds (optional).
+        target_sr: Target sample rate for the audio.
+        method: Method to use ('ffmpeg' or 'moviepy').
+        
+    Returns:
+        Tuple of (audio_data, sample_rate) where audio_data is a numpy array.
+    """
+    if method == 'moviepy':
+        return read_audio_moviepy(video_path, start_time, duration, target_sr)
+    else:  # default to ffmpeg
+        return read_audio_ffmpeg(video_path, start_time, duration, target_sr)

flaxdiff/data/sources/av_example.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+"""
+Example script demonstrating how to use the memory-leak-free audio-video reading functions.
+"""
+
+import os
+import time
+import numpy as np
+import matplotlib.pyplot as plt
+from av_utils import read_av_improved, read_av_batch
+from audio_utils import read_audio
+import argparse
+
+
+def visualize_av_data(audio_data, video_frames, output_path=None):
+    """
+    Visualize audio and video data.
+    
+    Args:
+        audio_data: Audio data as numpy array or list.
+        video_frames: Video frames as numpy array.
+        output_path: Path to save visualization (optional).
+    """
+    fig = plt.figure(figsize=(12, 6))
+    
+    # Number of frames to show
+    num_frames = min(4, len(video_frames))
+    
+    # Plot audio waveform
+    plt.subplot(2, num_frames, 1)
+    plt.plot(audio_data[:10000])
+    plt.title('Audio Waveform')
+    plt.grid(True)
+    
+    # Plot audio spectrogram
+    plt.subplot(2, num_frames, 2)
+    plt.specgram(audio_data, NFFT=1024, Fs=16000)
+    plt.title('Audio Spectrogram')
+    
+    # Plot sample frames
+    for i in range(num_frames):
+        plt.subplot(2, num_frames, num_frames+i+1)
+        plt.imshow(video_frames[i*len(video_frames)//num_frames])
+        plt.title(f'Frame {i*len(video_frames)//num_frames}')
+        plt.axis('off')
+    
+    plt.tight_layout()
+    
+    if output_path:
+        plt.savefig(output_path)
+        print(f"Visualization saved to {output_path}")
+    
+    plt.show()
+
+
+def benchmark_av_reading(video_path, num_iterations=10, use_batch=False):
+    """
+    Benchmark audio-video reading performance.
+    
+    Args:
+        video_path: Path to the video file.
+        num_iterations: Number of iterations for benchmarking.
+        use_batch: Whether to use batch reading.
+    """
+    print(f"Benchmarking {'batch' if use_batch else 'single'} reading...")
+    
+    # Perform warmup
+    if use_batch:
+        _ = read_av_batch([video_path])
+    else:
+        _ = read_av_improved(video_path)
+    
+    # Measure performance
+    start_time = time.time()
+    
+    for i in range(num_iterations):
+        if use_batch:
+            results = read_av_batch([video_path])
+        else:
+            audio, video = read_av_improved(video_path)
+    
+    end_time = time.time()
+    avg_time = (end_time - start_time) / num_iterations
+    
+    print(f"Average time per read: {avg_time:.4f} seconds")
+    
+    return avg_time
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Demo for memory-leak-free audio-video reading")
+    parser.add_argument("--video", "-v", required=True, help="Path to the video file")
+    parser.add_argument("--output", "-o", help="Path to save visualization")
+    parser.add_argument("--benchmark", "-b", action="store_true", help="Run benchmarks")
+    parser.add_argument("--iterations", "-i", type=int, default=10, help="Number of benchmark iterations")
+    
+    args = parser.parse_args()
+    
+    if not os.path.exists(args.video):
+        print(f"Error: Video file not found: {args.video}")
+        return
+    
+    # Load audio-video data
+    print(f"Reading audio-video data from {args.video}...")
+    audio, video = read_av_improved(args.video)
+    
+    print(f"Video shape: {video.shape}")
+    print(f"Audio length: {len(audio)}")
+    
+    # Visualize data
+    visualize_av_data(audio, video, args.output)
+    
+    # Run benchmarks if requested
+    if args.benchmark:
+        print("\nRunning benchmarks...")
+        single_time = benchmark_av_reading(args.video, args.iterations, use_batch=False)
+        batch_time = benchmark_av_reading(args.video, args.iterations, use_batch=True)
+        
+        print("\nBenchmark results:")
+        print(f"Single reading: {single_time:.4f} seconds per video")
+        print(f"Batch reading: {batch_time:.4f} seconds per video")
+    
+
+if __name__ == "__main__":
+    main()