Skip to content

Commit

Permalink
feat: add option to allow capture stderr from ffmpeg/ffprobe (#565)
Browse files Browse the repository at this point in the history
* feat: add option to allow capture stderr from ffmpeg/ffprobe

* add tests

* fix tests

* add comments

* fix decoding error

* fix

* fix
  • Loading branch information
Tao Peng authored Oct 11, 2022
1 parent 2bb2ef6 commit 2f3331c
Show file tree
Hide file tree
Showing 2 changed files with 151 additions and 29 deletions.
105 changes: 76 additions & 29 deletions mapillary_tools/ffmpeg.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,52 +34,117 @@ class Stream(TypedDict):
width: int


class ProbeFormat(TypedDict):
filename: str
duration: str


class ProbeOutput(TypedDict):
streams: T.List[Stream]
format: T.Dict


class FFmpegNotFoundError(Exception):
pass


_MAX_STDERR_LENGTH = 2048


def _truncate_begin(s: str) -> str:
if _MAX_STDERR_LENGTH < len(s):
return "..." + s[-_MAX_STDERR_LENGTH:]
else:
return s


def _truncate_end(s: str) -> str:
if _MAX_STDERR_LENGTH < len(s):
return s[:_MAX_STDERR_LENGTH] + "..."
else:
return s


class FFmpegCalledProcessError(Exception):
def __init__(self, ex: subprocess.CalledProcessError):
self.inner_ex = ex

def __str__(self) -> str:
msg = str(self.inner_ex)
if self.inner_ex.stderr is not None:
try:
stderr = self.inner_ex.stderr.decode("utf-8")
except UnicodeDecodeError:
stderr = str(self.inner_ex.stderr)
msg += f"\nSTDERR: {_truncate_begin(stderr)}"
return msg


class FFMPEG:
def __init__(
self, ffmpeg_path: str = "ffmpeg", ffprobe_path: str = "ffprobe"
self,
ffmpeg_path: str = "ffmpeg",
ffprobe_path: str = "ffprobe",
stderr: T.Optional[int] = None,
) -> None:
"""
ffmpeg_path: path to ffmpeg binary
ffprobe_path: path to ffprobe binary
stderr: param passed to subprocess.run to control whether to capture stderr
"""
self.ffmpeg_path = ffmpeg_path
self.ffprobe_path = ffprobe_path
self.stderr = stderr

def _run_ffprobe_json(self, cmd: T.List[str]) -> T.Dict:
full_cmd = [self.ffprobe_path, "-print_format", "json", *cmd]
LOG.info(f"Extracting video information: {' '.join(full_cmd)}")
try:
output = subprocess.check_output(full_cmd)
completed = subprocess.run(
full_cmd,
check=True,
stdout=subprocess.PIPE,
stderr=self.stderr,
)
except FileNotFoundError:
raise FFmpegNotFoundError(
f'The ffprobe command "{self.ffprobe_path}" not found'
)
except subprocess.CalledProcessError as ex:
raise FFmpegCalledProcessError(ex) from ex

try:
return json.loads(output)
stdout = completed.stdout.decode("utf-8")
except UnicodeDecodeError:
raise RuntimeError(
f"Error decoding ffprobe output as unicode: {_truncate_end(str(completed.stdout))}"
)

try:
output = json.loads(stdout)
except json.JSONDecodeError:
raise RuntimeError(
f"Error JSON decoding ffprobe output: {output.decode('utf-8')}"
f"Error JSON decoding ffprobe output: {_truncate_end(stdout)}"
)

# This check is for macOS:
# ffprobe -hide_banner -print_format json not_exists
# you will get exit code == 0 with the following stdout and stderr:
# {
# }
# not_exists: No such file or directory
if not output:
raise RuntimeError(
f"Empty JSON ffprobe output with STDERR: {_truncate_begin(str(completed.stderr))}"
)

return output

def _run_ffmpeg(self, cmd: T.List[str]) -> None:
full_cmd = [self.ffmpeg_path, *cmd]
LOG.info(f"Extracting frames: {' '.join(full_cmd)}")
try:
subprocess.check_call(full_cmd)
subprocess.run(full_cmd, check=True, stderr=self.stderr)
except FileNotFoundError:
raise FFmpegNotFoundError(
f'The ffmpeg command "{self.ffmpeg_path}" not found'
)
except subprocess.CalledProcessError as ex:
raise FFmpegCalledProcessError(ex) from ex

def probe_format_and_streams(self, video_path: Path) -> ProbeOutput:
cmd = [
Expand All @@ -90,24 +155,6 @@ def probe_format_and_streams(self, video_path: Path) -> ProbeOutput:
]
return T.cast(ProbeOutput, self._run_ffprobe_json(cmd))

def extract_stream(self, source: Path, dest: Path, stream_id: int) -> None:
cmd = [
"-hide_banner",
"-i",
str(source),
"-y", # overwrite - potentially dangerous
"-nostats",
"-codec",
"copy",
"-map",
f"0:{stream_id}",
"-f",
"rawvideo",
str(dest),
]

self._run_ffmpeg(cmd)

def extract_frames(
self,
video_path: Path,
Expand Down
75 changes: 75 additions & 0 deletions tests/unit/test_ffmpeg.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
import os
import subprocess
from pathlib import Path

import pytest

from mapillary_tools import ffmpeg


def ffmpeg_installed():
ffmpeg_path = os.getenv("MAPILLARY_TOOLS_FFMPEG_PATH", "ffmpeg")
ffprobe_path = os.getenv("MAPILLARY_TOOLS_FFPROBE_PATH", "ffprobe")
try:
subprocess.run([ffmpeg_path, "-version"])
# In Windows, ffmpeg is installed but ffprobe is not?
subprocess.run([ffprobe_path, "-version"])
except FileNotFoundError:
return False
return True


is_ffmpeg_installed = ffmpeg_installed()


def test_ffmpeg_not_exists():
if not is_ffmpeg_installed:
pytest.skip("ffmpeg not installed")

ff = ffmpeg.FFMPEG()
try:
ff.extract_frames(Path("not_exist_a"), Path("not_exist_b"), sample_interval=2)
except ffmpeg.FFmpegCalledProcessError as ex:
assert "STDERR:" not in str(ex)
else:
assert False, "FFmpegCalledProcessError not raised"

ff = ffmpeg.FFMPEG(stderr=subprocess.PIPE)
try:
ff.extract_frames(Path("not_exist_a"), Path("not_exist_b"), sample_interval=2)
except ffmpeg.FFmpegCalledProcessError as ex:
assert "STDERR:" in str(ex)
else:
assert False, "FFmpegCalledProcessError not raised"


def test_ffprobe_not_exists():
if not is_ffmpeg_installed:
pytest.skip("ffmpeg not installed")

ff = ffmpeg.FFMPEG()
try:
x = ff.probe_format_and_streams(Path("not_exist_a"))
except ffmpeg.FFmpegCalledProcessError as ex:
# exc from linux
assert "STDERR:" not in str(ex)
except RuntimeError as ex:
# exc from macos
assert "Empty JSON ffprobe output with STDERR: None" == str(ex)
else:
assert False, "RuntimeError not raised"

ff = ffmpeg.FFMPEG(stderr=subprocess.PIPE)
try:
x = ff.probe_format_and_streams(Path("not_exist_a"))
except ffmpeg.FFmpegCalledProcessError as ex:
# exc from linux
assert "STDERR:" in str(ex)
except RuntimeError as ex:
# exc from macos
assert (
"Empty JSON ffprobe output with STDERR: b'not_exist_a: No such file or directory"
in str(ex)
)
else:
assert False, "RuntimeError not raised"

0 comments on commit 2f3331c

Please sign in to comment.