Fix ffprobe output parsing and unquoted paths in shell commands

requirements.txt

@@ -1,5 +1,5 @@
-torch==2.8.0
-torchvision==0.23.0
+torch==2.9.0
+torchvision==0.24.0
 tqdm==4.67.1
 numpy==2.1.2
 pandas==2.2.3

utils/probe.py

@@ -16,6 +16,7 @@
 """
 
 import math
+import re
 import subprocess
 
 
@@ -39,7 +40,7 @@ def get_dimensions(
     result = subprocess.run(
         cmd,
         stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
+        stderr=subprocess.PIPE,
         check=True,
         text=True,
     )
@@ -72,7 +73,7 @@ def get_nb_frames(video_path, ffprobe_path="ffprobe") -> int | None:
     result = subprocess.run(
         cmd,
         stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
+        stderr=subprocess.PIPE,
         check=True,
         text=True,
     )
@@ -103,7 +104,7 @@ def get_r_frame_rate(video_path, ffprobe_path="ffprobe") -> int | None:
     result = subprocess.run(
         cmd,
         stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
+        stderr=subprocess.PIPE,
         check=True,
         text=True,
     )
@@ -137,10 +138,18 @@ def get_video_duration(video_path, ffprobe_path="ffprobe") -> float | None:
   ]
   try:
     result = subprocess.run(
-        cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, check=True
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        check=True,
+        text=True,
     )
-    duration = float(result.stdout)
-    return duration
+    matches = re.findall(r"^\s*(\d+\.?\d*)\s*$", result.stdout, re.MULTILINE)
+    if not matches:
+      print(f"Could not get duration for {video_path}")
+      return None
+    return float(matches[-1])
   except Exception as e:
     print(f"Error getting duration for {video_path}: {e}")
     return None
+

utils/video_reader.py

@@ -17,6 +17,7 @@
 
 import logging
 import os
+import shlex
 import subprocess
 import tempfile
 
@@ -73,10 +74,12 @@ def load_video_1p0(
       f"[tmp]scale={input_width_content}:{input_height_content}:flags=bilinear[out2]"
   )
   cmd = (
-      f"{ffmpeg_path}  -i {filepath} -filter_complex \"{filter_complex}\""
-      f" -map [out1] -r {video_fps} -f rawvideo -pix_fmt rgb24 -y {temp_filename}"
+      f"{ffmpeg_path}  -i {shlex.quote(filepath)} -filter_complex"
+      f" \"{filter_complex}\""
+      f" -map [out1] -r {video_fps} -f rawvideo -pix_fmt rgb24 -y"
+      f" {shlex.quote(temp_filename)}"
       f" -map [out2] -r {video_fps} -f rawvideo -pix_fmt rgb24 -y"
-      f" {temp_filename_small}"
+      f" {shlex.quote(temp_filename_small)}"
   )
 
   try:
@@ -190,9 +193,9 @@ def load_video_1p5(
   # Sample at constant frame rate, and save as RGB24 (RGBRGB...)
   fd, temp_filename = tempfile.mkstemp()
   cmd = (
-      f"{ffmpeg_path} -i {filepath} -vf"
+      f"{ffmpeg_path} -i {shlex.quote(filepath)} -vf"
       f" {transpose_param}scale=w={video_width}:h={video_height}:flags=bicubic,format=rgb24"
-      f" -r {video_fps} -f rawvideo -pix_fmt rgb24 -y {temp_filename}"
+      f" -r {video_fps} -f rawvideo -pix_fmt rgb24 -y {shlex.quote(temp_filename)}"
   )
 
   try:
@@ -246,3 +249,4 @@ def load_video_1p5(
   logging.info("Load %s done successfully.", filepath)
 
   return video, num_real_frames
+