2019-03-12 02:57:10 +00:00
|
|
|
"""Provides the worker thread needed for processing streams."""
|
2020-08-20 03:18:54 +00:00
|
|
|
from collections import deque
|
2019-03-12 02:57:10 +00:00
|
|
|
import io
|
|
|
|
import logging
|
|
|
|
|
2019-10-14 21:20:18 +00:00
|
|
|
import av
|
|
|
|
|
2020-10-25 02:55:12 +00:00
|
|
|
from .const import (
|
|
|
|
MAX_MISSING_DTS,
|
|
|
|
MAX_TIMESTAMP_GAP,
|
|
|
|
MIN_SEGMENT_DURATION,
|
|
|
|
PACKETS_TO_WAIT_FOR_AUDIO,
|
|
|
|
STREAM_TIMEOUT,
|
|
|
|
)
|
2019-03-12 02:57:10 +00:00
|
|
|
from .core import Segment, StreamBuffer
|
|
|
|
|
|
|
|
_LOGGER = logging.getLogger(__name__)
|
|
|
|
|
|
|
|
|
2020-08-20 03:18:54 +00:00
|
|
|
def create_stream_buffer(stream_output, video_stream, audio_stream, sequence):
|
2019-03-12 02:57:10 +00:00
|
|
|
"""Create a new StreamBuffer."""
|
2019-07-31 19:25:30 +00:00
|
|
|
|
2019-03-12 02:57:10 +00:00
|
|
|
segment = io.BytesIO()
|
2020-08-20 03:18:54 +00:00
|
|
|
container_options = (
|
|
|
|
stream_output.container_options(sequence)
|
|
|
|
if stream_output.container_options
|
|
|
|
else {}
|
|
|
|
)
|
2020-08-11 21:12:41 +00:00
|
|
|
output = av.open(
|
|
|
|
segment,
|
|
|
|
mode="w",
|
|
|
|
format=stream_output.format,
|
2020-09-11 18:07:45 +00:00
|
|
|
container_options={
|
|
|
|
"video_track_timescale": str(int(1 / video_stream.time_base)),
|
|
|
|
**container_options,
|
|
|
|
},
|
2020-08-11 21:12:41 +00:00
|
|
|
)
|
2019-04-05 06:40:22 +00:00
|
|
|
vstream = output.add_stream(template=video_stream)
|
2019-03-12 02:57:10 +00:00
|
|
|
# Check if audio is requested
|
|
|
|
astream = None
|
2020-08-20 03:18:54 +00:00
|
|
|
if audio_stream and audio_stream.name in stream_output.audio_codecs:
|
|
|
|
astream = output.add_stream(template=audio_stream)
|
|
|
|
return StreamBuffer(segment, output, vstream, astream)
|
2019-03-12 02:57:10 +00:00
|
|
|
|
|
|
|
|
2021-02-15 17:52:37 +00:00
|
|
|
class SegmentBuffer:
|
|
|
|
"""Buffer for writing a sequence of packets to the output as a segment."""
|
|
|
|
|
|
|
|
def __init__(self, video_stream, audio_stream, outputs_callback) -> None:
|
|
|
|
"""Initialize SegmentBuffer."""
|
|
|
|
self._video_stream = video_stream
|
|
|
|
self._audio_stream = audio_stream
|
|
|
|
self._outputs_callback = outputs_callback
|
|
|
|
# tuple of StreamOutput, StreamBuffer
|
|
|
|
self._outputs = []
|
|
|
|
self._sequence = 0
|
|
|
|
self._segment_start_pts = None
|
|
|
|
|
|
|
|
def reset(self, video_pts):
|
|
|
|
"""Initialize a new stream segment."""
|
|
|
|
# Keep track of the number of segments we've processed
|
|
|
|
self._sequence += 1
|
|
|
|
self._segment_start_pts = video_pts
|
|
|
|
|
|
|
|
# Fetch the latest StreamOutputs, which may have changed since the
|
|
|
|
# worker started.
|
|
|
|
self._outputs = []
|
|
|
|
for stream_output in self._outputs_callback().values():
|
|
|
|
if self._video_stream.name not in stream_output.video_codecs:
|
|
|
|
continue
|
|
|
|
buffer = create_stream_buffer(
|
|
|
|
stream_output, self._video_stream, self._audio_stream, self._sequence
|
|
|
|
)
|
|
|
|
self._outputs.append((buffer, stream_output))
|
|
|
|
|
|
|
|
def mux_packet(self, packet):
|
|
|
|
"""Mux a packet to the appropriate StreamBuffers."""
|
|
|
|
|
|
|
|
# Check for end of segment
|
|
|
|
if packet.stream == self._video_stream and packet.is_keyframe:
|
|
|
|
duration = (packet.pts - self._segment_start_pts) * packet.time_base
|
|
|
|
if duration >= MIN_SEGMENT_DURATION:
|
|
|
|
# Save segment to outputs
|
|
|
|
self.flush(duration)
|
|
|
|
|
|
|
|
# Reinitialize
|
|
|
|
self.reset(packet.pts)
|
|
|
|
|
|
|
|
# Mux the packet
|
|
|
|
for (buffer, _) in self._outputs:
|
|
|
|
if packet.stream == self._video_stream:
|
|
|
|
packet.stream = buffer.vstream
|
|
|
|
elif packet.stream == self._audio_stream:
|
|
|
|
packet.stream = buffer.astream
|
|
|
|
else:
|
|
|
|
continue
|
|
|
|
buffer.output.mux(packet)
|
|
|
|
|
|
|
|
def flush(self, duration):
|
|
|
|
"""Create a segment from the buffered packets and write to output."""
|
|
|
|
for (buffer, stream_output) in self._outputs:
|
|
|
|
buffer.output.close()
|
|
|
|
stream_output.put(Segment(self._sequence, buffer.segment, duration))
|
|
|
|
|
|
|
|
def close(self):
|
|
|
|
"""Close all StreamBuffers."""
|
|
|
|
for (buffer, _) in self._outputs:
|
|
|
|
buffer.output.close()
|
|
|
|
|
|
|
|
|
|
|
|
def stream_worker(source, options, outputs_callback, quit_event):
|
2019-03-12 02:57:10 +00:00
|
|
|
"""Handle consuming streams."""
|
2019-07-31 19:25:30 +00:00
|
|
|
|
2020-11-11 17:32:56 +00:00
|
|
|
try:
|
2021-02-15 17:52:37 +00:00
|
|
|
container = av.open(source, options=options, timeout=STREAM_TIMEOUT)
|
2020-11-11 17:32:56 +00:00
|
|
|
except av.AVError:
|
2021-02-15 17:52:37 +00:00
|
|
|
_LOGGER.error("Error opening stream %s", source)
|
2020-11-11 17:32:56 +00:00
|
|
|
return
|
2019-03-12 02:57:10 +00:00
|
|
|
try:
|
|
|
|
video_stream = container.streams.video[0]
|
|
|
|
except (KeyError, IndexError):
|
|
|
|
_LOGGER.error("Stream has no video")
|
2020-09-10 19:55:55 +00:00
|
|
|
container.close()
|
2019-03-12 02:57:10 +00:00
|
|
|
return
|
2020-08-20 03:18:54 +00:00
|
|
|
try:
|
|
|
|
audio_stream = container.streams.audio[0]
|
|
|
|
except (KeyError, IndexError):
|
|
|
|
audio_stream = None
|
2020-09-10 19:55:55 +00:00
|
|
|
# These formats need aac_adtstoasc bitstream filter, but auto_bsf not
|
|
|
|
# compatible with empty_moov and manual bitstream filters not in PyAV
|
|
|
|
if container.format.name in {"hls", "mpegts"}:
|
|
|
|
audio_stream = None
|
2020-09-24 12:35:52 +00:00
|
|
|
# Some audio streams do not have a profile and throw errors when remuxing
|
|
|
|
if audio_stream and audio_stream.profile is None:
|
|
|
|
audio_stream = None
|
2019-03-12 02:57:10 +00:00
|
|
|
|
2020-10-25 02:55:12 +00:00
|
|
|
# Iterator for demuxing
|
|
|
|
container_packets = None
|
2020-08-20 03:18:54 +00:00
|
|
|
# The decoder timestamps of the latest packet in each stream we processed
|
2020-11-16 20:13:33 +00:00
|
|
|
last_dts = {video_stream: float("-inf"), audio_stream: float("-inf")}
|
2020-07-13 13:47:33 +00:00
|
|
|
# Keep track of consecutive packets without a dts to detect end of stream.
|
2020-10-25 02:55:12 +00:00
|
|
|
missing_dts = 0
|
2020-08-20 03:18:54 +00:00
|
|
|
# Holds the buffers for each stream provider
|
2021-02-15 17:52:37 +00:00
|
|
|
segment_buffer = SegmentBuffer(video_stream, audio_stream, outputs_callback)
|
2020-08-20 03:18:54 +00:00
|
|
|
# The video pts at the beginning of the segment
|
|
|
|
segment_start_pts = None
|
|
|
|
# Because of problems 1 and 2 below, we need to store the first few packets and replay them
|
|
|
|
initial_packets = deque()
|
|
|
|
|
|
|
|
# Have to work around two problems with RTSP feeds in ffmpeg
|
|
|
|
# 1 - first frame has bad pts/dts https://trac.ffmpeg.org/ticket/5018
|
|
|
|
# 2 - seeking can be problematic https://trac.ffmpeg.org/ticket/7815
|
|
|
|
|
|
|
|
def peek_first_pts():
|
2020-11-16 20:13:33 +00:00
|
|
|
"""Initialize by peeking into the first few packets of the stream.
|
2020-08-20 03:18:54 +00:00
|
|
|
|
2020-11-16 20:13:33 +00:00
|
|
|
Deal with problem #1 above (bad first packet pts/dts) by recalculating using pts/dts from second packet.
|
|
|
|
Also load the first video keyframe pts into segment_start_pts and check if the audio stream really exists.
|
|
|
|
"""
|
|
|
|
nonlocal segment_start_pts, audio_stream, container_packets
|
|
|
|
missing_dts = 0
|
|
|
|
found_audio = False
|
2020-08-20 03:18:54 +00:00
|
|
|
try:
|
2020-10-25 02:55:12 +00:00
|
|
|
container_packets = container.demux((video_stream, audio_stream))
|
2020-11-16 20:13:33 +00:00
|
|
|
first_packet = None
|
2020-08-20 03:18:54 +00:00
|
|
|
# Get to first video keyframe
|
2020-11-16 20:13:33 +00:00
|
|
|
while first_packet is None:
|
2020-10-25 02:55:12 +00:00
|
|
|
packet = next(container_packets)
|
2020-09-21 01:26:24 +00:00
|
|
|
if (
|
2020-10-11 16:52:51 +00:00
|
|
|
packet.dts is None
|
2020-10-25 02:55:12 +00:00
|
|
|
): # Allow MAX_MISSING_DTS packets with no dts, raise error on the next one
|
|
|
|
if missing_dts >= MAX_MISSING_DTS:
|
|
|
|
raise StopIteration(
|
|
|
|
f"Invalid data - got {MAX_MISSING_DTS+1} packets with missing DTS while initializing"
|
|
|
|
)
|
|
|
|
missing_dts += 1
|
2020-10-11 16:52:51 +00:00
|
|
|
continue
|
2020-11-16 20:13:33 +00:00
|
|
|
if packet.stream == audio_stream:
|
|
|
|
found_audio = True
|
|
|
|
elif packet.is_keyframe: # video_keyframe
|
|
|
|
first_packet = packet
|
2020-08-20 03:18:54 +00:00
|
|
|
initial_packets.append(packet)
|
|
|
|
# Get first_pts from subsequent frame to first keyframe
|
2020-11-16 20:13:33 +00:00
|
|
|
while segment_start_pts is None or (
|
|
|
|
audio_stream
|
|
|
|
and not found_audio
|
|
|
|
and len(initial_packets) < PACKETS_TO_WAIT_FOR_AUDIO
|
|
|
|
):
|
2020-10-25 02:55:12 +00:00
|
|
|
packet = next(container_packets)
|
2020-10-11 16:52:51 +00:00
|
|
|
if (
|
|
|
|
packet.dts is None
|
2020-10-25 02:55:12 +00:00
|
|
|
): # Allow MAX_MISSING_DTS packet with no dts, raise error on the next one
|
|
|
|
if missing_dts >= MAX_MISSING_DTS:
|
|
|
|
raise StopIteration(
|
|
|
|
f"Invalid data - got {MAX_MISSING_DTS+1} packets with missing DTS while initializing"
|
|
|
|
)
|
|
|
|
missing_dts += 1
|
2020-10-11 16:52:51 +00:00
|
|
|
continue
|
2020-11-16 20:13:33 +00:00
|
|
|
if packet.stream == audio_stream:
|
|
|
|
found_audio = True
|
|
|
|
elif (
|
|
|
|
segment_start_pts is None
|
|
|
|
): # This is the second video frame to calculate first_pts from
|
|
|
|
segment_start_pts = packet.dts - packet.duration
|
|
|
|
first_packet.pts = segment_start_pts
|
|
|
|
first_packet.dts = segment_start_pts
|
2020-08-20 03:18:54 +00:00
|
|
|
initial_packets.append(packet)
|
2020-11-16 20:13:33 +00:00
|
|
|
if audio_stream and not found_audio:
|
2020-08-20 03:18:54 +00:00
|
|
|
_LOGGER.warning(
|
|
|
|
"Audio stream not found"
|
|
|
|
) # Some streams declare an audio stream and never send any packets
|
|
|
|
audio_stream = None
|
|
|
|
|
|
|
|
except (av.AVError, StopIteration) as ex:
|
|
|
|
_LOGGER.error(
|
|
|
|
"Error demuxing stream while finding first packet: %s", str(ex)
|
|
|
|
)
|
2020-09-03 16:22:00 +00:00
|
|
|
return False
|
|
|
|
return True
|
2020-08-20 03:18:54 +00:00
|
|
|
|
2020-09-03 16:22:00 +00:00
|
|
|
if not peek_first_pts():
|
|
|
|
container.close()
|
|
|
|
return
|
2020-11-16 20:13:33 +00:00
|
|
|
|
2021-02-15 17:52:37 +00:00
|
|
|
segment_buffer.reset(segment_start_pts)
|
2019-03-12 02:57:10 +00:00
|
|
|
|
|
|
|
while not quit_event.is_set():
|
|
|
|
try:
|
2020-08-20 03:18:54 +00:00
|
|
|
if len(initial_packets) > 0:
|
|
|
|
packet = initial_packets.popleft()
|
|
|
|
else:
|
2020-10-25 02:55:12 +00:00
|
|
|
packet = next(container_packets)
|
2019-03-12 02:57:10 +00:00
|
|
|
if packet.dts is None:
|
2020-10-25 02:55:12 +00:00
|
|
|
# Allow MAX_MISSING_DTS consecutive packets without dts. Terminate the stream on the next one.
|
|
|
|
if missing_dts >= MAX_MISSING_DTS:
|
|
|
|
raise StopIteration(
|
|
|
|
f"No dts in {MAX_MISSING_DTS+1} consecutive packets"
|
|
|
|
)
|
|
|
|
missing_dts += 1
|
2020-07-13 13:47:33 +00:00
|
|
|
continue
|
2020-10-25 02:55:12 +00:00
|
|
|
missing_dts = 0
|
2019-03-12 02:57:10 +00:00
|
|
|
except (av.AVError, StopIteration) as ex:
|
2019-03-21 14:31:55 +00:00
|
|
|
_LOGGER.error("Error demuxing stream: %s", str(ex))
|
2019-03-12 02:57:10 +00:00
|
|
|
break
|
|
|
|
|
2020-08-20 03:18:54 +00:00
|
|
|
# Discard packet if dts is not monotonic
|
|
|
|
if packet.dts <= last_dts[packet.stream]:
|
2020-10-16 15:48:45 +00:00
|
|
|
if (
|
|
|
|
packet.time_base * (last_dts[packet.stream] - packet.dts)
|
|
|
|
> MAX_TIMESTAMP_GAP
|
2020-09-30 15:45:59 +00:00
|
|
|
):
|
2020-10-16 15:48:45 +00:00
|
|
|
_LOGGER.warning(
|
|
|
|
"Timestamp overflow detected: last dts %s, dts = %s, resetting stream",
|
|
|
|
last_dts[packet.stream],
|
|
|
|
packet.dts,
|
|
|
|
)
|
|
|
|
break
|
2019-03-21 14:31:55 +00:00
|
|
|
continue
|
2020-08-20 03:18:54 +00:00
|
|
|
|
|
|
|
# Update last_dts processed
|
|
|
|
last_dts[packet.stream] = packet.dts
|
2021-02-15 17:52:37 +00:00
|
|
|
|
|
|
|
# Mux packets, and possibly write a segment to the output stream.
|
|
|
|
# This mutates packet timestamps and stream
|
|
|
|
segment_buffer.mux_packet(packet)
|
2020-05-22 16:13:37 +00:00
|
|
|
|
|
|
|
# Close stream
|
2021-02-15 17:52:37 +00:00
|
|
|
segment_buffer.close()
|
2020-05-22 16:13:37 +00:00
|
|
|
container.close()
|