diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index 75878d6a5..7fc0f6f76 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -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 = "" diff --git a/tests/process_data/test_video_input_validation.py b/tests/process_data/test_video_input_validation.py new file mode 100644 index 000000000..6e46aaad2 --- /dev/null +++ b/tests/process_data/test_video_input_validation.py @@ -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)