Add ability to render waveforms in restreamer

common/common/segments.py

@@ -326,9 +326,13 @@ def ffmpeg_cut_stdin(output_file, cut_start, duration, encode_args):
 		'ffmpeg',
 		'-hide_banner', '-loglevel', 'error', # suppress noisy output
 		'-i', '-',
-		'-ss', cut_start,
-		'-t', duration,
-	] + list(encode_args)
+	]
+	if cut_start is not None:
+		args += ['-ss', cut_start]
+	if duration is not None:
+		args += ['-t', duration]
+	args += list(encode_args)
+
 	if output_file is subprocess.PIPE:
 		args.append('-') # output to stdout
 	else:
@@ -363,6 +367,8 @@ def rough_cut_segments(segments, start, end):
 	This method works by simply concatenating all the segments, without any re-encoding.
 	"""
 	for segment in segments:
+		if segment is None:
+			continue
 		with open(segment.path, 'rb') as f:
 			for chunk in read_chunks(f):
 				yield chunk
@@ -511,3 +517,52 @@ def full_cut_segments(segments, start, end, encode_args, stream=False):
 					action()
 				except (OSError, IOError):
 					pass
+
+
+@timed('waveform')
+def render_segments_waveform(segments, size=(1024, 128), scale='sqrt', color='#000000'):
+	"""
+	Render an audio waveform of given list of segments. Yields chunks of PNG data.
+	Note we do not validate our inputs before passing them into an ffmpeg filtergraph.
+	Do not provide untrusted input without verifying, or else they can run arbitrary filters
+	(this MAY be fine but I wouldn't be shocked if some obscure filter lets them do arbitrary
+	filesystem writes).
+	"""
+	width, height = size
+
+	# Remove holes
+	segments = [segment for segment in segments if segment is not None]
+
+	ffmpeg = None
+	input_feeder = None
+	try:
+		args = [
+			# create waveform from input audio
+			'-filter_complex',
+			f'[0:a]showwavespic=size={width}x{height}:colors={color}:scale={scale}[out]',
+			# use created waveform as our output
+			'-map', '[out]',
+			# output as png
+			'-f', 'image2', '-c', 'png',
+		]
+		ffmpeg = ffmpeg_cut_stdin(subprocess.PIPE, cut_start=None, 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 rendering waveform: 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
+

restreamer/restreamer/main.py

@@ -17,7 +17,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
 from common.flask_stats import request_stats, after_request
-from common.segments import feed_input
+from common.segments import feed_input, render_segments_waveform
 
 from . import generate_hls
 
@@ -299,6 +299,42 @@ def cut(channel, quality):
 		return "Unknown type {!r}".format(type), 400
 
 
+@app.route('/waveform/<channel>/<quality>.png')
+@request_stats
+@has_path_args
+def generate_waveform(channel, quality):
+	"""
+	Returns a PNG image showing the audio waveform over the requested time period.
+	Params:
+		start, end: Required. The start and end times.
+			Must be in ISO 8601 format (ie. yyyy-mm-ddTHH:MM:SS) and UTC.
+			The returned image may extend beyond the requested start and end times by a few seconds.
+		size: The image size to render in form WIDTHxHEIGHT. Default 1024x64.
+	"""
+	start = dateutil.parse_utc_only(request.args['start'])
+	end = dateutil.parse_utc_only(request.args['end'])
+	if end <= start:
+		return "End must be after start", 400
+
+	size = request.args.get('size', '1024x64')
+	try:
+		width, height = map(int, size.split('x'))
+	except ValueError:
+		return "Invalid size", 400
+	if not ((0 < width <= 4096) and (0 < height <= 4096)):
+		return "Image size must be between 1x1 and 4096x4096", 400
+
+	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, start, end)
+	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(render_segments_waveform(segments, (width, height)), mimetype='image/png')
+
+
 @app.route('/generate_videos/<channel>/<quality>', methods=['POST'])
 @request_stats
 @has_path_args