Updating thumbnail generation to use database for templates and allow crop and

location to be varied

common/common/images.py

@@ -1,32 +1,20 @@
 
-import json
-import os
 import sys
 from io import BytesIO
 
 from PIL import Image
 
+
 """
-A template is two files:
-	NAME.png
-	NAME.json
-The image is the template image itself.
-The JSON file contains the following:
-	{
-		crop: BOX,
-		location: BOX,
-	}
-where BOX is a 4-tuple [left x, top y, right x, bottom y] describing a rectangle in image coordinates.
+A thumbnail template consists of a byte stream representation of a PNG template image along with two 4-tuples [left x, top y, right x, bottom y] describing the crop and location of the frame. 
 
 To create a thumbnail, the input frame is first cropped to the bounds of the "crop" box,
 then resized to the size of the "location" box, then pasted underneath the template image at that
 location within the template image.
 
-For example, a JSON file of:
-	{
-		"crop": [50, 100, 1870, 980],
-		"location": [320, 180, 1600, 900]
-	}
+For example
+crop = (50, 100, 1870, 980)
+location = (320, 180, 1600, 900)
 would crop the input frame from (50, 100) to (1870, 980), resize it to 720x1280,
 and place it at (320, 180).
 
@@ -34,18 +22,13 @@ If the original frame and the template differ in size, the frame is first resize
 This allows you to work with a consistent coordinate system regardless of the input frame size.
 """
 
-def compose_thumbnail_template(base_dir, template_name, frame_data):
-	template_path = os.path.join(base_dir, "thumbnail_templates", f"{template_name}.png")
-	info_path = os.path.join(base_dir, "thumbnail_templates", f"{template_name}.json")
+def compose_thumbnail_template(template_data, frame_data, crop, location):
 
-	template = Image.open(template_path)
 	# PIL can't load an image from a byte string directly, we have to pretend to be a file
+	template = Image.open(BytesIO(template_data))
 	frame = Image.open(BytesIO(frame_data))
 
-	with open(info_path) as f:
-		info = json.load(f)
-	crop = info['crop']
-	loc_left, loc_top, loc_right, loc_bottom = info['location']
+	loc_left, loc_top, loc_right, loc_bottom = location
 	location = loc_left, loc_top
 	location_size = loc_right - loc_left, loc_bottom - loc_top
 
@@ -72,11 +55,12 @@ def compose_thumbnail_template(base_dir, template_name, frame_data):
 	buf.seek(0)
 	return buf.read()
 
-
-def cli(template_name, image, base_dir="."):
-	with open(image, "rb") as f:
-		image = f.read()
-	thumbnail = compose_thumbnail_template(base_dir, template_name, image)
+def cli(template, frame, crop, location):
+	with open(template, "rb") as f:
+		template = f.read()
+	with open(frame, "rb") as f:
+		frame = f.read()
+	thumbnail = compose_thumbnail_template(template, frame, crop, location)
 	sys.stdout.buffer.write(thumbnail)
 
 

cutter/cutter/main.py

