Merge pull request #12 from humphrem/output-dir

Add -o, --output-dir option for common output clips dir

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",

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,24 +167,26 @@ 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.
 
         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