From 67913b97378cbb76d572ed77b655bdf82cefdff4 Mon Sep 17 00:00:00 2001 From: Breakthrough Date: Sat, 9 Nov 2024 21:46:18 -0500 Subject: [PATCH 1/4] [save-images] Add multithreaded version of save-images This improves performance by over 50% in some cases. This should also fix #450 since we use the Path module for files now instead of OpenCV's imwrite. --- scenedetect/_cli/__init__.py | 2 +- scenedetect/_cli/commands.py | 2 +- scenedetect/scene_manager.py | 265 ++++++++++++++++++++++++++++++++++- 3 files changed, 264 insertions(+), 5 deletions(-) diff --git a/scenedetect/_cli/__init__.py b/scenedetect/_cli/__init__.py index 0f02bd6c..8298bc58 100644 --- a/scenedetect/_cli/__init__.py +++ b/scenedetect/_cli/__init__.py @@ -1493,7 +1493,7 @@ def save_images_command( "num_images": ctx.config.get_value("save-images", "num-images", num_images), "output_dir": output, "scale": scale, - "show_progress": ctx.quiet_mode, + "show_progress": not ctx.quiet_mode, "width": width, } ctx.add_command(cli_commands.save_images, save_images_args) diff --git a/scenedetect/_cli/commands.py b/scenedetect/_cli/commands.py index 857507d3..357e0b16 100644 --- a/scenedetect/_cli/commands.py +++ b/scenedetect/_cli/commands.py @@ -30,7 +30,7 @@ write_scene_list_html, ) from scenedetect.scene_manager import ( - save_images as save_images_impl, + save_images_mt as save_images_impl, ) from scenedetect.video_splitter import split_video_ffmpeg, split_video_mkvmerge diff --git a/scenedetect/scene_manager.py b/scenedetect/scene_manager.py index ccc8d0ab..cbad0350 100644 --- a/scenedetect/scene_manager.py +++ b/scenedetect/scene_manager.py @@ -87,7 +87,7 @@ def on_new_scene(frame_img: numpy.ndarray, frame_num: int): import threading from enum import Enum from pathlib import Path -from typing import Callable, Dict, Iterable, List, Optional, TextIO, Tuple, Union +from typing import Callable, Dict, List, Optional, TextIO, Tuple, Union import cv2 import numpy as np @@ -400,8 +400,39 @@ def write_scene_list_html( page.save(output_html_filename) +def _scale_image( + image: cv2.Mat, + height: int, + width: int, + scale: float, + aspect_ratio: float, + interpolation: Interpolation, +) -> cv2.Mat: + # TODO: Combine this resize with the ones below. + if aspect_ratio is not None: + image = cv2.resize( + image, (0, 0), fx=aspect_ratio, fy=1.0, interpolation=interpolation.value + ) + image_height = image.shape[0] + image_width = image.shape[1] + + # Figure out what kind of resizing needs to be done + if height or width: + if height and not width: + factor = height / float(image_height) + width = int(factor * image_width) + if width and not height: + factor = width / float(image_width) + height = int(factor * image_height) + assert height > 0 and width > 0 + image = cv2.resize(image, (width, height), interpolation=interpolation.value) + elif scale: + image = cv2.resize(image, (0, 0), fx=scale, fy=scale, interpolation=interpolation.value) + return image + + # -# TODO(v1.0): Consider moving all post-processing functionality into a separate submodule. +# TODO(v1.0): Move post-processing functions into separate submodule. def save_images( scene_list: SceneList, video: VideoStream, @@ -487,7 +518,9 @@ def save_images( # Setup flags and init progress bar if available. completed = True - logger.info(f"Saving {num_images} images per scene to {output_dir}, format {image_extension}") + logger.info( + f"Saving {num_images} images per scene [format={image_extension}] {output_dir if output_dir else ''} " + ) progress_bar = None if show_progress: progress_bar = tqdm(total=len(scene_list) * num_images, unit="images", dynamic_ncols=True) @@ -610,6 +643,232 @@ def save_images( return image_filenames +def save_images_mt( + scene_list: SceneList, + video: VideoStream, + num_images: int = 3, + frame_margin: int = 1, + image_extension: str = "jpg", + encoder_param: int = 95, + image_name_template: str = "$VIDEO_NAME-Scene-$SCENE_NUMBER-$IMAGE_NUMBER", + output_dir: Optional[str] = None, + show_progress: Optional[bool] = False, + scale: Optional[float] = None, + height: Optional[int] = None, + width: Optional[int] = None, + interpolation: Interpolation = Interpolation.CUBIC, + video_manager=None, +) -> Dict[int, List[str]]: + """Save a set number of images from each scene, given a list of scenes + and the associated video/frame source. + + Arguments: + scene_list: A list of scenes (pairs of FrameTimecode objects) returned + from calling a SceneManager's detect_scenes() method. + video: A VideoStream object corresponding to the scene list. + Note that the video will be closed/re-opened and seeked through. + num_images: Number of images to generate for each scene. Minimum is 1. + frame_margin: Number of frames to pad each scene around the beginning + and end (e.g. moves the first/last image into the scene by N frames). + Can set to 0, but will result in some video files failing to extract + the very last frame. + image_extension: Type of image to save (must be one of 'jpg', 'png', or 'webp'). + encoder_param: Quality/compression efficiency, based on type of image: + 'jpg' / 'webp': Quality 0-100, higher is better quality. 100 is lossless for webp. + 'png': Compression from 1-9, where 9 achieves best filesize but is slower to encode. + image_name_template: Template to use for naming image files. Can use the template variables + $VIDEO_NAME, $SCENE_NUMBER, $IMAGE_NUMBER, $TIMECODE, $FRAME_NUMBER, $TIMESTAMP_MS. + Should not include an extension. + output_dir: Directory to output the images into. If not set, the output + is created in the working directory. + show_progress: If True, shows a progress bar if tqdm is installed. + scale: Optional factor by which to rescale saved images. A scaling factor of 1 would + not result in rescaling. A value < 1 results in a smaller saved image, while a + value > 1 results in an image larger than the original. This value is ignored if + either the height or width values are specified. + height: Optional value for the height of the saved images. Specifying both the height + and width will resize images to an exact size, regardless of aspect ratio. + Specifying only height will rescale the image to that number of pixels in height + while preserving the aspect ratio. + width: Optional value for the width of the saved images. Specifying both the width + and height will resize images to an exact size, regardless of aspect ratio. + Specifying only width will rescale the image to that number of pixels wide + while preserving the aspect ratio. + interpolation: Type of interpolation to use when resizing images. + video_manager: [DEPRECATED] DO NOT USE. For backwards compatibility only. + + Returns: + Dictionary of the format { scene_num : [image_paths] }, where scene_num is the + number of the scene in scene_list (starting from 1), and image_paths is a list of + the paths to the newly saved/created images. + + Raises: + ValueError: Raised if any arguments are invalid or out of range (e.g. + if num_images is negative). + """ + # TODO(v0.7): Add DeprecationWarning that `video_manager` will be removed in v0.8. + if video_manager is not None: + logger.error("`video_manager` argument is deprecated, use `video` instead.") + video = video_manager + + if not scene_list: + return {} + if num_images <= 0 or frame_margin < 0: + raise ValueError() + + # TODO: Validate that encoder_param is within the proper range. + # Should be between 0 and 100 (inclusive) for jpg/webp, and 1-9 for png. + imwrite_param = ( + [get_cv2_imwrite_params()[image_extension], encoder_param] + if encoder_param is not None + else [] + ) + + video.reset() + + # Setup flags and init progress bar if available. + completed = True + logger.info( + f"Saving {num_images} images per scene [format={image_extension}] {output_dir if output_dir else ''} " + ) + progress_bar = None + if show_progress: + progress_bar = tqdm(total=len(scene_list) * num_images, unit="images", dynamic_ncols=True) + + filename_template = Template(image_name_template) + + scene_num_format = "%0" + scene_num_format += str(max(3, math.floor(math.log(len(scene_list), 10)) + 1)) + "d" + image_num_format = "%0" + image_num_format += str(math.floor(math.log(num_images, 10)) + 2) + "d" + + framerate = scene_list[0][0].framerate + + # TODO(v1.0): Split up into multiple sub-expressions so auto-formatter works correctly. + timecode_list = [ + [ + FrameTimecode(int(f), fps=framerate) + for f in [ + # middle frames + a[len(a) // 2] + if (0 < j < num_images - 1) or num_images == 1 + # first frame + else min(a[0] + frame_margin, a[-1]) + if j == 0 + # last frame + else max(a[-1] - frame_margin, a[0]) + # for each evenly-split array of frames in the scene list + for j, a in enumerate(np.array_split(r, num_images)) + ] + ] + for i, r in enumerate( + [ + # pad ranges to number of images + r if 1 + r[-1] - r[0] >= num_images else list(r) + [r[-1]] * (num_images - len(r)) + # create range of frames in scene + for r in ( + range( + start.get_frames(), + start.get_frames() + + max( + 1, # guard against zero length scenes + end.get_frames() - start.get_frames(), + ), + ) + # for each scene in scene list + for start, end in scene_list + ) + ] + ) + ] + + image_filenames = {i: [] for i in range(len(timecode_list))} + aspect_ratio = video.aspect_ratio + if abs(aspect_ratio - 1.0) < 0.01: + aspect_ratio = None + + logger.debug("Writing images with template %s", filename_template.template) + + MAX_QUEUED_ENCODE_FRAMES = 4 + MAX_QUEUED_SAVE_IMAGES = 4 + + encode_queue = queue.Queue(MAX_QUEUED_ENCODE_FRAMES) + save_queue = queue.Queue(MAX_QUEUED_SAVE_IMAGES) + + def image_encode_thread( + encode_queue: queue.Queue, save_queue: queue.Queue, image_extension: str + ): + while True: + frame_im, dest_path = encode_queue.get() + if frame_im is None: + return + frame_im = _scale_image(frame_im, height, width, scale, aspect_ratio, interpolation) + (is_ok, encoded) = cv2.imencode(f".{image_extension}", frame_im, imwrite_param) + if not is_ok: + continue + save_queue.put((encoded, dest_path)) + + def save_files_thread(save_queue: queue.Queue, progress_bar: tqdm): + while True: + encoded, dest_path = save_queue.get() + if encoded is None: + return + if encoded is not False: + encoded.tofile(Path(dest_path)) + if progress_bar is not None: + progress_bar.update(1) + + encode_thread = threading.Thread( + target=image_encode_thread, + args=(encode_queue, save_queue, image_extension), + daemon=True, + ) + save_thread = threading.Thread( + target=save_files_thread, + args=(save_queue, progress_bar), + daemon=True, + ) + encode_thread.start() + save_thread.start() + for i, scene_timecodes in enumerate(timecode_list): + for j, image_timecode in enumerate(scene_timecodes): + video.seek(image_timecode) + frame_im = video.read() + if frame_im is not None and frame_im is not False: + # TODO: Add extension to template. + # TODO: Allow NUM to be a valid suffix in addition to NUMBER. + 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), + FRAME_NUMBER=image_timecode.get_frames(), + TIMESTAMP_MS=int(image_timecode.get_seconds() * 1000), + TIMECODE=image_timecode.get_timecode().replace(":", ";"), + ), + image_extension, + ) + image_filenames[i].append(file_path) + encode_queue.put((frame_im, get_and_create_path(file_path, output_dir))) + else: + completed = False + break + + # *WARNING*: We do not handle errors or exceptions yet, and this can deadlock on errors! + encode_queue.put((None, None)) + save_queue.put((None, None)) + encode_thread.join() + save_thread.join() + + if progress_bar is not None: + progress_bar.close() + + if not completed: + logger.error("Could not generate all output images.") + + return image_filenames + + ## ## SceneManager Class Implementation ## From 99eee133f3e55ff67111110f1524881a4a6dff1d Mon Sep 17 00:00:00 2001 From: Breakthrough Date: Wed, 13 Nov 2024 21:23:20 -0500 Subject: [PATCH 2/4] [save-images] Add new ImageExtractor class The save_images function was getting quite complex and difficult to maintain, especially with the multithreaded version. This breaks it out into an object with smaller functions. The existing save_images function can be implemented using this new object. --- scenedetect/__init__.py | 359 ++++++++++++++++++++++++++++++++--- scenedetect/scene_manager.py | 220 ++------------------- 2 files changed, 347 insertions(+), 232 deletions(-) diff --git a/scenedetect/__init__.py b/scenedetect/__init__.py index 544be977..4f8a6836 100644 --- a/scenedetect/__init__.py +++ b/scenedetect/__init__.py @@ -15,43 +15,56 @@ :class:`SceneManager `. """ +import math +import queue +import threading +import typing as ty +from dataclasses import dataclass from logging import getLogger -from typing import List, Optional, Tuple, Union +from pathlib import Path +from string import Template # OpenCV is a required package, but we don't have it as an explicit dependency since we # need to support both opencv-python and opencv-python-headless. Include some additional # context with the exception if this is the case. try: - import cv2 as _ + import cv2 except ModuleNotFoundError as ex: raise ModuleNotFoundError( "OpenCV could not be found, try installing opencv-python:\n\npip install opencv-python", name="cv2", ) from ex +import numpy as np -# Commonly used classes/functions exported under the `scenedetect` namespace for brevity. -from scenedetect.platform import init_logger # noqa: I001 -from scenedetect.frame_timecode import FrameTimecode -from scenedetect.video_stream import VideoStream, VideoOpenFailure -from scenedetect.video_splitter import split_video_ffmpeg, split_video_mkvmerge -from scenedetect.scene_detector import SceneDetector -from scenedetect.detectors import ( - ContentDetector, - AdaptiveDetector, - ThresholdDetector, - HistogramDetector, - HashDetector, -) from scenedetect.backends import ( AVAILABLE_BACKENDS, - VideoStreamCv2, + VideoCaptureAdapter, VideoStreamAv, + VideoStreamCv2, VideoStreamMoviePy, - VideoCaptureAdapter, ) -from scenedetect.stats_manager import StatsManager, StatsFileCorrupt -from scenedetect.scene_manager import SceneManager, save_images +from scenedetect.detectors import ( + AdaptiveDetector, + ContentDetector, + HashDetector, + HistogramDetector, + ThresholdDetector, +) +from scenedetect.frame_timecode import FrameTimecode + +# Commonly used classes/functions exported under the `scenedetect` namespace for brevity. +from scenedetect.platform import ( # noqa: I001 + get_and_create_path, + get_cv2_imwrite_params, + init_logger, + tqdm, +) +from scenedetect.scene_detector import SceneDetector +from scenedetect.scene_manager import Interpolation, SceneList, SceneManager, save_images +from scenedetect.stats_manager import StatsFileCorrupt, StatsManager from scenedetect.video_manager import VideoManager # [DEPRECATED] DO NOT USE. +from scenedetect.video_splitter import split_video_ffmpeg, split_video_mkvmerge +from scenedetect.video_stream import VideoOpenFailure, VideoStream # Used for module identification and when printing version & about info # (e.g. calling `scenedetect version` or `scenedetect about`). @@ -63,7 +76,7 @@ def open_video( path: str, - framerate: Optional[float] = None, + framerate: ty.Optional[float] = None, backend: str = "opencv", **kwargs, ) -> VideoStream: @@ -117,12 +130,12 @@ def open_video( def detect( video_path: str, detector: SceneDetector, - stats_file_path: Optional[str] = None, + stats_file_path: ty.Optional[str] = None, show_progress: bool = False, - start_time: Optional[Union[str, float, int]] = None, - end_time: Optional[Union[str, float, int]] = None, + start_time: ty.Optional[ty.Union[str, float, int]] = None, + end_time: ty.Optional[ty.Union[str, float, int]] = None, start_in_scene: bool = False, -) -> List[Tuple[FrameTimecode, FrameTimecode]]: +) -> SceneList: """Perform scene detection on a given video `path` using the specified `detector`. Arguments: @@ -143,7 +156,7 @@ def detect( will always be included until the first fade-out event is detected. Returns: - List of scenes (pairs of :class:`FrameTimecode` objects). + List of scenes as pairs of (start, end) :class:`FrameTimecode` objects. Raises: :class:`VideoOpenFailure`: `video_path` could not be opened. @@ -169,3 +182,299 @@ def detect( if scene_manager.stats_manager is not None: scene_manager.stats_manager.save_to_csv(csv_file=stats_file_path) return scene_manager.get_scene_list(start_in_scene=start_in_scene) + + +# TODO: Just merge these variables into the extractor. +@dataclass +class ImageExtractorConfig: + num_images: int = 3 + """Number of images to generate for each scene. Minimum is 1.""" + frame_margin: int = 1 + """Number of frames to pad each scene around the beginning + and end (e.g. moves the first/last image into the scene by N frames). + Can set to 0, but will result in some video files failing to extract + the very last frame.""" + image_extension: str = "jpg" + """Type of image to save (must be one of 'jpg', 'png', or 'webp').""" + encoder_param: int = 95 + """Quality/compression efficiency, based on type of image: + 'jpg' / 'webp': Quality 0-100, higher is better quality. 100 is lossless for webp. + 'png': Compression from 1-9, where 9 achieves best filesize but is slower to encode.""" + image_name_template: str = "$VIDEO_NAME-Scene-$SCENE_NUMBER-$IMAGE_NUMBER" + """Template to use for naming image files. Can use the template variables + $VIDEO_NAME, $SCENE_NUMBER, $IMAGE_NUMBER, $TIMECODE, $FRAME_NUMBER, $TIMESTAMP_MS. + Should not include an extension.""" + scale: ty.Optional[float] = None + """Optional factor by which to rescale saved images. A scaling factor of 1 would + not result in rescaling. A value < 1 results in a smaller saved image, while a + value > 1 results in an image larger than the original. This value is ignored if + either the height or width values are specified.""" + height: ty.Optional[int] = None + """Optional value for the height of the saved images. Specifying both the height + and width will resize images to an exact size, regardless of aspect ratio. + Specifying only height will rescale the image to that number of pixels in height + while preserving the aspect ratio.""" + width: ty.Optional[int] = None + """Optional value for the width of the saved images. Specifying both the width + and height will resize images to an exact size, regardless of aspect ratio. + Specifying only width will rescale the image to that number of pixels wide + while preserving the aspect ratio.""" + interpolation: Interpolation = Interpolation.CUBIC + """Type of interpolation to use when resizing images.""" + + +class ImageExtractor: + def __init__( + self, + num_images: int = 3, + frame_margin: int = 1, + image_extension: str = "jpg", + encoder_param: int = 95, + image_name_template: str = "$VIDEO_NAME-Scene-$SCENE_NUMBER-$IMAGE_NUMBER", + scale: ty.Optional[float] = None, + height: ty.Optional[int] = None, + width: ty.Optional[int] = None, + interpolation: Interpolation = Interpolation.CUBIC, + ): + """Helper type to handle saving images for a set of scenes. This object is *not* thread-safe. + + Arguments: + num_images: Number of images to generate for each scene. Minimum is 1. + frame_margin: Number of frames to pad each scene around the beginning + and end (e.g. moves the first/last image into the scene by N frames). + Can set to 0, but will result in some video files failing to extract + the very last frame. + image_extension: Type of image to save (must be one of 'jpg', 'png', or 'webp'). + encoder_param: Quality/compression efficiency, based on type of image: + 'jpg' / 'webp': Quality 0-100, higher is better quality. 100 is lossless for webp. + 'png': Compression from 1-9, where 9 achieves best filesize but is slower to encode. + image_name_template: Template to use for output filanames. Can use template variables + $VIDEO_NAME, $SCENE_NUMBER, $IMAGE_NUMBER, $TIMECODE, $FRAME_NUMBER, $TIMESTAMP_MS. + *NOTE*: Should not include the image extension (set `image_extension` instead). + scale: Optional factor by which to rescale saved images. A scaling factor of 1 would + not result in rescaling. A value < 1 results in a smaller saved image, while a + value > 1 results in an image larger than the original. This value is ignored if + either the height or width values are specified. + height: Optional value for the height of the saved images. Specifying both the height + and width will resize images to an exact size, regardless of aspect ratio. + Specifying only height will rescale the image to that number of pixels in height + while preserving the aspect ratio. + width: Optional value for the width of the saved images. Specifying both the width + and height will resize images to an exact size, regardless of aspect ratio. + Specifying only width will rescale the image to that number of pixels wide + while preserving the aspect ratio. + interpolation: Type of interpolation to use when resizing images. + """ + self._num_images = num_images + self._frame_margin = frame_margin + self._image_extension = image_extension + self._encoder_param = encoder_param + self._image_name_template = image_name_template + self._scale = scale + self._height = height + self._width = width + self._interpolation = interpolation + + def run( + self, + video: VideoStream, + scene_list: SceneList, + output_dir: ty.Optional[str] = None, + show_progress=False, + ) -> ty.Dict[int, ty.List[str]]: + if not scene_list: + return {} + if self._num_images <= 0 or self._frame_margin < 0: + raise ValueError() + + video.reset() + + # Setup flags and init progress bar if available. + completed = True + logger.info( + f"Saving {self._num_images} images per scene [format={self._image_extension}] {output_dir if output_dir else ''} " + ) + progress_bar = None + if show_progress: + progress_bar = tqdm( + total=len(scene_list) * self._num_images, unit="images", dynamic_ncols=True + ) + + filename_template = Template(self._image_name_template) + scene_num_format = "%0" + scene_num_format += str(max(3, math.floor(math.log(len(scene_list), 10)) + 1)) + "d" + image_num_format = "%0" + image_num_format += str(math.floor(math.log(self._num_images, 10)) + 2) + "d" + + timecode_list = self.generate_timecode_list(scene_list) + image_filenames = {i: [] for i in range(len(timecode_list))} + logger.debug("Writing images with template %s", filename_template.template) + + MAX_QUEUED_ENCODE_FRAMES = 4 + MAX_QUEUED_SAVE_IMAGES = 4 + encode_queue = queue.Queue(MAX_QUEUED_ENCODE_FRAMES) + save_queue = queue.Queue(MAX_QUEUED_SAVE_IMAGES) + encode_thread = threading.Thread( + target=self._image_encode_thread, + args=(video, encode_queue, save_queue, self._image_extension), + daemon=True, + ) + save_thread = threading.Thread( + target=self._save_files_thread, + args=(save_queue, progress_bar), + daemon=True, + ) + encode_thread.start() + save_thread.start() + + for i, scene_timecodes in enumerate(timecode_list): + for j, image_timecode in enumerate(scene_timecodes): + video.seek(image_timecode) + frame_im = video.read() + if frame_im is not None and frame_im is not False: + # TODO: Add extension to template. + # TODO: Allow NUM to be a valid suffix in addition to NUMBER. + 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), + FRAME_NUMBER=image_timecode.get_frames(), + TIMESTAMP_MS=int(image_timecode.get_seconds() * 1000), + TIMECODE=image_timecode.get_timecode().replace(":", ";"), + ), + self._image_extension, + ) + image_filenames[i].append(file_path) + encode_queue.put((frame_im, get_and_create_path(file_path, output_dir))) + else: + completed = False + break + + # *WARNING*: We do not handle errors or exceptions yet, and this can deadlock on errors! + encode_queue.put((None, None)) + save_queue.put((None, None)) + encode_thread.join() + save_thread.join() + if progress_bar is not None: + progress_bar.close() + if not completed: + logger.error("Could not generate all output images.") + + return image_filenames + + def _image_encode_thread( + self, + video: VideoStream, + encode_queue: queue.Queue, + save_queue: queue.Queue, + image_extension: str, + ): + aspect_ratio = video.aspect_ratio + if abs(aspect_ratio - 1.0) < 0.01: + aspect_ratio = None + # TODO: Validate that encoder_param is within the proper range. + # Should be between 0 and 100 (inclusive) for jpg/webp, and 1-9 for png. + imwrite_param = ( + [get_cv2_imwrite_params()[self._image_extension], self._encoder_param] + if self._encoder_param is not None + else [] + ) + while True: + frame_im, dest_path = encode_queue.get() + if frame_im is None: + return + frame_im = self.resize_image( + frame_im, + aspect_ratio, + ) + (is_ok, encoded) = cv2.imencode(f".{image_extension}", frame_im, imwrite_param) + if not is_ok: + continue + save_queue.put((encoded, dest_path)) + + def _save_files_thread(self, save_queue: queue.Queue, progress_bar: tqdm): + while True: + encoded, dest_path = save_queue.get() + if encoded is None: + return + if encoded is not False: + encoded.tofile(Path(dest_path)) + if progress_bar is not None: + progress_bar.update(1) + + def generate_timecode_list(self, scene_list: SceneList) -> ty.List[ty.Iterable[FrameTimecode]]: + """Generates a list of timecodes for each scene in `scene_list` based on the current config + parameters.""" + framerate = scene_list[0][0].framerate + # TODO(v1.0): Split up into multiple sub-expressions so auto-formatter works correctly. + return [ + ( + FrameTimecode(int(f), fps=framerate) + 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._frame_margin, a[-1]) + if j == 0 + # last frame + else max(a[-1] - self._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 r in ( + # pad ranges to number of images + r + if 1 + r[-1] - r[0] >= self._num_images + else list(r) + [r[-1]] * (self._num_images - len(r)) + # create range of frames in scene + for r in ( + range( + start.get_frames(), + start.get_frames() + + max( + 1, # guard against zero length scenes + end.get_frames() - start.get_frames(), + ), + ) + # for each scene in scene list + for start, end in scene_list + ) + ) + ] + + def resize_image( + self, + image: cv2.Mat, + aspect_ratio: float, + ) -> cv2.Mat: + """Resizes the given `image` according to the current config parameters. `aspect_ratio` is + used to correct for non-square pixels.""" + # TODO: Combine this resize with the ones below. + if aspect_ratio is not None: + image = cv2.resize( + image, (0, 0), fx=aspect_ratio, fy=1.0, interpolation=self._interpolation.value + ) + image_height = image.shape[0] + image_width = image.shape[1] + # Figure out what kind of resizing needs to be done + if self._height or self._width: + if self._height and not self._width: + factor = self._height / float(image_height) + width = int(factor * image_width) + if self._width and not self._height: + factor = width / float(image_width) + height = int(factor * image_height) + assert height > 0 and width > 0 + image = cv2.resize(image, (width, height), interpolation=self._interpolation.value) + elif self._scale: + image = cv2.resize( + image, + (0, 0), + fx=self._scale, + fy=self._scale, + interpolation=self._interpolation.value, + ) + return image diff --git a/scenedetect/scene_manager.py b/scenedetect/scene_manager.py index cbad0350..a37793e5 100644 --- a/scenedetect/scene_manager.py +++ b/scenedetect/scene_manager.py @@ -659,214 +659,20 @@ def save_images_mt( interpolation: Interpolation = Interpolation.CUBIC, video_manager=None, ) -> Dict[int, List[str]]: - """Save a set number of images from each scene, given a list of scenes - and the associated video/frame source. - - Arguments: - scene_list: A list of scenes (pairs of FrameTimecode objects) returned - from calling a SceneManager's detect_scenes() method. - video: A VideoStream object corresponding to the scene list. - Note that the video will be closed/re-opened and seeked through. - num_images: Number of images to generate for each scene. Minimum is 1. - frame_margin: Number of frames to pad each scene around the beginning - and end (e.g. moves the first/last image into the scene by N frames). - Can set to 0, but will result in some video files failing to extract - the very last frame. - image_extension: Type of image to save (must be one of 'jpg', 'png', or 'webp'). - encoder_param: Quality/compression efficiency, based on type of image: - 'jpg' / 'webp': Quality 0-100, higher is better quality. 100 is lossless for webp. - 'png': Compression from 1-9, where 9 achieves best filesize but is slower to encode. - image_name_template: Template to use for naming image files. Can use the template variables - $VIDEO_NAME, $SCENE_NUMBER, $IMAGE_NUMBER, $TIMECODE, $FRAME_NUMBER, $TIMESTAMP_MS. - Should not include an extension. - output_dir: Directory to output the images into. If not set, the output - is created in the working directory. - show_progress: If True, shows a progress bar if tqdm is installed. - scale: Optional factor by which to rescale saved images. A scaling factor of 1 would - not result in rescaling. A value < 1 results in a smaller saved image, while a - value > 1 results in an image larger than the original. This value is ignored if - either the height or width values are specified. - height: Optional value for the height of the saved images. Specifying both the height - and width will resize images to an exact size, regardless of aspect ratio. - Specifying only height will rescale the image to that number of pixels in height - while preserving the aspect ratio. - width: Optional value for the width of the saved images. Specifying both the width - and height will resize images to an exact size, regardless of aspect ratio. - Specifying only width will rescale the image to that number of pixels wide - while preserving the aspect ratio. - interpolation: Type of interpolation to use when resizing images. - video_manager: [DEPRECATED] DO NOT USE. For backwards compatibility only. - - Returns: - Dictionary of the format { scene_num : [image_paths] }, where scene_num is the - number of the scene in scene_list (starting from 1), and image_paths is a list of - the paths to the newly saved/created images. - - Raises: - ValueError: Raised if any arguments are invalid or out of range (e.g. - if num_images is negative). - """ - # TODO(v0.7): Add DeprecationWarning that `video_manager` will be removed in v0.8. - if video_manager is not None: - logger.error("`video_manager` argument is deprecated, use `video` instead.") - video = video_manager - - if not scene_list: - return {} - if num_images <= 0 or frame_margin < 0: - raise ValueError() - - # TODO: Validate that encoder_param is within the proper range. - # Should be between 0 and 100 (inclusive) for jpg/webp, and 1-9 for png. - imwrite_param = ( - [get_cv2_imwrite_params()[image_extension], encoder_param] - if encoder_param is not None - else [] - ) - - video.reset() - - # Setup flags and init progress bar if available. - completed = True - logger.info( - f"Saving {num_images} images per scene [format={image_extension}] {output_dir if output_dir else ''} " - ) - progress_bar = None - if show_progress: - progress_bar = tqdm(total=len(scene_list) * num_images, unit="images", dynamic_ncols=True) - - filename_template = Template(image_name_template) - - scene_num_format = "%0" - scene_num_format += str(max(3, math.floor(math.log(len(scene_list), 10)) + 1)) + "d" - image_num_format = "%0" - image_num_format += str(math.floor(math.log(num_images, 10)) + 2) + "d" - - framerate = scene_list[0][0].framerate - - # TODO(v1.0): Split up into multiple sub-expressions so auto-formatter works correctly. - timecode_list = [ - [ - FrameTimecode(int(f), fps=framerate) - for f in [ - # middle frames - a[len(a) // 2] - if (0 < j < num_images - 1) or num_images == 1 - # first frame - else min(a[0] + frame_margin, a[-1]) - if j == 0 - # last frame - else max(a[-1] - frame_margin, a[0]) - # for each evenly-split array of frames in the scene list - for j, a in enumerate(np.array_split(r, num_images)) - ] - ] - for i, r in enumerate( - [ - # pad ranges to number of images - r if 1 + r[-1] - r[0] >= num_images else list(r) + [r[-1]] * (num_images - len(r)) - # create range of frames in scene - for r in ( - range( - start.get_frames(), - start.get_frames() - + max( - 1, # guard against zero length scenes - end.get_frames() - start.get_frames(), - ), - ) - # for each scene in scene list - for start, end in scene_list - ) - ] - ) - ] - - image_filenames = {i: [] for i in range(len(timecode_list))} - aspect_ratio = video.aspect_ratio - if abs(aspect_ratio - 1.0) < 0.01: - aspect_ratio = None - - logger.debug("Writing images with template %s", filename_template.template) - - MAX_QUEUED_ENCODE_FRAMES = 4 - MAX_QUEUED_SAVE_IMAGES = 4 - - encode_queue = queue.Queue(MAX_QUEUED_ENCODE_FRAMES) - save_queue = queue.Queue(MAX_QUEUED_SAVE_IMAGES) - - def image_encode_thread( - encode_queue: queue.Queue, save_queue: queue.Queue, image_extension: str - ): - while True: - frame_im, dest_path = encode_queue.get() - if frame_im is None: - return - frame_im = _scale_image(frame_im, height, width, scale, aspect_ratio, interpolation) - (is_ok, encoded) = cv2.imencode(f".{image_extension}", frame_im, imwrite_param) - if not is_ok: - continue - save_queue.put((encoded, dest_path)) - - def save_files_thread(save_queue: queue.Queue, progress_bar: tqdm): - while True: - encoded, dest_path = save_queue.get() - if encoded is None: - return - if encoded is not False: - encoded.tofile(Path(dest_path)) - if progress_bar is not None: - progress_bar.update(1) - - encode_thread = threading.Thread( - target=image_encode_thread, - args=(encode_queue, save_queue, image_extension), - daemon=True, - ) - save_thread = threading.Thread( - target=save_files_thread, - args=(save_queue, progress_bar), - daemon=True, + import scenedetect + + extractor = scenedetect.ImageExtractor( + num_images, + frame_margin, + image_extension, + encoder_param, + image_name_template, + scale, + height, + width, + interpolation, ) - encode_thread.start() - save_thread.start() - for i, scene_timecodes in enumerate(timecode_list): - for j, image_timecode in enumerate(scene_timecodes): - video.seek(image_timecode) - frame_im = video.read() - if frame_im is not None and frame_im is not False: - # TODO: Add extension to template. - # TODO: Allow NUM to be a valid suffix in addition to NUMBER. - 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), - FRAME_NUMBER=image_timecode.get_frames(), - TIMESTAMP_MS=int(image_timecode.get_seconds() * 1000), - TIMECODE=image_timecode.get_timecode().replace(":", ";"), - ), - image_extension, - ) - image_filenames[i].append(file_path) - encode_queue.put((frame_im, get_and_create_path(file_path, output_dir))) - else: - completed = False - break - - # *WARNING*: We do not handle errors or exceptions yet, and this can deadlock on errors! - encode_queue.put((None, None)) - save_queue.put((None, None)) - encode_thread.join() - save_thread.join() - - if progress_bar is not None: - progress_bar.close() - - if not completed: - logger.error("Could not generate all output images.") - - return image_filenames + return extractor.run(video, scene_list, output_dir, show_progress) ## From 55dec7f8590fec903442d874da588c82a779bc71 Mon Sep 17 00:00:00 2001 From: Breakthrough Date: Sat, 16 Nov 2024 00:52:00 -0500 Subject: [PATCH 3/4] [save-images] Make all threads exception-safe Ensure errors are re-raised safely from worker threads by using non-blocking puts and monitoring a common error queue. --- scenedetect/__init__.py | 302 -------------------------- scenedetect/_cli/commands.py | 4 +- scenedetect/scene_manager.py | 406 ++++++++++++++++++++++++++++------- tests/test_scene_manager.py | 4 +- 4 files changed, 337 insertions(+), 379 deletions(-) diff --git a/scenedetect/__init__.py b/scenedetect/__init__.py index 4f8a6836..4066b18e 100644 --- a/scenedetect/__init__.py +++ b/scenedetect/__init__.py @@ -15,14 +15,8 @@ :class:`SceneManager `. """ -import math -import queue -import threading import typing as ty -from dataclasses import dataclass from logging import getLogger -from pathlib import Path -from string import Template # OpenCV is a required package, but we don't have it as an explicit dependency since we # need to support both opencv-python and opencv-python-headless. Include some additional @@ -182,299 +176,3 @@ def detect( if scene_manager.stats_manager is not None: scene_manager.stats_manager.save_to_csv(csv_file=stats_file_path) return scene_manager.get_scene_list(start_in_scene=start_in_scene) - - -# TODO: Just merge these variables into the extractor. -@dataclass -class ImageExtractorConfig: - num_images: int = 3 - """Number of images to generate for each scene. Minimum is 1.""" - frame_margin: int = 1 - """Number of frames to pad each scene around the beginning - and end (e.g. moves the first/last image into the scene by N frames). - Can set to 0, but will result in some video files failing to extract - the very last frame.""" - image_extension: str = "jpg" - """Type of image to save (must be one of 'jpg', 'png', or 'webp').""" - encoder_param: int = 95 - """Quality/compression efficiency, based on type of image: - 'jpg' / 'webp': Quality 0-100, higher is better quality. 100 is lossless for webp. - 'png': Compression from 1-9, where 9 achieves best filesize but is slower to encode.""" - image_name_template: str = "$VIDEO_NAME-Scene-$SCENE_NUMBER-$IMAGE_NUMBER" - """Template to use for naming image files. Can use the template variables - $VIDEO_NAME, $SCENE_NUMBER, $IMAGE_NUMBER, $TIMECODE, $FRAME_NUMBER, $TIMESTAMP_MS. - Should not include an extension.""" - scale: ty.Optional[float] = None - """Optional factor by which to rescale saved images. A scaling factor of 1 would - not result in rescaling. A value < 1 results in a smaller saved image, while a - value > 1 results in an image larger than the original. This value is ignored if - either the height or width values are specified.""" - height: ty.Optional[int] = None - """Optional value for the height of the saved images. Specifying both the height - and width will resize images to an exact size, regardless of aspect ratio. - Specifying only height will rescale the image to that number of pixels in height - while preserving the aspect ratio.""" - width: ty.Optional[int] = None - """Optional value for the width of the saved images. Specifying both the width - and height will resize images to an exact size, regardless of aspect ratio. - Specifying only width will rescale the image to that number of pixels wide - while preserving the aspect ratio.""" - interpolation: Interpolation = Interpolation.CUBIC - """Type of interpolation to use when resizing images.""" - - -class ImageExtractor: - def __init__( - self, - num_images: int = 3, - frame_margin: int = 1, - image_extension: str = "jpg", - encoder_param: int = 95, - image_name_template: str = "$VIDEO_NAME-Scene-$SCENE_NUMBER-$IMAGE_NUMBER", - scale: ty.Optional[float] = None, - height: ty.Optional[int] = None, - width: ty.Optional[int] = None, - interpolation: Interpolation = Interpolation.CUBIC, - ): - """Helper type to handle saving images for a set of scenes. This object is *not* thread-safe. - - Arguments: - num_images: Number of images to generate for each scene. Minimum is 1. - frame_margin: Number of frames to pad each scene around the beginning - and end (e.g. moves the first/last image into the scene by N frames). - Can set to 0, but will result in some video files failing to extract - the very last frame. - image_extension: Type of image to save (must be one of 'jpg', 'png', or 'webp'). - encoder_param: Quality/compression efficiency, based on type of image: - 'jpg' / 'webp': Quality 0-100, higher is better quality. 100 is lossless for webp. - 'png': Compression from 1-9, where 9 achieves best filesize but is slower to encode. - image_name_template: Template to use for output filanames. Can use template variables - $VIDEO_NAME, $SCENE_NUMBER, $IMAGE_NUMBER, $TIMECODE, $FRAME_NUMBER, $TIMESTAMP_MS. - *NOTE*: Should not include the image extension (set `image_extension` instead). - scale: Optional factor by which to rescale saved images. A scaling factor of 1 would - not result in rescaling. A value < 1 results in a smaller saved image, while a - value > 1 results in an image larger than the original. This value is ignored if - either the height or width values are specified. - height: Optional value for the height of the saved images. Specifying both the height - and width will resize images to an exact size, regardless of aspect ratio. - Specifying only height will rescale the image to that number of pixels in height - while preserving the aspect ratio. - width: Optional value for the width of the saved images. Specifying both the width - and height will resize images to an exact size, regardless of aspect ratio. - Specifying only width will rescale the image to that number of pixels wide - while preserving the aspect ratio. - interpolation: Type of interpolation to use when resizing images. - """ - self._num_images = num_images - self._frame_margin = frame_margin - self._image_extension = image_extension - self._encoder_param = encoder_param - self._image_name_template = image_name_template - self._scale = scale - self._height = height - self._width = width - self._interpolation = interpolation - - def run( - self, - video: VideoStream, - scene_list: SceneList, - output_dir: ty.Optional[str] = None, - show_progress=False, - ) -> ty.Dict[int, ty.List[str]]: - if not scene_list: - return {} - if self._num_images <= 0 or self._frame_margin < 0: - raise ValueError() - - video.reset() - - # Setup flags and init progress bar if available. - completed = True - logger.info( - f"Saving {self._num_images} images per scene [format={self._image_extension}] {output_dir if output_dir else ''} " - ) - progress_bar = None - if show_progress: - progress_bar = tqdm( - total=len(scene_list) * self._num_images, unit="images", dynamic_ncols=True - ) - - filename_template = Template(self._image_name_template) - scene_num_format = "%0" - scene_num_format += str(max(3, math.floor(math.log(len(scene_list), 10)) + 1)) + "d" - image_num_format = "%0" - image_num_format += str(math.floor(math.log(self._num_images, 10)) + 2) + "d" - - timecode_list = self.generate_timecode_list(scene_list) - image_filenames = {i: [] for i in range(len(timecode_list))} - logger.debug("Writing images with template %s", filename_template.template) - - MAX_QUEUED_ENCODE_FRAMES = 4 - MAX_QUEUED_SAVE_IMAGES = 4 - encode_queue = queue.Queue(MAX_QUEUED_ENCODE_FRAMES) - save_queue = queue.Queue(MAX_QUEUED_SAVE_IMAGES) - encode_thread = threading.Thread( - target=self._image_encode_thread, - args=(video, encode_queue, save_queue, self._image_extension), - daemon=True, - ) - save_thread = threading.Thread( - target=self._save_files_thread, - args=(save_queue, progress_bar), - daemon=True, - ) - encode_thread.start() - save_thread.start() - - for i, scene_timecodes in enumerate(timecode_list): - for j, image_timecode in enumerate(scene_timecodes): - video.seek(image_timecode) - frame_im = video.read() - if frame_im is not None and frame_im is not False: - # TODO: Add extension to template. - # TODO: Allow NUM to be a valid suffix in addition to NUMBER. - 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), - FRAME_NUMBER=image_timecode.get_frames(), - TIMESTAMP_MS=int(image_timecode.get_seconds() * 1000), - TIMECODE=image_timecode.get_timecode().replace(":", ";"), - ), - self._image_extension, - ) - image_filenames[i].append(file_path) - encode_queue.put((frame_im, get_and_create_path(file_path, output_dir))) - else: - completed = False - break - - # *WARNING*: We do not handle errors or exceptions yet, and this can deadlock on errors! - encode_queue.put((None, None)) - save_queue.put((None, None)) - encode_thread.join() - save_thread.join() - if progress_bar is not None: - progress_bar.close() - if not completed: - logger.error("Could not generate all output images.") - - return image_filenames - - def _image_encode_thread( - self, - video: VideoStream, - encode_queue: queue.Queue, - save_queue: queue.Queue, - image_extension: str, - ): - aspect_ratio = video.aspect_ratio - if abs(aspect_ratio - 1.0) < 0.01: - aspect_ratio = None - # TODO: Validate that encoder_param is within the proper range. - # Should be between 0 and 100 (inclusive) for jpg/webp, and 1-9 for png. - imwrite_param = ( - [get_cv2_imwrite_params()[self._image_extension], self._encoder_param] - if self._encoder_param is not None - else [] - ) - while True: - frame_im, dest_path = encode_queue.get() - if frame_im is None: - return - frame_im = self.resize_image( - frame_im, - aspect_ratio, - ) - (is_ok, encoded) = cv2.imencode(f".{image_extension}", frame_im, imwrite_param) - if not is_ok: - continue - save_queue.put((encoded, dest_path)) - - def _save_files_thread(self, save_queue: queue.Queue, progress_bar: tqdm): - while True: - encoded, dest_path = save_queue.get() - if encoded is None: - return - if encoded is not False: - encoded.tofile(Path(dest_path)) - if progress_bar is not None: - progress_bar.update(1) - - def generate_timecode_list(self, scene_list: SceneList) -> ty.List[ty.Iterable[FrameTimecode]]: - """Generates a list of timecodes for each scene in `scene_list` based on the current config - parameters.""" - framerate = scene_list[0][0].framerate - # TODO(v1.0): Split up into multiple sub-expressions so auto-formatter works correctly. - return [ - ( - FrameTimecode(int(f), fps=framerate) - 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._frame_margin, a[-1]) - if j == 0 - # last frame - else max(a[-1] - self._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 r in ( - # pad ranges to number of images - r - if 1 + r[-1] - r[0] >= self._num_images - else list(r) + [r[-1]] * (self._num_images - len(r)) - # create range of frames in scene - for r in ( - range( - start.get_frames(), - start.get_frames() - + max( - 1, # guard against zero length scenes - end.get_frames() - start.get_frames(), - ), - ) - # for each scene in scene list - for start, end in scene_list - ) - ) - ] - - def resize_image( - self, - image: cv2.Mat, - aspect_ratio: float, - ) -> cv2.Mat: - """Resizes the given `image` according to the current config parameters. `aspect_ratio` is - used to correct for non-square pixels.""" - # TODO: Combine this resize with the ones below. - if aspect_ratio is not None: - image = cv2.resize( - image, (0, 0), fx=aspect_ratio, fy=1.0, interpolation=self._interpolation.value - ) - image_height = image.shape[0] - image_width = image.shape[1] - # Figure out what kind of resizing needs to be done - if self._height or self._width: - if self._height and not self._width: - factor = self._height / float(image_height) - width = int(factor * image_width) - if self._width and not self._height: - factor = width / float(image_width) - height = int(factor * image_height) - assert height > 0 and width > 0 - image = cv2.resize(image, (width, height), interpolation=self._interpolation.value) - elif self._scale: - image = cv2.resize( - image, - (0, 0), - fx=self._scale, - fy=self._scale, - interpolation=self._interpolation.value, - ) - return image diff --git a/scenedetect/_cli/commands.py b/scenedetect/_cli/commands.py index 357e0b16..74a65386 100644 --- a/scenedetect/_cli/commands.py +++ b/scenedetect/_cli/commands.py @@ -29,9 +29,7 @@ write_scene_list, write_scene_list_html, ) -from scenedetect.scene_manager import ( - save_images_mt as save_images_impl, -) +from scenedetect.scene_manager import save_images as save_images_impl from scenedetect.video_splitter import split_video_ffmpeg, split_video_mkvmerge logger = logging.getLogger("pyscenedetect") diff --git a/scenedetect/scene_manager.py b/scenedetect/scene_manager.py index a37793e5..67bb5fb7 100644 --- a/scenedetect/scene_manager.py +++ b/scenedetect/scene_manager.py @@ -85,9 +85,10 @@ def on_new_scene(frame_img: numpy.ndarray, frame_num: int): import queue import sys import threading +import typing as ty from enum import Enum from pathlib import Path -from typing import Callable, Dict, List, Optional, TextIO, Tuple, Union +from string import Template import cv2 import numpy as np @@ -100,17 +101,17 @@ def on_new_scene(frame_img: numpy.ndarray, frame_num: int): SimpleTableRow, ) from scenedetect.frame_timecode import FrameTimecode -from scenedetect.platform import Template, get_and_create_path, get_cv2_imwrite_params, tqdm +from scenedetect.platform import get_and_create_path, get_cv2_imwrite_params, tqdm from scenedetect.scene_detector import SceneDetector, SparseSceneDetector from scenedetect.stats_manager import StatsManager from scenedetect.video_stream import VideoStream logger = logging.getLogger("pyscenedetect") -SceneList = List[Tuple[FrameTimecode, FrameTimecode]] +SceneList = ty.List[ty.Tuple[FrameTimecode, FrameTimecode]] """Type hint for a list of scenes in the form (start time, end time).""" -CutList = List[FrameTimecode] +CutList = ty.List[FrameTimecode] """Type hint for a list of cuts, where each timecode represents the first frame of a new shot.""" # TODO: This value can and should be tuned for performance improvements as much as possible, @@ -166,9 +167,9 @@ def compute_downscale_factor(frame_width: int, effective_width: int = DEFAULT_MI def get_scenes_from_cuts( cut_list: CutList, - start_pos: Union[int, FrameTimecode], - end_pos: Union[int, FrameTimecode], - base_timecode: Optional[FrameTimecode] = None, + start_pos: ty.Union[int, FrameTimecode], + end_pos: ty.Union[int, FrameTimecode], + base_timecode: ty.Optional[FrameTimecode] = None, ) -> SceneList: """Returns a list of tuples of start/end FrameTimecodes for each scene based on a list of detected scene cuts/breaks. @@ -212,11 +213,14 @@ def get_scenes_from_cuts( return scene_list +# TODO(v1.0): Move post-processing functionality into separate submodule. + + def write_scene_list( - output_csv_file: TextIO, + output_csv_file: ty.TextIO, scene_list: SceneList, include_cut_list: bool = True, - cut_list: Optional[CutList] = None, + cut_list: ty.Optional[CutList] = None, col_separator: str = ",", row_separator: str = "\n", ): @@ -279,12 +283,12 @@ def write_scene_list( def write_scene_list_html( output_html_filename: str, scene_list: SceneList, - cut_list: Optional[CutList] = None, + cut_list: ty.Optional[CutList] = None, css: str = None, css_class: str = "mytable", - image_filenames: Optional[Dict[int, List[str]]] = None, - image_width: Optional[int] = None, - image_height: Optional[int] = None, + image_filenames: ty.Optional[ty.Dict[int, ty.List[str]]] = None, + image_width: ty.Optional[int] = None, + image_height: ty.Optional[int] = None, ): """Writes the given list of scenes to an output file handle in html format. @@ -402,10 +406,10 @@ def write_scene_list_html( def _scale_image( image: cv2.Mat, - height: int, - width: int, - scale: float, aspect_ratio: float, + height: ty.Optional[int], + width: ty.Optional[int], + scale: ty.Optional[float], interpolation: Interpolation, ) -> cv2.Mat: # TODO: Combine this resize with the ones below. @@ -431,8 +435,283 @@ def _scale_image( return image -# -# TODO(v1.0): Move post-processing functions into separate submodule. +class _ImageExtractor: + def __init__( + self, + num_images: int = 3, + frame_margin: int = 1, + image_extension: str = "jpg", + imwrite_param: ty.Dict[str, ty.Union[int, None]] = None, + image_name_template: str = "$VIDEO_NAME-Scene-$SCENE_NUMBER-$IMAGE_NUMBER", + scale: ty.Optional[float] = None, + height: ty.Optional[int] = None, + width: ty.Optional[int] = None, + interpolation: Interpolation = Interpolation.CUBIC, + ): + """Multi-threaded implementation of save-images functionality. Uses background threads to + handle image encoding and saving images to disk to improve parallelism. + + This object is thread-safe. + + Arguments: + num_images: Number of images to generate for each scene. Minimum is 1. + frame_margin: Number of frames to pad each scene around the beginning + and end (e.g. moves the first/last image into the scene by N frames). + Can set to 0, but will result in some video files failing to extract + the very last frame. + image_extension: Type of image to save (must be one of 'jpg', 'png', or 'webp'). + encoder_param: Quality/compression efficiency, based on type of image: + 'jpg' / 'webp': Quality 0-100, higher is better quality. 100 is lossless for webp. + 'png': Compression from 1-9, where 9 achieves best filesize but is slower to encode. + image_name_template: Template to use for output filanames. Can use template variables + $VIDEO_NAME, $SCENE_NUMBER, $IMAGE_NUMBER, $TIMECODE, $FRAME_NUMBER, $TIMESTAMP_MS. + *NOTE*: Should not include the image extension (set `image_extension` instead). + scale: Optional factor by which to rescale saved images. A scaling factor of 1 would + not result in rescaling. A value < 1 results in a smaller saved image, while a + value > 1 results in an image larger than the original. This value is ignored if + either the height or width values are specified. + height: Optional value for the height of the saved images. Specifying both the height + and width will resize images to an exact size, regardless of aspect ratio. + Specifying only height will rescale the image to that number of pixels in height + while preserving the aspect ratio. + width: Optional value for the width of the saved images. Specifying both the width + and height will resize images to an exact size, regardless of aspect ratio. + Specifying only width will rescale the image to that number of pixels wide + while preserving the aspect ratio. + interpolation: Type of interpolation to use when resizing images. + """ + self._num_images = num_images + self._frame_margin = frame_margin + self._image_extension = image_extension + self._image_name_template = image_name_template + self._scale = scale + self._height = height + self._width = width + self._interpolation = interpolation + self._imwrite_param = imwrite_param if imwrite_param else {} + + def run( + self, + video: VideoStream, + scene_list: SceneList, + output_dir: ty.Optional[str] = None, + show_progress=False, + ) -> ty.Dict[int, ty.List[str]]: + """Run image extraction on `video` using the current parameters. Thread-safe. + + Arguments: + video: The video to process. + scene_list: The scenes detected in the video. + output_dir: Directory to write files to. + show_progress: If `true` and tqdm is available, shows a progress bar. + """ + # Setup flags and init progress bar if available. + completed = True + logger.info( + f"Saving {self._num_images} images per scene [format={self._image_extension}] {output_dir if output_dir else ''} " + ) + progress_bar = None + if show_progress: + progress_bar = tqdm( + total=len(scene_list) * self._num_images, unit="images", dynamic_ncols=True + ) + + timecode_list = self.generate_timecode_list(scene_list) + image_filenames = {i: [] for i in range(len(timecode_list))} + + filename_template = Template(self._image_name_template) + logger.debug("Writing images with template %s", filename_template.template) + scene_num_format = "%0" + scene_num_format += str(max(3, math.floor(math.log(len(scene_list), 10)) + 1)) + "d" + image_num_format = "%0" + image_num_format += str(math.floor(math.log(self._num_images, 10)) + 2) + "d" + + def format_filename(scene_number: int, image_number: int, image_timecode: FrameTimecode): + return "%s.%s" % ( + filename_template.safe_substitute( + VIDEO_NAME=video.name, + SCENE_NUMBER=scene_num_format % (scene_number + 1), + IMAGE_NUMBER=image_num_format % (image_number + 1), + FRAME_NUMBER=image_timecode.get_frames(), + TIMESTAMP_MS=int(image_timecode.get_seconds() * 1000), + TIMECODE=image_timecode.get_timecode().replace(":", ";"), + ), + self._image_extension, + ) + + MAX_QUEUED_ENCODE_FRAMES = 4 + MAX_QUEUED_SAVE_IMAGES = 4 + encode_queue = queue.Queue(MAX_QUEUED_ENCODE_FRAMES) + save_queue = queue.Queue(MAX_QUEUED_SAVE_IMAGES) + error_queue = queue.Queue(2) # Queue size must be the same as the # of worker threads! + + def check_error_queue(): + try: + return error_queue.get(block=False) + except queue.Empty: + pass + return None + + def launch_thread(callable, *args, **kwargs): + def capture_errors(callable, *args, **kwargs): + try: + return callable(*args, **kwargs) + # Errors we capture in `error_queue` will be re-raised by this thread. + except: # noqa: E722 + error_queue.put(sys.exc_info()) + return None + + thread = threading.Thread( + target=capture_errors, + args=( + callable, + *args, + ), + kwargs=kwargs, + daemon=True, + ) + thread.start() + return thread + + def checked_put(work_queue: queue.Queue, item: ty.Any): + error = None + while True: + try: + work_queue.put(item, timeout=0.1) + return + except queue.Full: + error = check_error_queue() + if error is not None: + break + continue + raise error[1].with_traceback(error[2]) + + encode_thread = launch_thread( + self.image_encode_thread, + video, + encode_queue, + save_queue, + ) + save_thread = launch_thread(self.image_save_thread, save_queue, progress_bar) + + for i, scene_timecodes in enumerate(timecode_list): + for j, timecode in enumerate(scene_timecodes): + video.seek(timecode) + frame_im = video.read() + if frame_im is not None and frame_im is not False: + file_path = format_filename(i, j, timecode) + image_filenames[i].append(file_path) + checked_put( + encode_queue, (frame_im, get_and_create_path(file_path, output_dir)) + ) + else: + completed = False + break + + checked_put(encode_queue, (None, None)) + encode_thread.join() + checked_put(save_queue, (None, None)) + save_thread.join() + + error = check_error_queue() + if error is not None: + raise error[1].with_traceback(error[2]) + + if progress_bar is not None: + progress_bar.close() + if not completed: + logger.error("Could not generate all output images.") + + return image_filenames + + def image_encode_thread( + self, + video: VideoStream, + encode_queue: queue.Queue, + save_queue: queue.Queue, + ): + aspect_ratio = video.aspect_ratio + if abs(aspect_ratio - 1.0) < 0.01: + aspect_ratio = None + # TODO: Validate that encoder_param is within the proper range. + # Should be between 0 and 100 (inclusive) for jpg/webp, and 1-9 for png. + while True: + frame_im, dest_path = encode_queue.get() + if frame_im is None: + return + frame_im = self.resize_image( + frame_im, + aspect_ratio, + ) + (is_ok, encoded) = cv2.imencode( + f".{self._image_extension}", frame_im, self._imwrite_param + ) + if not is_ok: + continue + save_queue.put((encoded, dest_path)) + + def image_save_thread(self, save_queue: queue.Queue, progress_bar: tqdm): + while True: + encoded, dest_path = save_queue.get() + if encoded is None: + return + if encoded is not False: + encoded.tofile(Path(dest_path)) + if progress_bar is not None: + progress_bar.update(1) + + def generate_timecode_list(self, scene_list: SceneList) -> ty.List[ty.Iterable[FrameTimecode]]: + """Generates a list of timecodes for each scene in `scene_list` based on the current config + parameters.""" + framerate = scene_list[0][0].framerate + # TODO(v1.0): Split up into multiple sub-expressions so auto-formatter works correctly. + return [ + ( + FrameTimecode(int(f), fps=framerate) + 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._frame_margin, a[-1]) + if j == 0 + # last frame + else max(a[-1] - self._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 r in ( + # pad ranges to number of images + r + if 1 + r[-1] - r[0] >= self._num_images + else list(r) + [r[-1]] * (self._num_images - len(r)) + # create range of frames in scene + for r in ( + range( + start.get_frames(), + start.get_frames() + + max( + 1, # guard against zero length scenes + end.get_frames() - start.get_frames(), + ), + ) + # for each scene in scene list + for start, end in scene_list + ) + ) + ] + + def resize_image( + self, + image: cv2.Mat, + aspect_ratio: float, + ) -> cv2.Mat: + return _scale_image( + image, aspect_ratio, self._height, self._width, self._scale, self._interpolation + ) + + def save_images( scene_list: SceneList, video: VideoStream, @@ -441,14 +720,15 @@ def save_images( image_extension: str = "jpg", encoder_param: int = 95, image_name_template: str = "$VIDEO_NAME-Scene-$SCENE_NUMBER-$IMAGE_NUMBER", - output_dir: Optional[str] = None, - show_progress: Optional[bool] = False, - scale: Optional[float] = None, - height: Optional[int] = None, - width: Optional[int] = None, + output_dir: ty.Optional[str] = None, + show_progress: ty.Optional[bool] = False, + scale: ty.Optional[float] = None, + height: ty.Optional[int] = None, + width: ty.Optional[int] = None, interpolation: Interpolation = Interpolation.CUBIC, + threading: bool = True, video_manager=None, -) -> Dict[int, List[str]]: +) -> ty.Dict[int, ty.List[str]]: """Save a set number of images from each scene, given a list of scenes and the associated video/frame source. @@ -485,6 +765,7 @@ def save_images( Specifying only width will rescale the image to that number of pixels wide while preserving the aspect ratio. interpolation: Type of interpolation to use when resizing images. + threading: Offload image encoding and disk IO to background threads to improve performance. video_manager: [DEPRECATED] DO NOT USE. For backwards compatibility only. Returns: @@ -513,9 +794,22 @@ def save_images( if encoder_param is not None else [] ) - video.reset() + if threading: + extractor = _ImageExtractor( + num_images, + frame_margin, + image_extension, + imwrite_param, + image_name_template, + scale, + height, + width, + interpolation, + ) + return extractor.run(video, scene_list, output_dir, show_progress) + # Setup flags and init progress bar if available. completed = True logger.info( @@ -643,38 +937,6 @@ def save_images( return image_filenames -def save_images_mt( - scene_list: SceneList, - video: VideoStream, - num_images: int = 3, - frame_margin: int = 1, - image_extension: str = "jpg", - encoder_param: int = 95, - image_name_template: str = "$VIDEO_NAME-Scene-$SCENE_NUMBER-$IMAGE_NUMBER", - output_dir: Optional[str] = None, - show_progress: Optional[bool] = False, - scale: Optional[float] = None, - height: Optional[int] = None, - width: Optional[int] = None, - interpolation: Interpolation = Interpolation.CUBIC, - video_manager=None, -) -> Dict[int, List[str]]: - import scenedetect - - extractor = scenedetect.ImageExtractor( - num_images, - frame_margin, - image_extension, - encoder_param, - image_name_template, - scale, - height, - width, - interpolation, - ) - return extractor.run(video, scene_list, output_dir, show_progress) - - ## ## SceneManager Class Implementation ## @@ -688,7 +950,7 @@ class SceneManager: def __init__( self, - stats_manager: Optional[StatsManager] = None, + stats_manager: ty.Optional[StatsManager] = None, ): """ Arguments: @@ -697,7 +959,7 @@ def __init__( """ self._cutting_list = [] self._event_list = [] - self._detector_list: List[SceneDetector] = [] + self._detector_list: ty.List[SceneDetector] = [] self._sparse_detector_list = [] # TODO(v1.0): This class should own a StatsManager instead of taking an optional one. # Expose a new `stats_manager` @property from the SceneManager, and either change the @@ -706,16 +968,16 @@ def __init__( # TODO(v1.0): This class should own a VideoStream as well, instead of passing one # to the detect_scenes method. If concatenation is required, it can be implemented as # a generic VideoStream wrapper. - self._stats_manager: Optional[StatsManager] = stats_manager + self._stats_manager: ty.Optional[StatsManager] = stats_manager # Position of video that was first passed to detect_scenes. self._start_pos: FrameTimecode = None # Position of video on the last frame processed by detect_scenes. self._last_pos: FrameTimecode = None # Size of the decoded frames. - self._frame_size: Tuple[int, int] = None + self._frame_size: ty.Tuple[int, int] = None self._frame_size_errors: int = 0 - self._base_timecode: Optional[FrameTimecode] = None + self._base_timecode: ty.Optional[FrameTimecode] = None self._downscale: int = 1 self._auto_downscale: bool = True # Interpolation method to use when downscaling. Defaults to linear interpolation @@ -740,7 +1002,7 @@ def interpolation(self, value: Interpolation): self._interpolation = value @property - def stats_manager(self) -> Optional[StatsManager]: + def stats_manager(self) -> ty.Optional[StatsManager]: """Getter for the StatsManager associated with this SceneManager, if any.""" return self._stats_manager @@ -823,7 +1085,7 @@ def clear_detectors(self) -> None: self._sparse_detector_list.clear() def get_scene_list( - self, base_timecode: Optional[FrameTimecode] = None, start_in_scene: bool = False + self, base_timecode: ty.Optional[FrameTimecode] = None, start_in_scene: bool = False ) -> SceneList: """Return a list of tuples of start/end FrameTimecodes for each detected scene. @@ -855,7 +1117,7 @@ def get_scene_list( scene_list = [] return sorted(self._get_event_list() + scene_list) - def _get_cutting_list(self) -> List[int]: + def _get_cutting_list(self) -> ty.List[int]: """Return a sorted list of unique frame numbers of any detected scene cuts.""" if not self._cutting_list: return [] @@ -876,7 +1138,7 @@ def _process_frame( self, frame_num: int, frame_im: np.ndarray, - callback: Optional[Callable[[np.ndarray, int], None]] = None, + callback: ty.Optional[ty.Callable[[np.ndarray, int], None]] = None, ) -> bool: """Add any cuts detected with the current frame to the cutting list. Returns True if any new cuts were detected, False otherwise.""" @@ -917,12 +1179,12 @@ def stop(self) -> None: def detect_scenes( self, video: VideoStream = None, - duration: Optional[FrameTimecode] = None, - end_time: Optional[FrameTimecode] = None, + duration: ty.Optional[FrameTimecode] = None, + end_time: ty.Optional[FrameTimecode] = None, frame_skip: int = 0, show_progress: bool = False, - callback: Optional[Callable[[np.ndarray, int], None]] = None, - frame_source: Optional[VideoStream] = None, + callback: ty.Optional[ty.Callable[[np.ndarray, int], None]] = None, + frame_source: ty.Optional[VideoStream] = None, ) -> int: """Perform scene detection on the given video using the added SceneDetectors, returning the number of frames processed. Results can be obtained by calling :meth:`get_scene_list` or @@ -1150,7 +1412,7 @@ def _decode_thread( def get_cut_list( self, - base_timecode: Optional[FrameTimecode] = None, + base_timecode: ty.Optional[FrameTimecode] = None, show_warning: bool = True, ) -> CutList: """[DEPRECATED] Return a list of FrameTimecodes of the detected scene changes/cuts. @@ -1180,7 +1442,7 @@ def get_cut_list( logger.error("`get_cut_list()` is deprecated and will be removed in a future release.") return self._get_cutting_list() - def get_event_list(self, base_timecode: Optional[FrameTimecode] = None) -> SceneList: + def get_event_list(self, base_timecode: ty.Optional[FrameTimecode] = None) -> SceneList: """[DEPRECATED] DO NOT USE. Get a list of start/end timecodes of sparse detection events. diff --git a/tests/test_scene_manager.py b/tests/test_scene_manager.py index 16683bce..58593485 100644 --- a/tests/test_scene_manager.py +++ b/tests/test_scene_manager.py @@ -116,7 +116,7 @@ def test_save_images(test_video_file): total_images = 0 for scene_number in image_filenames: for path in image_filenames[scene_number]: - assert os.path.exists(path) + assert os.path.exists(path), f"expected {path} to exist" total_images += 1 assert total_images == len(glob.glob(image_name_glob)) @@ -151,7 +151,7 @@ def test_save_images_zero_width_scene(test_video_file): total_images = 0 for scene_number in image_filenames: for path in image_filenames[scene_number]: - assert os.path.exists(path) + assert os.path.exists(path), f"expected {path} to exist" total_images += 1 assert total_images == len(glob.glob(image_name_glob)) finally: From a9f012edc13f7e44780f14dba2f1e09bc23715ca Mon Sep 17 00:00:00 2001 From: Breakthrough Date: Sat, 23 Nov 2024 21:58:56 -0500 Subject: [PATCH 4/4] [save-images] Add new config option for setting threading mode --- scenedetect.cfg | 13 ++-- scenedetect/__init__.py | 44 +++++------- scenedetect/_cli/__init__.py | 1 + scenedetect/_cli/commands.py | 2 + scenedetect/_cli/config.py | 1 + scenedetect/scene_manager.py | 2 +- tests/test_scene_manager.py | 136 ++++++++++++++++++++++------------- website/pages/changelog.md | 9 ++- 8 files changed, 124 insertions(+), 84 deletions(-) diff --git a/scenedetect.cfg b/scenedetect.cfg index 205c614f..2cb1037b 100644 --- a/scenedetect.cfg +++ b/scenedetect.cfg @@ -220,22 +220,25 @@ # Image quality (jpeg/webp). Default is 95 for jpeg, 100 for webp #quality = 95 -# Compression amount for png images (0 to 9). Does not affect quality. +# Compression amount for png images (0 to 9). Only affects size, not quality. #compression = 3 -# Number of frames to skip at beginning/end of scene. +# Number of frames to ignore around each scene cut when selecting frames. #frame-margin = 1 -# Factor to resize images by (0.5 = half, 1.0 = same, 2.0 = double). +# Resize by scale factor (0.5 = half, 1.0 = same, 2.0 = double). #scale = 1.0 -# Override image height and/or width. Mutually exclusive with scale. +# Resize to specified height, width, or both. Mutually exclusive with scale. #height = 0 #width = 0 -# Method to use for image scaling (nearest, linear, cubic, area, lanczos4). +# Method to use for scaling (nearest, linear, cubic, area, lanczos4). #scale-method = linear +# Use separate threads for encoding and disk IO. Can improve performance. +#threading = yes + [export-html] # Filename format of created HTML file. Can use $VIDEO_NAME in the name. diff --git a/scenedetect/__init__.py b/scenedetect/__init__.py index 4066b18e..1463bc8f 100644 --- a/scenedetect/__init__.py +++ b/scenedetect/__init__.py @@ -22,43 +22,37 @@ # need to support both opencv-python and opencv-python-headless. Include some additional # context with the exception if this is the case. try: - import cv2 + import cv2 as _ except ModuleNotFoundError as ex: raise ModuleNotFoundError( "OpenCV could not be found, try installing opencv-python:\n\npip install opencv-python", name="cv2", ) from ex -import numpy as np -from scenedetect.backends import ( - AVAILABLE_BACKENDS, - VideoCaptureAdapter, - VideoStreamAv, - VideoStreamCv2, - VideoStreamMoviePy, -) +# Commonly used classes/functions exported under the `scenedetect` namespace for brevity. +# Note that order of importants is important! +from scenedetect.platform import init_logger # noqa: I001 +from scenedetect.frame_timecode import FrameTimecode +from scenedetect.video_stream import VideoStream, VideoOpenFailure +from scenedetect.video_splitter import split_video_ffmpeg, split_video_mkvmerge +from scenedetect.scene_detector import SceneDetector from scenedetect.detectors import ( - AdaptiveDetector, ContentDetector, - HashDetector, - HistogramDetector, + AdaptiveDetector, ThresholdDetector, + HistogramDetector, + HashDetector, ) -from scenedetect.frame_timecode import FrameTimecode - -# Commonly used classes/functions exported under the `scenedetect` namespace for brevity. -from scenedetect.platform import ( # noqa: I001 - get_and_create_path, - get_cv2_imwrite_params, - init_logger, - tqdm, +from scenedetect.backends import ( + AVAILABLE_BACKENDS, + VideoStreamCv2, + VideoStreamAv, + VideoStreamMoviePy, + VideoCaptureAdapter, ) -from scenedetect.scene_detector import SceneDetector -from scenedetect.scene_manager import Interpolation, SceneList, SceneManager, save_images -from scenedetect.stats_manager import StatsFileCorrupt, StatsManager +from scenedetect.stats_manager import StatsManager, StatsFileCorrupt +from scenedetect.scene_manager import SceneManager, save_images, SceneList, CutList, Interpolation from scenedetect.video_manager import VideoManager # [DEPRECATED] DO NOT USE. -from scenedetect.video_splitter import split_video_ffmpeg, split_video_mkvmerge -from scenedetect.video_stream import VideoOpenFailure, VideoStream # Used for module identification and when printing version & about info # (e.g. calling `scenedetect version` or `scenedetect about`). diff --git a/scenedetect/_cli/__init__.py b/scenedetect/_cli/__init__.py index 8298bc58..cfe5fe84 100644 --- a/scenedetect/_cli/__init__.py +++ b/scenedetect/_cli/__init__.py @@ -1494,6 +1494,7 @@ def save_images_command( "output_dir": output, "scale": scale, "show_progress": not ctx.quiet_mode, + "threading": ctx.config.get_value("save-images", "threading"), "width": width, } ctx.add_command(cli_commands.save_images, save_images_args) diff --git a/scenedetect/_cli/commands.py b/scenedetect/_cli/commands.py index 74a65386..9f6b5d32 100644 --- a/scenedetect/_cli/commands.py +++ b/scenedetect/_cli/commands.py @@ -177,6 +177,7 @@ def save_images( height: int, width: int, interpolation: Interpolation, + threading: bool, ): """Handles the `save-images` command.""" del cuts # save-images only uses scenes. @@ -195,6 +196,7 @@ def save_images( height=height, width=width, interpolation=interpolation, + threading=threading, ) # Save the result for use by `export-html` if required. context.save_images_result = (images, output_dir) diff --git a/scenedetect/_cli/config.py b/scenedetect/_cli/config.py index c53dcb7c..76327a62 100644 --- a/scenedetect/_cli/config.py +++ b/scenedetect/_cli/config.py @@ -361,6 +361,7 @@ def format(self, timecode: FrameTimecode) -> str: "quality": RangeValue(_PLACEHOLDER, min_val=0, max_val=100), "scale": 1.0, "scale-method": Interpolation.LINEAR, + "threading": True, "width": 0, }, "save-qp": { diff --git a/scenedetect/scene_manager.py b/scenedetect/scene_manager.py index 67bb5fb7..58cf6726 100644 --- a/scenedetect/scene_manager.py +++ b/scenedetect/scene_manager.py @@ -213,7 +213,7 @@ def get_scenes_from_cuts( return scene_list -# TODO(v1.0): Move post-processing functionality into separate submodule. +# TODO(#463): Move post-processing functionality into separate submodule. def write_scene_list( diff --git a/tests/test_scene_manager.py b/tests/test_scene_manager.py index 58593485..036c29e6 100644 --- a/tests/test_scene_manager.py +++ b/tests/test_scene_manager.py @@ -18,6 +18,7 @@ import glob import os import os.path +from pathlib import Path from typing import List from scenedetect.backends.opencv import VideoStreamCv2 @@ -84,7 +85,7 @@ def test_get_scene_list_start_in_scene(test_video_file): assert scene_list[0][1] == end_time -def test_save_images(test_video_file): +def test_save_images(test_video_file, tmp_path: Path): """Test scenedetect.scene_manager.save_images function.""" video = VideoStreamCv2(test_video_file) sm = SceneManager() @@ -97,66 +98,101 @@ def test_save_images(test_video_file): "$TIMESTAMP_MS.$TIMECODE" ) - try: - video_fps = video.frame_rate - scene_list = [ - (FrameTimecode(start, video_fps), FrameTimecode(end, video_fps)) - for start, end in [(0, 100), (200, 300), (300, 400)] - ] + video_fps = video.frame_rate + scene_list = [ + (FrameTimecode(start, video_fps), FrameTimecode(end, video_fps)) + for start, end in [(0, 100), (200, 300), (300, 400)] + ] + + image_filenames = save_images( + scene_list=scene_list, + output_dir=tmp_path, + video=video, + num_images=3, + image_extension="jpg", + image_name_template=image_name_template, + threading=False, + ) - image_filenames = save_images( - scene_list=scene_list, - video=video, - num_images=3, - image_extension="jpg", - image_name_template=image_name_template, - ) + # Ensure images got created, and the proper number got created. + total_images = 0 + for scene_number in image_filenames: + for path in image_filenames[scene_number]: + assert tmp_path.joinpath(path).exists(), f"expected {path} to exist" + total_images += 1 - # Ensure images got created, and the proper number got created. - total_images = 0 - for scene_number in image_filenames: - for path in image_filenames[scene_number]: - assert os.path.exists(path), f"expected {path} to exist" - total_images += 1 + assert total_images == len([path for path in tmp_path.glob(image_name_glob)]) - assert total_images == len(glob.glob(image_name_glob)) - finally: - for path in glob.glob(image_name_glob): - os.remove(path) +def test_save_images_singlethreaded(test_video_file, tmp_path: Path): + """Test scenedetect.scene_manager.save_images function.""" + video = VideoStreamCv2(test_video_file) + sm = SceneManager() + sm.add_detector(ContentDetector()) + + image_name_glob = "scenedetect.tempfile.*.jpg" + image_name_template = ( + "scenedetect.tempfile." + "$SCENE_NUMBER.$IMAGE_NUMBER.$FRAME_NUMBER." + "$TIMESTAMP_MS.$TIMECODE" + ) + + video_fps = video.frame_rate + scene_list = [ + (FrameTimecode(start, video_fps), FrameTimecode(end, video_fps)) + for start, end in [(0, 100), (200, 300), (300, 400)] + ] + + image_filenames = save_images( + scene_list=scene_list, + output_dir=tmp_path, + video=video, + num_images=3, + image_extension="jpg", + image_name_template=image_name_template, + threading=True, + ) + + # Ensure images got created, and the proper number got created. + total_images = 0 + for scene_number in image_filenames: + for path in image_filenames[scene_number]: + assert tmp_path.joinpath(path).exists(), f"expected {path} to exist" + total_images += 1 + + assert total_images == len([path for path in tmp_path.glob(image_name_glob)]) # TODO: Test other functionality against zero width scenes. -def test_save_images_zero_width_scene(test_video_file): +def test_save_images_zero_width_scene(test_video_file, tmp_path: Path): """Test scenedetect.scene_manager.save_images guards against zero width scenes.""" video = VideoStreamCv2(test_video_file) image_name_glob = "scenedetect.tempfile.*.jpg" image_name_template = "scenedetect.tempfile.$SCENE_NUMBER.$IMAGE_NUMBER" - try: - video_fps = video.frame_rate - scene_list = [ - (FrameTimecode(start, video_fps), FrameTimecode(end, video_fps)) - for start, end in [(0, 0), (1, 1), (2, 3)] - ] - NUM_IMAGES = 10 - image_filenames = save_images( - scene_list=scene_list, - video=video, - num_images=10, - image_extension="jpg", - image_name_template=image_name_template, - ) - assert len(image_filenames) == 3 - assert all(len(image_filenames[scene]) == NUM_IMAGES for scene in image_filenames) - total_images = 0 - for scene_number in image_filenames: - for path in image_filenames[scene_number]: - assert os.path.exists(path), f"expected {path} to exist" - total_images += 1 - assert total_images == len(glob.glob(image_name_glob)) - finally: - for path in glob.glob(image_name_glob): - os.remove(path) + + video_fps = video.frame_rate + scene_list = [ + (FrameTimecode(start, video_fps), FrameTimecode(end, video_fps)) + for start, end in [(0, 0), (1, 1), (2, 3)] + ] + NUM_IMAGES = 10 + image_filenames = save_images( + scene_list=scene_list, + output_dir=tmp_path, + video=video, + num_images=10, + image_extension="jpg", + image_name_template=image_name_template, + ) + assert len(image_filenames) == 3 + assert all(len(image_filenames[scene]) == NUM_IMAGES for scene in image_filenames) + total_images = 0 + for scene_number in image_filenames: + for path in image_filenames[scene_number]: + assert tmp_path.joinpath(path).exists(), f"expected {path} to exist" + total_images += 1 + + assert total_images == len([path for path in tmp_path.glob(image_name_glob)]) # TODO: This would be more readable if the callbacks were defined within the test case, e.g. diff --git a/website/pages/changelog.md b/website/pages/changelog.md index 9c24b2e5..ccbc2ada 100644 --- a/website/pages/changelog.md +++ b/website/pages/changelog.md @@ -588,15 +588,18 @@ Development - [feature] Add ability to configure CSV separators for rows/columns in config file [#423](https://github.com/Breakthrough/PySceneDetect/issues/423) - [feature] Add new `--show` flag to `export-html` command to launch browser after processing [#442](https://github.com/Breakthrough/PySceneDetect/issues/442) - [general] Timecodes of the form `MM:SS[.nnn]` are now processed correctly [#443](https://github.com/Breakthrough/PySceneDetect/issues/443) - - [bugfix] Fix `save-images`/`save_images()` not working correctly with UTF-8 paths [#450](https://github.com/Breakthrough/PySceneDetect/issues/455) + - [bugfix] Fix `save-images`/`save_images()` not working correctly with UTF-8 paths [#450](https://github.com/Breakthrough/PySceneDetect/issues/450) + - [improvement] Add new `threading` option to `save-images`/`save_images()` [#456](https://github.com/Breakthrough/PySceneDetect/issues/456) + - Enabled by default, offloads image encoding and disk IO to separate threads + - Improves performance by up to 50% in some cases - [bugfix] Fix crash when using `save-images`/`save_images()` with OpenCV backend [#455](https://github.com/Breakthrough/PySceneDetect/issues/455) - [bugfix] Fix new detectors not working with `default-detector` config option - [improvement] The `export-html` command now implicitly invokes `save-images` with default parameters - - The output of the `export-html` command will always use the result of the `save-images` command that *precedes* it + - The output of the `export-html` command will always use the result of the `save-images` command that *precedes* it - [general] Updates to Windows distributions: - The MoviePy backend is now included with Windows distributions - Bundled Python interpreter is now Python 3.13 - Updated PyAV 10 -> 13.1.0 and OpenCV 4.10.0.82 -> 4.10.0.84 - [improvement] `save_to_csv` now works with paths from `pathlib` - [api] The `save_to_csv` function now works correctly with paths from the `pathlib` module - - [api] Add `col_separator` and `row_separator` args to `write_scene_list` function in `scenedetect.scene_manager` + - [api] Add `col_separator` and `row_separator` args to `write_scene_list` function in `scenedetect.scene_manager` \ No newline at end of file