Skip to content

Commit

Permalink
ffmpeg_reader: set default streams automatically if unspecified
Browse files Browse the repository at this point in the history
  • Loading branch information
bzczb committed Jan 28, 2023
1 parent 858bb81 commit 03d3e0f
Showing 1 changed file with 34 additions and 30 deletions.
64 changes: 34 additions & 30 deletions moviepy/video/io/ffmpeg_reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -344,8 +344,9 @@ def _reset_state(self):
# should be ignored
self._inside_output = False

# flag which indicates that a default stream has not been found yet
self._default_stream_found = False
# map from stream type to default stream
# if a default stream is not indicated, pick the first one available
self._default_streams = {}

# current input file, stream and chapter, which will be built at runtime
self._current_input_file = {"streams": []}
Expand Down Expand Up @@ -444,20 +445,13 @@ def parse(self):
"stream_number": stream_number,
"stream_type": stream_type_lower,
"language": language,
"default": not self._default_stream_found
or line.endswith("(default)"),
"default": (stream_type_lower not in self._default_streams) and line.endswith("(default)"),
}
self._default_stream_found = True

# for default streams, set their numbers globally, so it's
# easy to get without iterating all
if self._current_stream["default"]:
self.result[
f"default_{stream_type_lower}_input_number"
] = input_number
self.result[
f"default_{stream_type_lower}_stream_number"
] = stream_number
self._default_streams[stream_type_lower] = self._current_stream

# exit chapter
if self._current_chapter:
Expand All @@ -476,21 +470,20 @@ def parse(self):
input_number
]

# add new input file to self.result
# add new input file to result
self.result["inputs"].append(self._current_input_file)
self._current_input_file = {"input_number": input_number}

# parse relevant data by stream type
try:
global_data, stream_data = self.parse_data_by_stream_type(
stream_data = self.parse_data_by_stream_type(
stream_type, line
)
except NotImplementedError as exc:
warnings.warn(
f"{str(exc)}\nffmpeg output:\n\n{self.infos}", UserWarning
)
else:
self.result.update(global_data)
self._current_stream.update(stream_data)
elif line.startswith(" Metadata:"):
# enter group " Metadata:"
Expand Down Expand Up @@ -553,7 +546,7 @@ def parse(self):
self._last_metadata_field_added = field
self._current_chapter["metadata"][field] = value

# last input file, must be included in self.result
# last input file, must be included in the result
if self._current_input_file:
self._current_input_file["streams"].append(self._current_stream)
# include their chapters, if there are
Expand All @@ -563,6 +556,27 @@ def parse(self):
]
self.result["inputs"].append(self._current_input_file)

# set any missing default automatically
for stream in self._current_input_file["streams"]:
if stream["stream_type"] not in self._default_streams:
self._default_streams[stream["stream_type"]] = stream
stream["default"] = True

# set some global info based on the defaults
for stream_type_lower, stream_data in self._default_streams.items():
self.result[f"default_{stream_type_lower}_input_number"] = stream_data["input_number"]
self.result[f"default_{stream_type_lower}_stream_number"] = stream_data["stream_number"]

if stream_type_lower == "audio":
self.result["audio_found"] = True
self.result["audio_fps"] = stream_data["fps"]
self.result["audio_bitrate"] = stream_data["bitrate"]
elif stream_type_lower == "video":
self.result["video_found"] = True
self.result["video_size"] = stream_data.get("size", None)
self.result["video_bitrate"] = stream_data.get("bitrate", None)
self.result["video_fps"] = stream_data["fps"]

# some video duration utilities
if self.result["video_found"] and self.check_duration:
self.result["video_n_frames"] = int(
Expand Down Expand Up @@ -601,7 +615,7 @@ def parse_data_by_stream_type(self, stream_type, line):
return {
"Audio": self.parse_audio_stream_data,
"Video": self.parse_video_stream_data,
"Data": lambda _line: ({}, {}),
"Data": lambda _line: {},
}[stream_type](line)
except KeyError:
raise NotImplementedError(
Expand All @@ -611,7 +625,7 @@ def parse_data_by_stream_type(self, stream_type, line):

def parse_audio_stream_data(self, line):
"""Parses data from "Stream ... Audio" line."""
global_data, stream_data = ({"audio_found": True}, {})
stream_data = {}
try:
stream_data["fps"] = int(re.search(r" (\d+) Hz", line).group(1))
except (AttributeError, ValueError):
Expand All @@ -622,14 +636,11 @@ def parse_audio_stream_data(self, line):
stream_data["bitrate"] = (
int(match_audio_bitrate.group(1)) if match_audio_bitrate else None
)
if self._current_stream["default"]:
global_data["audio_fps"] = stream_data["fps"]
global_data["audio_bitrate"] = stream_data["bitrate"]
return (global_data, stream_data)
return stream_data

def parse_video_stream_data(self, line):
"""Parses data from "Stream ... Video" line."""
global_data, stream_data = ({"video_found": True}, {})
stream_data = {}

try:
match_video_size = re.search(r" (\d+)x(\d+)[,\s]", line)
Expand Down Expand Up @@ -681,14 +692,7 @@ def parse_video_stream_data(self, line):
fps = x * coef
stream_data["fps"] = fps

if self._current_stream["default"] or "video_size" not in self.result:
global_data["video_size"] = stream_data.get("size", None)
if self._current_stream["default"] or "video_bitrate" not in self.result:
global_data["video_bitrate"] = stream_data.get("bitrate", None)
if self._current_stream["default"] or "video_fps" not in self.result:
global_data["video_fps"] = stream_data["fps"]

return (global_data, stream_data)
return stream_data

def parse_fps(self, line):
"""Parses number of FPS from a line of the ``ffmpeg -i`` command output."""
Expand Down

0 comments on commit 03d3e0f

Please sign in to comment.