@@ -72,6 +72,8 @@ CUT_JOB_PARAMS = [
 	"thumbnail_time",
 	"thumbnail_template",
 	"thumbnail_image",
+	"thumbnail_crop",
+	"thumbnail_location",
 ]
 CutJob = namedtuple('CutJob', [
 	"id",
@@ -82,6 +84,24 @@ CutJob = namedtuple('CutJob', [
 	# params which map directly from DB columns
 ] + CUT_JOB_PARAMS)
 
+def get_template(dbmanager, name, crop, location):
+	"""Fetch the thumbnail template and any missing parameters from the database"""
+	with dbmanager.get_conn() as conn:
+		query = """
+			SELECT image, crop, location FROM templates WHERE name = %s
+		"""
+		results = database.query(conn, query, name)
+		row = results.fetchone()
+		if row is None:
+			raise ValueError('Template {} not found'.format(name))
+		row = row._asdict()
+		if not crop:
+			crop = row['crop']
+		if not location:
+			location = row['location']
+
+		return row['image'], crop, location
+
 
 def get_duration(job):
 	"""Get total video duration of a job, in seconds"""
@@ -456,7 +476,8 @@ class Cutter(object):
 			if job.thumbnail_mode == 'BARE':
 				image_data = frame
 			elif job.thumbnail_mode == 'TEMPLATE':
-				image_data = compose_thumbnail_template(self.segments_path, job.thumbnail_template, frame)
+				template, crop, location = get_template(self.dbmanager, job.thumbnail_template, job.thumbnail_crop, job.thumbnail_location)
+				image_data = compose_thumbnail_template(template, frame, crop, location)
 			else:
 				# shouldn't be able to happen given database constraints
 				assert False, "Bad thumbnail mode: {}".format(job.thumbnail_mode)
@@ -707,6 +728,8 @@ UPDATE_JOB_PARAMS = [
 	"thumbnail_time",
 	"thumbnail_template",
 	"thumbnail_image",
+	"thumbnail_crop",
+	"thumbnail_location",
 	"thumbnail_last_written",
 ]
 
@@ -766,7 +789,8 @@ class VideoUpdater(object):
 								if job.thumbnail_mode == 'BARE':
 									thumbnail_image = frame
 								elif job.thumbnail_mode == 'TEMPLATE':
-									thumbnail_image = compose_thumbnail_template(self.segments_path, job.thumbnail_template, frame)
+									template, crop, location = get_template(self.dbmanager, job.thumbnail_template, job.thumbnail_crop, job.thumbnail_location)
+									thumbnail_image = compose_thumbnail_template(template, frame, crop, location)
 								else:
 									assert False, "Bad thumbnail mode: {}".format(job.thumbnail_mode)
 								updates['thumbnail_image'] = thumbnail_image

restreamer/restreamer/main.py

@@ -8,6 +8,7 @@ import os
 import subprocess
 from uuid import uuid4
 
+import base64
 import gevent
 import gevent.backdoor
 import gevent.event
@@ -15,7 +16,7 @@ import prometheus_client as prom
 from flask import Flask, url_for, request, abort, Response
 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 import database, 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.images import compose_thumbnail_template
 from common.segments import smart_cut_segments, feed_input, render_segments_waveform, extract_frame, list_segment_files, get_best_segments_for_frame
@@ -174,20 +175,6 @@ def list_extras(dir):
 	return json.dumps(result)
 
 
-@app.route('/thumbnail-templates')
-def list_thumbnail_templates():
-	"""List available thumbnail templates. Returns a JSON list of names."""
-	path = os.path.join(
-		app.static_folder,
-		"thumbnail_templates",
-	)
-	return json.dumps([
-		os.path.splitext(filename)[0]
-		for filename in listdir(path)
-		if os.path.splitext(filename)[1] == ".png"
-	])
-
-
 def time_range_for_quality(channel, quality):
 	"""Returns earliest and latest times that the given quality has segments for
 	(up to hour resolution), or 404 if it doesn't exist / is empty."""
@@ -493,23 +480,43 @@ def get_frame(channel, quality):
 
 @app.route('/thumbnail/<channel>/<quality>.png')
 @request_stats
-@has_path_args
 def get_thumbnail(channel, quality):
 	"""
 	Returns a PNG image which is a preview of how a thumbnail will be generated.
-	Params:
-		timestamp: Required. The frame to use as the thumbnail image.
-			Must be in ISO 8601 format (ie. yyyy-mm-ddTHH:MM:SS) and UTC.
-		template: Required. The template name to use.
-			Must be one of the template names (without file extension) as returned
-			by GET /files/thumbnail_templates
 	"""
-	template_name = request.args['template']
-	template_path = os.path.join(app.static_folder, "thumbnail_templates", f"{template_name}.png")
-	if not os.path.exists(template_path):
-		return "No such template", 404
 
-	timestamp = dateutil.parse_utc_only(request.args['timestamp'])
+	template_params = request.json
+
+	if template_params['image']:
+		template = base64.b64decode(template_params['image'])
+
+		crop = template_params['crop']
+		location = template_params['location']
+
+	else:
+		if app.db_manager is None:
+			return 'A database connection is required to generate thumbnails', 501
+
+		with app.db_manager.get_conn() as conn:
+			query = """
+				SELECT image, crop, location FROM templates WHERE name = %s
+			"""
+			results = database.query(conn, query, template_params['name'])
+			row = results.fetchone()
+			if row is None:
+				return 'Template {} not found'.format(template_params['name']), 404
+			row = row._asdict()
+			template = row['image']
+			if not template_params['crop']:
+				crop = row['crop']
+			else:
+				crop = template_params['crop']
+			if not template_params['location']:
+				location = row['location']
+			else:
+				location = template_params['location']
+
+	timestamp = dateutil.parse_utc_only(template_params['timestamp'])
 
 	hours_path = os.path.join(app.static_folder, channel, quality)
 	if not os.path.isdir(hours_path):
@@ -520,7 +527,7 @@ def get_thumbnail(channel, quality):
 		return "We have no content available within the requested time range.", 406
 
 	frame = b''.join(extract_frame(segments, timestamp))
-	template = compose_thumbnail_template(app.static_folder, template_name, frame)
+	template = compose_thumbnail_template(template, frame, crop, location)
 
 	return Response(template, mimetype='image/png')
 
@@ -656,13 +663,18 @@ def generate_videos(channel, quality):
 	return ''
 
 
-def main(host='0.0.0.0', port=8000, base_dir='.', backdoor_port=0):
+def main(host='0.0.0.0', port=8000, base_dir='.', backdoor_port=0, connection_string=''):
 	app.static_folder = base_dir
 	server = WSGIServer((host, port), cors(app))
 
 	PromLogCountsHandler.install()
 	install_stacksampler()
 
+	if connection_string:
+		app.db_manager = database.DBManager(dsn=connection_string)
+	else:
+		app.db_manager = None
+
 	if backdoor_port:
 		gevent.backdoor.BackdoorServer(('127.0.0.1', backdoor_port), locals=locals()).start()