From 9dc371a1bc8283e9834af2cc69e1a67c8010dc63 Mon Sep 17 00:00:00 2001 From: TankNee Date: Mon, 4 Dec 2023 21:30:31 +0800 Subject: [PATCH] Refactor image and video classes to support parallel and ordered processing --- Katna/image_selector.py | 6 ++++-- Katna/video.py | 29 +++++++++++++++++++++++------ 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/Katna/image_selector.py b/Katna/image_selector.py index 09f40a6..a505e9e 100644 --- a/Katna/image_selector.py +++ b/Katna/image_selector.py @@ -30,9 +30,10 @@ class ImageSelector(object): :type object: class:`Object` """ - def __init__(self, n_processes=1): + def __init__(self, n_processes=1, ordered=False): # Setting number of processes for Multiprocessing Pool Object self.n_processes = n_processes + self.ordered = ordered # Setting for optimum Brightness values self.min_brightness_value = config.ImageSelector.min_brightness_value @@ -297,7 +298,8 @@ def select_best_frames(self, input_key_frames, number_of_frames): selected_images_index = self.__get_best_images_index_from_each_cluster__( filtered_key_frames, files_clusters_index_array ) - + if self.ordered: + selected_images_index = sorted(selected_images_index) for index in selected_images_index: img = filtered_key_frames[index] filtered_images_list.append(img) diff --git a/Katna/video.py b/Katna/video.py index 80b114d..d550c08 100644 --- a/Katna/video.py +++ b/Katna/video.py @@ -39,12 +39,24 @@ class Video(object): :type object: class:`Object` """ - def __init__(self, autoflip_build_path=None, autoflip_model_path=None): + def __init__(self, autoflip_build_path=None, autoflip_model_path=None, parallel: bool=True, ordered: bool=False): + """ + :param autoflip_build_path: path to the mediapipe autoflip build + :type autoflip_build_path: str + :param autoflip_model_path: path to the mediapipe autoflip model + :type autoflip_model_path: str + :param parallel: if True, then parallel processing will be used for keyframe extraction, defaults to True + :type parallel: bool, optional + :param ordered: if True, then keyframes will be extracted in order, defaults to False + :type ordered: bool, optional + """ # Find out location of ffmpeg binary on system helper._set_ffmpeg_binary_path() # If the duration of the clipped video is less than **min_video_duration** # then, the clip will be added with the previous clipped self._min_video_duration = config.Video.min_video_duration + self.parallel = parallel + self.ordered = ordered # Calculating optimum number of processes for multiprocessing self.n_processes = cpu_count() // 2 - 1 @@ -176,15 +188,20 @@ def _extract_keyframes_from_video(self, no_of_frames, file_path): # Passing all the clipped videos for the frame extraction using map function of the # multiprocessing pool - with self.pool_extractor: - extracted_candidate_frames = self.pool_extractor.map( - frame_extractor.extract_candidate_frames, chunked_videos - ) + if self.parallel: + with self.pool_extractor: + extracted_candidate_frames = self.pool_extractor.map( + frame_extractor.extract_candidate_frames, chunked_videos + ) + else: + extracted_candidate_frames = [] + for chunked_video in chunked_videos: + extracted_candidate_frames.append(frame_extractor.extract_candidate_frames(chunked_video)) # Converting the nested list of extracted frames into 1D list extracted_candidate_frames = functools.reduce(operator.iconcat, extracted_candidate_frames, []) self._remove_clips(chunked_videos) - image_selector = ImageSelector(self.n_processes) + image_selector = ImageSelector(self.n_processes, self.ordered) top_frames = image_selector.select_best_frames( extracted_candidate_frames, no_of_frames