Skip to content

Commit

Permalink
simplify
Browse files Browse the repository at this point in the history
  • Loading branch information
@mthrok
mthrok committed May 19, 2022
1 parent dd96163 commit abe8ce9
Showing 1 changed file with 20 additions and 98 deletions.
118 changes: 20 additions & 98 deletions torchaudio/io/_stream_reader.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -187,11 +187,9 @@ def __init__(self, src: str, fmt: Optional[str], option: Optional[Dict[str, str]
i = torch.ops.torchaudio.ffmpeg_streamer_find_best_video_stream(self._s)
self.default_video_stream = None if i < 0 else i

@property
def num_src_streams(self) -> int:
return torch.ops.torchaudio.ffmpeg_streamer_num_src_streams(self._s)

@property
def num_out_streams(self) -> int:
return torch.ops.torchaudio.ffmpeg_streamer_num_out_streams(self._s)

Expand All @@ -209,9 +207,9 @@ def add_audio_stream(
stream_index: int,
frames_per_chunk: int,
buffer_chunk_size: int,
filter_desc: Optional[str],
decoder: Optional[str],
decoder_options: Optional[Dict[str, str]],
filter_desc: Optional[str],
):
torch.ops.torchaudio.ffmpeg_streamer_add_audio_stream(
self._s,
Expand All @@ -228,10 +226,10 @@ def add_video_stream(
stream_index: int,
frames_per_chunk: int,
buffer_chunk_size: int,
filter_desc: Optional[str],
decoder: Optional[str],
decoder_options: Optional[Dict[str, str]],
hw_accel: Optional[str],
filter_desc: Optional[str],
):
torch.ops.torchaudio.ffmpeg_streamer_add_video_stream(
self._s,
Expand Down Expand Up @@ -260,91 +258,6 @@ def pop_chunks(self) -> Tuple[Optional[torch.Tensor]]:
return torch.ops.torchaudio.ffmpeg_streamer_pop_chunks(self._s)


class _PyBindBackend:
def __init__(
self,
src,
fmt: Optional[str],
option: Optional[Dict[str, str]],
buffer_size: int,
):
self._s = torchaudio._torchaudio_ffmpeg.StreamReaderFileObj(src, fmt, option, buffer_size)
i = self._s.find_best_audio_stream()
self.default_audio_stream = None if i < 0 else i
i = self._s.find_best_video_stream()
self.default_video_stream = None if i < 0 else i

@property
def num_src_streams(self) -> int:
return self._s.num_src_streams()

@property
def num_out_streams(self) -> int:
return self._s.num_out_streams()

def get_src_stream_info(self, i: int) -> torchaudio.io.StreamReaderSourceStream:
return self._s.get_src_stream_info(i)

def get_out_stream_info(self, i: int) -> torchaudio.io.StreamReaderOutputStream:
return self._s.get_out_stream_info(i)

def seek(self, timestamp: float):
self._s.seek(timestamp)

def add_audio_stream(
self,
stream_index: int,
frames_per_chunk: int,
buffer_chunk_size: int,
decoder: Optional[str],
decoder_options: Optional[Dict[str, str]],
filter_desc: Optional[str],
):
self._s.add_audio_stream(
stream_index,
frames_per_chunk,
buffer_chunk_size,
filter_desc,
decoder,
decoder_options or {},
)

def add_video_stream(
self,
stream_index: int,
frames_per_chunk: int,
buffer_chunk_size: int,
decoder: Optional[str],
decoder_options: Optional[Dict[str, str]],
hw_accel: Optional[str],
filter_desc: Optional[str],
):
self._s.add_video_stream(
stream_index,
frames_per_chunk,
buffer_chunk_size,
filter_desc,
decoder,
decoder_options or {},
hw_accel,
)

def remove_stream(self, i: int):
self._s.remove_stream(i)

def process_packet(self, timeout: Optional[float], backoff: float):
return self._s.process_packet(timeout, backoff)

def process_all_packets(self):
self._s.process_all_packets()

def is_buffer_ready(self) -> bool:
return self._s.is_buffer_ready()

def pop_chunks(self) -> Tuple[Optional[torch.Tensor]]:
return self._s.pop_chunks()


class StreamReader:
"""Fetch and decode audio/video streams chunk by chunk.
Expand Down Expand Up @@ -412,8 +325,17 @@ def __init__(
):
if isinstance(src, str):
self._be = _TorchBindBackend(src, format, option)
i = torch.ops.torchaudio.ffmpeg_streamer_find_best_audio_stream(self._be._s)
self._default_audio_stream = None if i < 0 else i
i = torch.ops.torchaudio.ffmpeg_streamer_find_best_video_stream(self._be._s)
self._default_video_stream = None if i < 0 else i
elif hasattr(src, "read"):
self._be = _PyBindBackend(src, format, option, buffer_size)
self._be = torchaudio._torchaudio_ffmpeg.StreamReaderFileObj(
src, format, option, buffer_size)
i = self._be.find_best_audio_stream()
self._default_audio_stream = None if i < 0 else i
i = self._be.find_best_video_stream()
self._default_video_stream = None if i < 0 else i
else:
raise ValueError("`src` must be either string or file-like object.")

Expand All @@ -423,31 +345,31 @@ def num_src_streams(self):
:type: int
"""
return self._be.num_src_streams
return self._be.num_src_streams()

@property
def num_out_streams(self):
"""Number of output streams configured by client code.
:type: int
"""
return self._be.num_out_streams
return self._be.num_out_streams()

@property
def default_audio_stream(self):
"""The index of default audio stream. ``None`` if there is no audio stream
:type: Optional[int]
"""
return self._be.default_audio_stream
return self._default_audio_stream

@property
def default_video_stream(self):
"""The index of default video stream. ``None`` if there is no video stream
:type: Optional[int]
"""
return self._be.default_video_stream
return self._default_video_stream

def get_src_stream_info(self, i: int) -> torchaudio.io.StreamReaderSourceStream:
"""Get the metadata of source stream
Expand Down Expand Up @@ -649,9 +571,9 @@ def add_audio_stream(
i,
frames_per_chunk,
buffer_chunk_size,
decoder,
decoder_options,
filter_desc,
decoder,
decoder_options or {},
)

def add_video_stream(
Expand Down Expand Up @@ -740,10 +662,10 @@ def add_video_stream(
i,
frames_per_chunk,
buffer_chunk_size,
filter_desc,
decoder,
decoder_options,
decoder_options or {},
hw_accel,
filter_desc,
)

def remove_stream(self, i: int):
Expand Down

0 comments on commit abe8ce9

Please sign in to comment.