diff --git a/action.py b/action.py index d2a7ac0..31df9d4 100755 --- a/action.py +++ b/action.py @@ -9,7 +9,12 @@ import logging import argparse -from clip_manager import ClipManager, remove_clips_dir +from clip_manager import ( + ClipManager, + remove_clips_dir, + remove_output_dir, + create_output_dir, +) from yolo_fish_detector import YoloFishDetector from megadetector_detector import MegadetectorDetector from utils import * @@ -47,7 +52,15 @@ def load_detector(environment, logger): # Defining the function process_frames, called in main def process_frames( - video_path, cap, detector, clips, fps, total_frames, frames_to_skip, logger, args + video_path, + cap, + detector, + clips, + fps, + total_frames, + frames_to_skip, + logger, + args, ): """ Process frames from a video file and create clips based on detections. @@ -64,7 +77,7 @@ def process_frames( args (argparse.Namespace): The command line arguments. Returns: - clip_count (int): The number of clips created. + None """ confidence_threshold = args.confidence buffer_seconds = args.buffer @@ -82,7 +95,6 @@ def process_frames( detection_event = False frame_count = 0 - clip_count = 0 # Loop over all frames in the video while cap.isOpened(): @@ -126,11 +138,9 @@ def process_frames( logger.info( f"Detection period ended: {format_time(detection_end_time)} (duration={format_time(detection_end_time - detection_start_time)}, max confidence={format_percent(detection_highest_confidence)})" ) - clip_count += 1 clips.create_clip( detection_start_time, detection_end_time, - clip_count, video_path, ) @@ -178,14 +188,11 @@ def process_frames( logger.info( f"Detection period ended: {format_time(detection_end_time)} (duration={format_time(detection_end_time - detection_start_time)}, max confidence={format_percent(detection_highest_confidence)})" ) - clip_count += 1 clips.create_clip( detection_start_time, detection_end_time, - clip_count, video_path, ) - return clip_count # Main part of program to do setup and start processing frames in each file @@ -208,6 +215,7 @@ def main(args): buffer_seconds = args.buffer min_detection_duration = args.min_duration delete_clips = args.delete_clips + output_dir = args.output_dir environment = args.environment # Validate argument parameters from user before using them @@ -230,9 +238,17 @@ def main(args): cap = None clips = None + # Initialize the output_dir if specified + if output_dir: + # If `-d`` was specified, delete old clips first + if delete_clips: + remove_output_dir(output_dir, logger) + # Create the output directory if it doesn't exist + create_output_dir(output_dir) + try: # Create a queue manager for clips to be processed by ffmpeg - clips = ClipManager(logger) + clips = ClipManager(logger, output_dir) # Load YOLO-Fish or Megadetector, based on `-e` value detector = load_detector(environment, logger) @@ -249,8 +265,9 @@ def main(args): file_start_time = time.time() - # If the user requests it via -d flag, remove old clips first - if delete_clips: + # If the user requests it via -d flag, and isn't using a common output_dir + # remove old clips first + if not output_dir and delete_clips: remove_clips_dir(video_path, logger) # Setup video capture for this video file @@ -267,8 +284,14 @@ def main(args): f"Using confidence threshold {confidence_threshold}, minimum clip duration of {min_detection_duration} seconds, and {buffer_seconds} seconds of buffer." ) + # If we're not using a common clips dir, reset the counter for future clips + if not output_dir: + clips.reset_clip_count() + + clip_count_before = clips.get_clip_count() + # Process the video's frames into clips - clip_count = process_frames( + process_frames( video_path, cap, detector, @@ -279,19 +302,26 @@ def main(args): logger, args, ) + + clip_count_after = clips.get_clip_count() + clips_processed = clip_count_after - clip_count_before + file_end_time = time.time() logger.info( - f"Finished file {i} of {len(video_paths)}: {video_path} (total time to process file {format_time(file_end_time - file_start_time)}). Processed {total_frames} frames into {clip_count} clips" + f"Finished file {i} of {len(video_paths)}: {video_path} (total time to process file {format_time(file_end_time - file_start_time)}). Processed {total_frames} frames into {clips_processed} clips" ) # Clean-up the resources we have open, if necessary if cap is not None: cap.release() + cv2.destroyAllWindows() + cv2.waitKey(1) + # Keep track of total time to process all files, recording end time total_time_end = time.time() logger.info( - f"\nFinished processing, total time for {len(video_paths)} files: {format_time(total_time_end - total_time_start)}" + f"\nFinished. Total time for {len(video_paths)} files: {format_time(total_time_end - total_time_start)}" ) except KeyboardInterrupt: @@ -305,6 +335,7 @@ def main(args): cap.release() cv2.destroyAllWindows() + cv2.waitKey(1) # Wait for the ffmpeg clip queue to complete before we exit if clips is not None: @@ -371,6 +402,12 @@ def main(args): dest="delete_clips", help="Whether to delete previous clips before processing video", ) + parser.add_argument( + "-o", + "--output-dir", + dest="output_dir", + help="Output directory to use for all clips", + ) parser.add_argument( "-s", "--show-detections", diff --git a/clip_manager.py b/clip_manager.py index d09de7f..e051c45 100644 --- a/clip_manager.py +++ b/clip_manager.py @@ -32,23 +32,53 @@ def get_clips_dir(video_path): def remove_clips_dir(video_path, logger): """ Delete the old clips directory for the given video_path, - if exists. Args: video_path (str): The path to the video file. logger (Logger): The logger object. + Returns: + None + """ + clips_dir = get_clips_dir(video_path) + remove_output_dir(clips_dir, logger) + + +def remove_output_dir(output_dir, logger): + """ + Delete a common output dir for video clips. + + Args: + output_dir (str): The path to remove. + logger (Logger): The logger object. + Returns: None """ try: - clips_dir = get_clips_dir(video_path) - logger.debug(f"Removing old clips from {clips_dir}") - shutil.rmtree(clips_dir) + logger.debug(f"Removing old clips from {output_dir}") + shutil.rmtree(output_dir, ignore_errors=True) except Exception as e: logger.warning(f"Unable to remove old clips: {e}") +def create_output_dir(output_dir): + """ + Make sure the output dir exists + + Args: + output_dir (str): The path to create. + + Returns: + None + """ + if os.path.exists(output_dir): + if not os.path.isdir(output_dir): + raise ValueError(f"{output_dir} is a file, not a directory.") + else: + os.makedirs(output_dir, exist_ok=True) + + class ClipManager: """ A class used to manage the creation of video clips. Video clips @@ -57,26 +87,31 @@ class ClipManager: Attributes: logger (Logger): The logger object. + output_dir (string): Optional base path for storing clips clip_queue (Queue): The queue of clips. stop_event (Event): The stop event for the clip process. clip_process (Process): The clip process. + clip_count (int): The current clip count. """ - def __init__(self, logger): + def __init__(self, logger, output_dir): """ The constructor for the ClipManager class. It manages the queue and process for ffmpeg. Args: logger (Logger): The logger object. + output_dir (string): Optional base path for storing clips """ self.logger = logger + self.output_dir = output_dir self.clip_queue = Queue() self.stop_event = Event() self.clip_process = Process( target=self.create_clip_process, args=(self.clip_queue, self.stop_event) ) self.clip_process.start() + self.clip_count = 0 def create_clip_process(self, queue, stop_event): """ @@ -101,8 +136,12 @@ def create_clip_process(self, queue, stop_event): # Create a clip for the given detection period with ffmpeg clip_duration = clip_end_time - clip_start_time - clip_filename = f"{get_clips_dir(video_path)}/{(clip_count):03}-{format_time(clip_start_time, '_')}-{format_time(clip_end_time, '_')}.mp4" - os.makedirs(os.path.dirname(clip_filename), exist_ok=True) + base_dir = ( + self.output_dir if self.output_dir else get_clips_dir(video_path) + ) + clip_filename = f"{base_dir}/{(clip_count):04}-{format_time(clip_start_time, '_')}-{format_time(clip_end_time, '_')}.mp4" + create_output_dir(os.path.dirname(clip_filename)) + subprocess.run( [ "ffmpeg", @@ -128,7 +167,7 @@ def create_clip_process(self, queue, stop_event): except Empty: continue - def create_clip(self, clip_start_time, clip_end_time, clip_count, video_path): + def create_clip(self, clip_start_time, clip_end_time, video_path): """ Add a clip request to the queue, which will eventually create a new clip file. @@ -136,16 +175,18 @@ def create_clip(self, clip_start_time, clip_end_time, clip_count, video_path): Args: clip_start_time (float): The start time of the clip. clip_end_time (float): The end time of the clip. - clip_count (int): The count of the clip. video_path (str): The path to the video file. Returns: None """ + self.clip_count += 1 self.logger.debug( - f"Creating clip {clip_count}: start={format_time(clip_start_time)}, end={format_time(clip_end_time)}" + f"Creating clip {self.clip_count}: start={format_time(clip_start_time)}, end={format_time(clip_end_time)}" + ) + self.clip_queue.put( + (clip_start_time, clip_end_time, self.clip_count, video_path) ) - self.clip_queue.put((clip_start_time, clip_end_time, clip_count, video_path)) def stop(self): """ @@ -169,3 +210,22 @@ def cleanup(self): self.logger.debug("Waiting for remaining clips to be saved") self.clip_queue.put((None, None, None, None)) self.clip_process.join() + + def reset_clip_count(self): + """ + Reset the clip count to 0. + + Returns: + None + """ + self.logger.debug("Resetting clip manager clip_count to 0") + self.clip_count = 0 + + def get_clip_count(self): + """ + Get the current clip count. + + Returns: + int: The current clip count. + """ + return self.clip_count