Update hls_download.py

hls_download.py

@@ -1,39 +1,40 @@
-import subprocess
-import os
-
-def download_clips(stream_url, out_dir, start_time, end_time, resize=True, use_60fps=False):
-    # 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',
-                '-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}")
-
+import subprocess
+import os
+
+def download_clips(stream_url, out_dir, start_time, end_time, resize=True, use_60fps=False):
+    # 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',
+                '-allowed_extensions', 'ALL',
+                '-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