diff --git a/scenedetect/cli/__init__.py b/scenedetect/cli/__init__.py index 664c57cf..f0ba85ff 100644 --- a/scenedetect/cli/__init__.py +++ b/scenedetect/cli/__init__.py @@ -702,12 +702,16 @@ def split_video_command(ctx, output, filename, high_quality, override_args, quie 'PNG compression rate, from 0-9. Higher values produce smaller files but result' ' in longer compression time. This setting does not affect image quality, only' ' file size. [default: 3]') +@click.option( + '--image-frame-margin', metavar='N', default=0, + type=click.INT, help= + 'Number of frames to ignore at the beginning and end of scenes when saving images') @click.pass_context -def save_images_command(ctx, output, filename, num_images, jpeg, webp, quality, png, compression): +def save_images_command(ctx, output, filename, num_images, jpeg, webp, quality, png, compression, image_frame_margin): """ Create images for each detected scene. """ if ctx.obj.save_images: duplicate_command(ctx, 'save-images') - ctx.obj.save_images_command(num_images, output, filename, jpeg, webp, quality, png, compression) + ctx.obj.save_images_command(num_images, output, filename, jpeg, webp, quality, png, compression, image_frame_margin) @@ -745,4 +749,3 @@ def colors_command(ctx): add_cli_command(scenedetect_cli, split_video_command) add_cli_command(scenedetect_cli, export_html_command) - diff --git a/scenedetect/cli/context.py b/scenedetect/cli/context.py index 305909da..aa01811a 100644 --- a/scenedetect/cli/context.py +++ b/scenedetect/cli/context.py @@ -38,6 +38,7 @@ import time import math from string import Template +import numpy as np # Third-Party Library Imports import click @@ -71,6 +72,8 @@ from scenedetect.platform import check_opencv_ffmpeg_dll +from scenedetect.frame_timecode import FrameTimecode + def get_plural(val_list): """ Get Plural: Helper function to return 's' if a list has more than one (1) element, otherwise returns ''. @@ -190,35 +193,43 @@ def _generate_images(self, scene_list, video_name, image_num_format += str(math.floor(math.log(self.num_images, 10)) + 2) + 'd' timecode_list = dict() - self.image_filenames = dict() - - for i in range(len(scene_list)): - timecode_list[i] = [] - self.image_filenames[i] = [] - - if self.num_images == 1: - for i, (start_time, end_time) in enumerate(scene_list): - duration = end_time - start_time - timecode_list[i].append(start_time + int(duration.get_frames() / 2)) - else: - middle_images = self.num_images - 2 - for i, (start_time, end_time) in enumerate(scene_list): - timecode_list[i].append(start_time) - - if middle_images > 0: - duration = (end_time.get_frames() - 1) - start_time.get_frames() - duration_increment = None - duration_increment = int(duration / (middle_images + 1)) - for j in range(middle_images): - timecode_list[i].append(start_time + ((j+1) * duration_increment)) - - # End FrameTimecode is always the same frame as the next scene's start_time - # (one frame past the end), so we need to subtract 1 here. - timecode_list[i].append(end_time - 1) - - for i in timecode_list: - for j, image_timecode in enumerate(timecode_list[i]): + fps = scene_list[0][0].framerate + + timecode_list = [ + [ + FrameTimecode(int(f), fps=fps) for f in [ + # middle frames + a[len(a)//2] if (0 < j < self.num_images-1) or self.num_images == 1 + + # first frame + else min(a[0] + self.image_frame_margin, a[-1]) if j == 0 + + # last frame + else max(a[-1] - self.image_frame_margin, a[0]) + + # for each evenly-split array of frames in the scene list + for j, a in enumerate(np.array_split(r, self.num_images)) + ] + ] + for i, r in enumerate([ + # pad ranges to number of images + r + if r.stop-r.start >= self.num_images + else list(r) + [r.stop-1] * (self.num_images - len(r)) + # create range of frames in scene + for r in ( + range(start.get_frames(), end.get_frames()) + # for each scene in scene list + for start, end in scene_list + ) + ]) + ] + + self.image_filenames = { i: [] for i in range(len(timecode_list)) } + + for i, tl in enumerate(timecode_list): + for j, image_timecode in enumerate(tl): self.video_manager.seek(image_timecode) self.video_manager.grab() ret_val, frame_im = self.video_manager.retrieve() @@ -226,7 +237,8 @@ def _generate_images(self, scene_list, video_name, file_path = '%s.%s' % (filename_template.safe_substitute( VIDEO_NAME=video_name, SCENE_NUMBER=scene_num_format % (i + 1), - IMAGE_NUMBER=image_num_format % (j + 1)), + IMAGE_NUMBER=image_num_format % (j + 1), + FRAME_NUMBER=image_timecode.get_frames()), self.image_extension) self.image_filenames[i].append(file_path) cv2.imwrite( @@ -641,7 +653,7 @@ def export_html_command(self, filename, no_images, image_width, image_height): def save_images_command(self, num_images, output, name_format, jpeg, webp, quality, - png, compression): + png, compression, image_frame_margin): # type: (int, str, str, bool, bool, int, bool, int) -> None """ Save Images Command: Parses all options/arguments passed to the save-images command, or with respect to the CLI, this function processes [save-images options] when calling: @@ -676,6 +688,7 @@ def save_images_command(self, num_images, output, name_format, jpeg, webp, quali self.image_param = compression if png else quality self.image_name_format = name_format self.num_images = num_images + self.image_frame_margin = image_frame_margin image_type = 'JPEG' if self.image_extension == 'jpg' else self.image_extension.upper() image_param_type = '' @@ -691,4 +704,3 @@ def save_images_command(self, num_images, output, name_format, jpeg, webp, quali logging.error('Multiple image type flags set for save-images command.') raise click.BadParameter( 'Only one image type (JPG/PNG/WEBP) can be specified.', param_hint='save-images') -