From 9f9ef66a851ebec5356ab9e4cd3f6ab9bcf79bfd Mon Sep 17 00:00:00 2001
From: Mike Lang <mikelang3000@gmail.com>
Date: Mon, 7 Feb 2022 20:48:26 +1100
Subject: [PATCH] Add endpoint to get a given frame of video

---
 common/common/segments.py     | 53 +++++++++++++++++++++++++++++++++++
 restreamer/restreamer/main.py | 25 ++++++++++++++++-
 2 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/common/common/segments.py b/common/common/segments.py
index c81cee8..7790225 100644
--- a/common/common/segments.py
+++ b/common/common/segments.py
@@ -587,3 +587,56 @@ def render_segments_waveform(segments, size=(1024, 128), scale='sqrt', color='#0
 				except (OSError, IOError):
 					pass
 
+
+@timed('extract_frame')
+def extract_frame(segments, timestamp):
+	"""
+	Extract the frame at TIMESTAMP within SEGMENT, yielding it as chunks of PNG data.
+	"""
+
+	# Remove holes
+	segments = [segment for segment in segments if segment is not None]
+
+	# Find segment containing timestamp
+	segments = [
+		segment for segment in segments
+		if segment.start <= timestamp < segment.end
+	]
+	if not segments:
+		raise ValueError("No data at timestamp within segment list")
+	if len(segments) != 1:
+		raise ValueError("Segment list contains overlap at timestamp")
+	(segment,) = segments
+
+	# "cut" input so that first frame is our target frame
+	cut_start = (timestamp - segment.start).total_seconds()
+
+	ffmpeg = None
+	input_feeder = None
+	try:
+		args = [
+			# get a single frame
+			'-vframes', '1',
+			# output as png
+			'-f', 'image2', '-c', 'png',
+		]
+		ffmpeg = ffmpeg_cut_stdin(subprocess.PIPE, cut_start=cut_start, duration=None, encode_args=args)
+		input_feeder = gevent.spawn(feed_input, segments, ffmpeg.stdin)
+
+		for chunk in read_chunks(ffmpeg.stdout):
+			yield chunk
+
+		# check if any errors occurred in input writing, or if ffmpeg exited non-success.
+		if ffmpeg.wait() != 0:
+			raise Exception("Error while extracting frame: ffmpeg exited {}".format(ffmpeg.returncode))
+		input_feeder.get() # re-raise any errors from feed_input()
+	finally:
+		# if something goes wrong, try to clean up ignoring errors
+		if input_feeder is not None:
+			input_feeder.kill()
+		if ffmpeg is not None and ffmpeg.poll() is None:
+			for action in (ffmpeg.kill, ffmpeg.stdin.close, ffmpeg.stdout.close):
+				try:
+					action()
+				except (OSError, IOError):
+					pass
diff --git a/restreamer/restreamer/main.py b/restreamer/restreamer/main.py
index 5269693..4706810 100644
--- a/restreamer/restreamer/main.py
+++ b/restreamer/restreamer/main.py
@@ -16,7 +16,7 @@ from gevent.pywsgi import WSGIServer
 
 from common import dateutil, get_best_segments, rough_cut_segments, fast_cut_segments, full_cut_segments, PromLogCountsHandler, install_stacksampler, serve_with_graceful_shutdown
 from common.flask_stats import request_stats, after_request
-from common.segments import feed_input, render_segments_waveform
+from common.segments import feed_input, render_segments_waveform, extract_frame
 
 from . import generate_hls
 
@@ -356,6 +356,29 @@ def generate_waveform(channel, quality):
 	return Response(render_segments_waveform(segments, (width, height)), mimetype='image/png')
 
 
+@app.route('/frame/<channel>/<quality>.png')
+@request_stats
+@has_path_args
+def get_frame(channel, quality):
+	"""
+	Returns a PNG image for the frame at the specific timestamp given.
+	Params:
+		timestamp: Required. The timestamp to get.
+			Must be in ISO 8601 format (ie. yyyy-mm-ddTHH:MM:SS) and UTC.
+	"""
+	timestamp = dateutil.parse_utc_only(request.args['timestamp'])
+
+	hours_path = os.path.join(app.static_folder, channel, quality)
+	if not os.path.isdir(hours_path):
+		abort(404)
+
+	segments = get_best_segments(hours_path, timestamp, timestamp)
+	if not any(segment is not None for segment in segments):
+		return "We have no content available within the requested time range.", 406
+
+	return Response(extract_frame(segments, timestamp), mimetype='image/png')
+
+
 @app.route('/generate_videos/<channel>/<quality>', methods=['POST'])
 @request_stats
 @has_path_args