hls_download.py

import subprocess
import os

def download_clips(stream_url, out_dir, start_time, end_time, resize=True, use_60fps=False, use_cuda=True):
    # remove all .mp4 files in out_dir to avoid confusion
    if len(os.listdir(out_dir)) > 5:
        for f in os.listdir(out_dir):
            if f.endswith('.mp4'):
                os.remove(os.path.join(out_dir, f))
    output_file = os.path.join(out_dir, f"train_{len(os.listdir(out_dir))}.mp4")
    if resize:  # resize and convert to 30 fps
        ffmpeg_cmd = [
                'ffmpeg',
                '-hwaccel', 'cuda' if use_cuda else 'none',
                '-i', stream_url,
                #'-c:v', 'libx264',
                '-crf', '23',
                '-r', '30' if not use_60fps else '60',
                '-maxrate', '2M',
                '-bufsize', '4M',
                '-vf', f"scale=-2:300",
                '-ss', start_time] + (['-to', end_time] if end_time != '' else []) + [
                '-c:a', 'aac',
                '-b:a', '128k',
                '-y',
                output_file
            ]
        print(' '.join(ffmpeg_cmd))
        try:
            subprocess.run(ffmpeg_cmd, check=True, capture_output=True, text=True)
        except subprocess.CalledProcessError as e:
            print(f"Error occurred: {e}")
            print(f"ffmpeg output: {e.output}")
            return output_file
    # else:
    #     os.rename(tmp_file, output_file)
    print(f"Converted {output_file}")

    return output_file