diff --git a/mapillary_tools/commands/sample_video.py b/mapillary_tools/commands/sample_video.py index 2573f800..1d3c3570 100644 --- a/mapillary_tools/commands/sample_video.py +++ b/mapillary_tools/commands/sample_video.py @@ -66,6 +66,19 @@ def add_basic_arguments(self, parser: argparse.ArgumentParser): default=False, required=False, ) + group.add_argument( + "--source_frame_names", + help=( + "Rename distance-sampled JPEGs from sequential 000001, 000002, … " + "to 0-based source frame numbers (…_000000.jpg, …_033717.jpg). " + "Gaps in the numbers are skipped video time. Mapillary only " + "needs unique names; interval sampling stays sequential. " + "[default: %(default)s]" + ), + action="store_true", + default=constants.SOURCE_FRAME_NAMES, + required=False, + ) def run(self, vars_args: dict): video_import_path: Path = vars_args["video_import_path"] diff --git a/mapillary_tools/constants.py b/mapillary_tools/constants.py index cab9e9b1..4a5dc4f2 100644 --- a/mapillary_tools/constants.py +++ b/mapillary_tools/constants.py @@ -91,6 +91,9 @@ def _parse_scaled_integers( VIDEO_SAMPLE_INTERVAL = float(os.getenv(_ENV_PREFIX + "VIDEO_SAMPLE_INTERVAL", -1)) # In meters VIDEO_SAMPLE_DISTANCE = float(os.getenv(_ENV_PREFIX + "VIDEO_SAMPLE_DISTANCE", 3)) +SOURCE_FRAME_NAMES: bool = _yes_or_no( + os.getenv(_ENV_PREFIX + "SOURCE_FRAME_NAMES", "NO") +) VIDEO_DURATION_RATIO = float(os.getenv(_ENV_PREFIX + "VIDEO_DURATION_RATIO", 1)) FFPROBE_PATH: str = os.getenv(_ENV_PREFIX + "FFPROBE_PATH", "ffprobe") FFMPEG_PATH: str = os.getenv(_ENV_PREFIX + "FFMPEG_PATH", "ffmpeg") diff --git a/mapillary_tools/ffmpeg.py b/mapillary_tools/ffmpeg.py index 0264e2b2..0fe68fdb 100644 --- a/mapillary_tools/ffmpeg.py +++ b/mapillary_tools/ffmpeg.py @@ -280,6 +280,7 @@ def extract_specified_frames( sample_dir: Path, frame_indices: set[int], stream_specifier: int | str = "v", + source_frame_names: bool = False, ) -> None: """ Extract specific frames from video by frame number using select filter. @@ -294,14 +295,19 @@ def extract_specified_frames( stream_specifier: Stream specifier to target specific stream(s). Can be an integer (stream index) or "v" (all video streams) See https://ffmpeg.org/ffmpeg.html#Stream-specifiers-1 + source_frame_names: If true, rename sequential ``-start_number`` + files onto 0-based source frame indices after extract. Raises: FFmpegNotFoundError: If ffmpeg binary is not found FFmpegCalledProcessError: If ffmpeg command fails Note: - Frame indices are 0-based but ffmpeg output files are numbered starting from 1. - Creates temporary filter script file on Windows to avoid command line length limits. + Frame indices are 0-based. FFmpeg writes sequential files + (``…_000001.jpg``, ``…_000002.jpg``). With + ``source_frame_names`` they become ``…_000000.jpg``, + ``…_033717.jpg``. Creates a temporary filter script file on + Windows to avoid command line length limits. """ self._validate_stream_specifier(stream_specifier) @@ -361,6 +367,41 @@ def extract_specified_frames( except FileNotFoundError: pass + if source_frame_names: + self._rename_extracted_to_source_indices( + sample_prefix, stream_specifier, sorted(frame_indices) + ) + + @classmethod + def _rename_extracted_to_source_indices( + cls, + sample_prefix: Path, + stream_specifier: int | str, + sorted_indices: list[int], + ) -> None: + """Map ffmpeg's sequential ``-start_number`` files onto source frame numbers.""" + if not sorted_indices: + return + + def path_for(n: int) -> Path: + return Path(f"{sample_prefix}_{stream_specifier}_{n:06d}{cls.FRAME_EXT}") + + sources = [path_for(i) for i in range(1, len(sorted_indices) + 1)] + dests = [path_for(idx) for idx in sorted_indices] + if sources == dests: + return + missing = [src for src in sources if not src.is_file()] + if missing: + raise RuntimeError( + f"expected {len(sorted_indices)} extracted frames under {sample_prefix.parent}, " + f"missing {missing[0].name}" + ) + tmps = [src.with_name(src.name + ".mlytmp") for src in sources] + for src, tmp in zip(sources, tmps): + src.rename(tmp) + for tmp, dst in zip(tmps, dests): + tmp.rename(dst) + @classmethod def sort_selected_samples( cls, @@ -434,7 +475,9 @@ def iterate_samples( Yields: Tuple containing: - stream_specifier (str): Stream specifier (number or "v") - - frame_idx (int): Frame index (0-based or 1-based depending on extraction method) + - frame_idx (int): Number parsed from the filename (sequential + 1-based after extract_specified_frames, or 0-based source + frame with ``source_frame_names``; 1-based for interval sampling) - sample_path (Path): Path to the frame image file Note: diff --git a/mapillary_tools/sample_video.py b/mapillary_tools/sample_video.py index 1e3e7764..c34fb988 100644 --- a/mapillary_tools/sample_video.py +++ b/mapillary_tools/sample_video.py @@ -60,8 +60,11 @@ def sample_video( video_start_time: str | None = None, skip_sample_errors: bool = False, rerun: bool = False, + source_frame_names: bool | None = None, ) -> None: video_dir, video_list = _normalize_path(video_import_path, skip_subfolders) + if source_frame_names is None: + source_frame_names = constants.SOURCE_FRAME_NAMES if not xor(0 <= video_sample_distance, 0 < video_sample_interval): raise exceptions.MapillaryBadParameterError( @@ -112,6 +115,7 @@ def sample_video( sample_dir, sample_distance=video_sample_distance, start_time=video_start_time_dt, + source_frame_names=source_frame_names, ) else: assert 0 < video_sample_interval, ( @@ -281,6 +285,7 @@ def _sample_single_video_by_distance( sample_dir: Path, sample_distance: float, start_time: datetime.datetime | None = None, + source_frame_names: bool = False, ) -> None: ffmpeg = ffmpeglib.FFMPEG(constants.FFMPEG_PATH, constants.FFPROBE_PATH) @@ -327,6 +332,7 @@ def _sample_single_video_by_distance( wip_dir, frame_indices=set(sorted_sample_indices), stream_specifier=str(video_stream_idx), + source_frame_names=source_frame_names, ) frame_samples = ffmpeglib.FFMPEG.sort_selected_samples( @@ -336,14 +342,26 @@ def _sample_single_video_by_distance( raise exceptions.MapillaryVideoError( f"Expect {len(sorted_sample_indices)} samples but extracted {len(frame_samples)} samples" ) - for idx, (frame_idx_1based, sample_paths) in enumerate(frame_samples): - assert len(sample_paths) == 1, ( - "Expect 1 sample path at {frame_idx_1based} but got {sample_paths}" - ) - if idx + 1 != frame_idx_1based: - raise exceptions.MapillaryVideoError( - f"Expect {sample_paths[0]} to be {idx + 1}th sample but got {frame_idx_1based}" + if source_frame_names: + for (file_idx, sample_paths), sample_idx in zip( + frame_samples, sorted_sample_indices + ): + assert len(sample_paths) == 1, ( + f"Expect 1 sample path at {file_idx} but got {sample_paths}" + ) + if file_idx != sample_idx: + raise exceptions.MapillaryVideoError( + f"Expect {sample_paths[0]} to be source frame {sample_idx} but got {file_idx}" + ) + else: + for idx, (frame_idx_1based, sample_paths) in enumerate(frame_samples): + assert len(sample_paths) == 1, ( + f"Expect 1 sample path at {frame_idx_1based} but got {sample_paths}" ) + if idx + 1 != frame_idx_1based: + raise exceptions.MapillaryVideoError( + f"Expect {sample_paths[0]} to be {idx + 1}th sample but got {frame_idx_1based}" + ) for (_, sample_paths), sample_idx in zip(frame_samples, sorted_sample_indices): if sample_paths[0] is None: diff --git a/tests/unit/test_ffmpeg.py b/tests/unit/test_ffmpeg.py index c79fc8ce..6aa8c78a 100644 --- a/tests/unit/test_ffmpeg.py +++ b/tests/unit/test_ffmpeg.py @@ -130,6 +130,57 @@ def test_ffmpeg_extract_specified_frames_empty_ok(setup_data: py.path.local): assert len(results) == 0 +def test_ffmpeg_extract_specified_frames_source_names(setup_data: py.path.local): + pytest_skip_if_not_ffmpeg_installed() + + ff = ffmpeg.FFMPEG() + + video_path = Path(setup_data.join("videos/sample-5s.mp4")) + + sample_dir = Path(setup_data.join("videos/samples_source_names")) + sample_dir.mkdir() + + ff.extract_specified_frames( + video_path, sample_dir, frame_indices={2, 9}, source_frame_names=True + ) + + results = list(ff.sort_selected_samples(sample_dir, video_path)) + assert [file_idx for file_idx, _ in results] == [2, 9] + for file_idx, frame_paths in results: + assert frame_paths[0] is not None + assert frame_paths[0].name.endswith(f"_{file_idx:06d}.jpg") + + +def test_rename_extracted_to_source_indices(tmp_path: Path): + video_stem = "GX040129" + prefix = tmp_path / video_stem + spec = 0 + sequential = [1, 2, 3] + source_frames = [0, 14, 33717] + for n in sequential: + (tmp_path / f"{video_stem}_{spec}_{n:06d}.jpg").write_bytes(b"x" * n) + + ffmpeg.FFMPEG._rename_extracted_to_source_indices(prefix, spec, source_frames) + + names = sorted(p.name for p in tmp_path.glob("*.jpg")) + assert names == [f"{video_stem}_{spec}_{idx:06d}.jpg" for idx in source_frames] + assert (tmp_path / f"{video_stem}_{spec}_{14:06d}.jpg").read_bytes() == b"xx" + + +def test_rename_extracted_to_source_indices_collision(tmp_path: Path): + video_stem = "clip" + prefix = tmp_path / video_stem + spec = "v" + (tmp_path / f"{video_stem}_{spec}_000001.jpg").write_bytes(b"a") + (tmp_path / f"{video_stem}_{spec}_000002.jpg").write_bytes(b"b") + + ffmpeg.FFMPEG._rename_extracted_to_source_indices(prefix, spec, [2, 9]) + + assert (tmp_path / f"{video_stem}_{spec}_000002.jpg").read_bytes() == b"a" + assert (tmp_path / f"{video_stem}_{spec}_000009.jpg").read_bytes() == b"b" + assert not (tmp_path / f"{video_stem}_{spec}_000001.jpg").exists() + + def test_probe_format_and_streams_ok(setup_data: py.path.local): pytest_skip_if_not_ffmpeg_installed() diff --git a/tests/unit/test_sample_video.py b/tests/unit/test_sample_video.py index 0743eeb4..1c9ef6ad 100644 --- a/tests/unit/test_sample_video.py +++ b/tests/unit/test_sample_video.py @@ -446,6 +446,7 @@ def fake_extract_frames( sample_dir: Path, frame_indices: set[int], stream_specifier: str = "v", + **_kwargs, ) -> None: _create_fake_frames( sample_dir,