From f814945dbd4bb988d0887caa69b1a474a371702f Mon Sep 17 00:00:00 2001 From: Christopher Usher Date: Sat, 19 Oct 2024 19:23:12 -0700 Subject: [PATCH] Updating thumbnail generation to use database for templates and allow crop and location to be varied --- common/common/images.py | 46 ++++++++--------------- cutter/cutter/main.py | 28 +++++++++++++- restreamer/restreamer/main.py | 70 ++++++++++++++++++++--------------- 3 files changed, 82 insertions(+), 62 deletions(-) diff --git a/common/common/images.py b/common/common/images.py index 0cd9eb2..32f2ecc 100644 --- a/common/common/images.py +++ b/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,14 +55,15 @@ 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) if __name__ == '__main__': import argh - argh.dispatch_command(cli) + argh.dispatch_command(cli) diff --git a/cutter/cutter/main.py b/cutter/cutter/main.py index 100f5de..aff1366 100644 --- a/cutter/cutter/main.py +++ b/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 diff --git a/restreamer/restreamer/main.py b/restreamer/restreamer/main.py index 69f665b..df34160 100644 --- a/restreamer/restreamer/main.py +++ b/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//.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()