tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Check the source video before removing previously extracted frames

Browse Source
This commit is contained in:
Afloat 2026-11-07 02:54:57 +08:00
parent 50e0e3c70c
commit ef30284229
2 changed files with 31 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
nerfstudio/process_data/process_data_utils.py
View File

@ -145,13 +145,6 @@ def convert_video_to_images(
A tuple containing summary of the conversion and the number of extracted frames.
"""
# If keep_image_dir is False, then remove the output image directory and its downscaled versions
if not keep_image_dir:
for i in range(num_downscales + 1):
dir_to_remove = image_dir if i == 0 else f"{image_dir}_{2**i}"
shutil.rmtree(dir_to_remove, ignore_errors=True)
image_dir.mkdir(exist_ok=True, parents=True)
for i in crop_factor:
if i < 0 or i > 1:
CONSOLE.print("[bold red]Error: Invalid crop factor. All crops must be in [0,1].")
@ -171,6 +164,13 @@ def convert_video_to_images(
sys.exit(1)
CONSOLE.print("Number of frames in video:", num_frames)
# If keep_image_dir is False, then remove the output image directory and its downscaled versions
if not keep_image_dir:
for i in range(num_downscales + 1):
dir_to_remove = image_dir if i == 0 else f"{image_dir}_{2**i}"
shutil.rmtree(dir_to_remove, ignore_errors=True)
image_dir.mkdir(exist_ok=True, parents=True)
ffmpeg_cmd = f'ffmpeg -i "{video_path}"'
crop_cmd = ""

24
tests/process_data/test_video_input_validation.py Normal file
View File

@ -0,0 +1,24 @@
from pathlib import Path
from unittest import mock
import pytest
from nerfstudio.process_data.process_data_utils import convert_video_to_images
@pytest.mark.parametrize("empty_video", [False, True])
def test_invalid_video_preserves_existing_frames(tmp_path: Path, empty_video: bool):
image_dir = tmp_path / "images"
old_frames = [directory / "frame_00001.png" for directory in (image_dir, tmp_path / "images_2")]
for frame in old_frames:
frame.parent.mkdir()
frame.write_bytes(b"previous frame")
video = tmp_path / "capture.mp4"
if empty_video:
video.touch()
with mock.patch("nerfstudio.process_data.process_data_utils.get_num_frames_in_video", return_value=0):
with pytest.raises(SystemExit):
convert_video_to_images(video, image_dir, num_frames_target=3, num_downscales=1)
assert all(frame.read_bytes() == b"previous frame" for frame in old_frames)