diff --git a/requirements.txt b/requirements.txt index 9adcb82..8d77a79 100644 --- a/requirements.txt +++ b/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 diff --git a/utils/probe.py b/utils/probe.py index ec7b485..5d89c1e 100644 --- a/utils/probe.py +++ b/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 + diff --git a/utils/video_reader.py b/utils/video_reader.py index 50ed956..e716637 100644 --- a/utils/video_reader.py +++ b/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 +