megadetector 5.0.8__py3-none-any.whl → 5.0.9__py3-none-any.whl

detection/process_video.py

@@ -1,11 +1,20 @@
-########
-#
-# process_video.py
-#
-# Split a video (or folder of videos) into frames, run the frames through run_detector_batch.py,
-# and optionally stitch together results into a new video with detection boxes.
-#
-########
+"""
+
+process_video.py
+
+Splits a video (or folder of videos) into frames, runs the frames through run_detector_batch.py,
+and optionally stitches together results into a new video with detection boxes.
+
+Operates by separating the video into frames, typically sampling every Nth frame, and writing
+those frames to disk, before running MD.  This approach clearly has a downside: it requires
+a bunch more disk space, compared to extracting frames and running MD on them without ever
+writing them to disk.  The upside, though, is that this approach allows you to run repeat
+detection elimination after running MegaDetector, and it allows allows more efficient re-use
+of frames if you end up running MD more than once, or running multiple versions of MD.
+
+TODO: optionally skip writing frames to disk, and process frames in memory.
+
+"""
 
 #%% Imports
 
@@ -29,73 +38,113 @@ from uuid import uuid1
 from detection.video_utils import default_fourcc
 
 
-#%% Options classes
+#%% Classes
 
 class ProcessVideoOptions:
-
-    # Can be a model filename (.pt or .pb) or a model name (e.g. "MDV5A")
+    """
+    Options controlling the behavior of process_video()
+    """
+    
+    #: Can be a model filename (.pt or .pb) or a model name (e.g. "MDV5A")
     model_file = 'MDV5A'
     
-    # Can be a file or a folder
+    #: Video (of folder of videos) to process
     input_video_file = ''
 
+    #: .json file to which we should write results
     output_json_file = None
     
-    # Only relevant if render_output_video is True
+    #: File to which we should write a video with boxes, only relevant if 
+    #: render_output_video is True
     output_video_file = None
     
-    # Folder to use for extracted frames
+    #: Folder to use for extracted frames; will use a folder in system temp space
+    #: if this is None
     frame_folder = None
     
-    # Folder to use for rendered frames (if rendering output video)
+    # Folder to use for rendered frames (if rendering output video); will use a folder 
+    #: in system temp space if this is None
     frame_rendering_folder = None
     
-    # Should we render a video with detection boxes?
-    #
-    # Only supported when processing a single video, not a folder.
+    #: Should we render a video with detection boxes?
+    #:
+    #: Only supported when processing a single video, not a folder.
     render_output_video = False
     
-    # If we are rendering boxes to a new video, should we keep the temporary
-    # rendered frames?
+    #: If we are rendering boxes to a new video, should we keep the temporary
+    #: rendered frames?
     keep_rendered_frames = False
     
-    # Should we keep the extracted frames?
+    #: Should we keep the extracted frames?
     keep_extracted_frames = False
     
-    # Should we delete the entire folder the extracted frames are written to?
-    #
-    # By default, we delete the frame files but leave the (probably-empty) folder in place.
+    #: Should we delete the entire folder the extracted frames are written to?
+    #:
+    #: By default, we delete the frame files but leave the (probably-empty) folder in place, 
+    #: for no reason other than being paranoid about deleting folders.
     force_extracted_frame_folder_deletion = False
     
-    # Should we delete the entire folder the rendered frames are written to?
-    #
-    # By default, we delete the frame files but leave the (probably-empty) folder in place.
+    #: Should we delete the entire folder the rendered frames are written to?
+    #:
+    #: By default, we delete the frame files but leave the (probably-empty) folder in place,
+    #: for no reason other than being paranoid about deleting folders.
     force_rendered_frame_folder_deletion = False
-        
+     
+    #: If we've already run MegaDetector on this video or folder of videos, i.e. if we 
+    #: find a corresponding MD results file, should we re-use it?  Defaults to reprocessing.
     reuse_results_if_available = False
+    
+    #: If we've already split this video or folder of videos into frames, should we 
+    #: we re-use those extracted frames?  Defaults to reprocessing.
     reuse_frames_if_available = False
     
+    #: If [input_video_file] is a folder, should we search for videos recursively?
     recursive = False 
+    
+    #: Enable additional debug console output
     verbose = False
     
+    #: fourcc code to use for writing videos; only relevant if render_output_video is True
     fourcc = None
 
+    #: Confidence threshold to use for writing videos with boxes, only relevant if
+    #: if render_output_video is True.  Defaults to choosing a reasonable threshold
+    #: based on the model version.
     rendering_confidence_threshold = None
+    
+    #: Detections below this threshold will not be included in the output file.
     json_confidence_threshold = 0.005
+    
+    #: Sample every Nth frame; set to None (default) or 1 to sample every frame.  Typically
+    #: we sample down to around 3 fps, so for typical 30 fps videos, frame_sample=10 is a 
+    #: typical value.
     frame_sample = None
     
+    #: Number of workers to use for parallelization; set to <= 1 to disable parallelization
     n_cores = 1
 
+    #: For debugging only, stop processing after a certain number of frames.
     debug_max_frames = -1
     
+    #: File containing non-standard categories, typically only used if you're running a non-MD
+    #: detector.
     class_mapping_filename = None
 
+# ...class ProcessVideoOptions
+
 
 #%% Functions
 
 def process_video(options):
     """
-    Process a single video
+    Process a single video through MD, optionally writing a new video with boxes
+    
+    Args: 
+        options (ProcessVideoOptions): all the parameters used to control this process,
+            including filenames; see ProcessVideoOptions for details
+            
+    Returns:
+        dict: frame-level MegaDetector results, identical to what's in the output .json file
     """
 
     if options.output_json_file is None:
@@ -229,7 +278,11 @@ def process_video(options):
 
 def process_video_folder(options):
     """
-    Process a folder of videos    
+    Process a folder of videos through MD
+    
+    Args: 
+        options (ProcessVideoOptions): all the parameters used to control this process,
+            including filenames; see ProcessVideoOptions for details            
     """
     
     ## Validate options
@@ -428,8 +481,7 @@ def process_video_folder(options):
             print('Warning: error deleting frames from folder {}:\n{}'.format(
                 frame_output_folder,str(e)))
             pass
-        
-    
+
 # ...process_video_folder()
 
 
@@ -547,7 +599,7 @@ def main():
     default_options = ProcessVideoOptions()
     
     parser = argparse.ArgumentParser(description=(
-        'Run MegaDetector on each frame in a video (or every Nth frame), optionally '\
+        'Run MegaDetector on each frame (or every Nth frame) in a video (or folder of videos), optionally '\
         'producing a new video with detections annotated'))
 
     parser.add_argument('model_file', type=str,

detection/pytorch_detector.py

@@ -1,12 +1,12 @@
-########
-#
-# pytorch_detector.py
-#
-# Module to run MegaDetector v5, a PyTorch YOLOv5 animal detection model.
-#
-########
+"""
+
+pytorch_detector.py
 
-#%% Imports
+Module to run MegaDetector v5, a PyTorch YOLOv5 animal detection model.
+
+"""
+
+#%% Imports and constants
 
 import torch
 import numpy as np
@@ -104,12 +104,19 @@ print(f'Using PyTorch version {torch.__version__}')
 
 class PTDetector:
 
-    IMAGE_SIZE = 1280  # image size used in training
+    #: Image size passed to YOLOv5's letterbox() function; 1280 means "1280 on the long side, preserving 
+    #: aspect ratio"
+    #:
+    #: :meta private:
+    IMAGE_SIZE = 1280
+    
+    #: Stride size passed to YOLOv5's letterbox() function
+    #:
+    #: :meta private:
     STRIDE = 64
 
-    def __init__(self, model_path: str, 
-                 force_cpu: bool = False,
-                 use_model_native_classes: bool = False):
+    def __init__(self, model_path, force_cpu=False, use_model_native_classes= False):
+        
         self.device = 'cpu'
         if not force_cpu:
             if torch.cuda.is_available():
@@ -162,21 +169,26 @@ class PTDetector:
                                       detection_threshold=0.00001, image_size=None,
                                       skip_image_resizing=False):
         """
-        Apply the detector to an image.
+        Applies the detector to an image.
 
         Args:
-            img_original: the PIL Image object with EXIF rotation taken into account
-            image_id: a path to identify the image; will be in the "file" field of the output object
-            detection_threshold: confidence above which to include the detection proposal
-            skip_image_resizing: whether to skip internal image resizing and rely on external resizing
+            img_original (Image): the PIL Image object with EXIF rotation taken into account
+            image_id (str, optional): a path to identify the image; will be in the "file" field 
+                of the output object
+            detection_threshold (float, optional): only detections above this confidence threshold 
+                will be included in the return value
+            image_size (tuple, optional): image size to use for inference, only mess with this
+                if (a) you're using a model other than MegaDetector or (b) you know what you're
+                doing
+            skip_image_resizing (bool, optional): whether to skip internal image resizing (and rely on external 
+                resizing)
 
         Returns:
-        A dict with the following fields, see the 'images' key in https://github.com/agentmorris/MegaDetector/tree/master/api/batch_processing#batch-processing-api-output-format
-            - 'file' (always present)
-            - 'max_detection_conf' (removed from MegaDetector output by default, but generated here)
-            - 'detections', which is a list of detection objects containing keys 'category', 
-              'conf' and 'bbox'
-            - 'failure'
+            dict: a dictionary with the following fields:
+                - 'file' (filename, always present)
+                - 'max_detection_conf' (removed from MegaDetector output files by default, but generated here)
+                - 'detections' (a list of detection objects containing keys 'category', 'conf', and 'bbox')
+                - 'failure' (a failure string, or None if everything went fine)
         """
 
         result = {
@@ -297,13 +309,19 @@ class PTDetector:
 
         return result
 
+    # ...def generate_detections_one_image(...)
+
+# ...class PTDetector
+
 
 #%% Command-line driver
 
+# For testing only... you don't really want to run this module directly.
+
 if __name__ == '__main__':
-    
-    # For testing only... you don't really want to run this module directly
 
+    pass
+    
     #%%
     
     import md_visualization.visualization_utils as vis_utils

detection/run_detector.py

@@ -1,40 +1,26 @@
-########
-#
-# run_detector.py
-#
-# Module to run an animal detection model on images.
-# 
-# The main function in this script also renders the predicted
-# bounding boxes on images and saves the resulting images (with bounding boxes).
-# 
-# This script is not a good way to process lots of images (tens of thousands,
-# say). It does not facilitate checkpointing the results so if it crashes you
-# would have to start from scratch. If you want to run a detector (e.g., ours)
-# on lots of images, you should check out run_detector_batch.py.
-# 
-# To run this script, we recommend you set up a conda virtual environment
-# following instructions in the Installation section on the main README, using
-# `environment-detector.yml` as the environment file where asked.
-# 
-# This is a good way to test our detector on a handful of images and get
-# super-satisfying, graphical results.  It's also a good way to see how fast a
-# detector model will run on a particular machine.
-# 
-# If you would like to *not* use the GPU on the machine, set the environment
-# variable CUDA_VISIBLE_DEVICES to "-1".
-# 
-# If no output directory is specified, writes detections for c:\foo\bar.jpg to
-# c:\foo\bar_detections.jpg.
-# 
-# This script will only consider detections with > 0.005 confidence at all times.
-# The `threshold` you provide is only for rendering the results. If you need to
-# see lower-confidence detections, you can change
-# DEFAULT_OUTPUT_CONFIDENCE_THRESHOLD.
-# 
-# Reference:
-# https://github.com/tensorflow/models/blob/master/research/object_detection/inference/detection_inference.py
-# 
-########
+"""
+
+run_detector.py
+
+Module to run an animal detection model on images.  The main function in this script also renders 
+the predicted bounding boxes on images and saves the resulting images (with bounding boxes).
+
+**This script is not a good way to process lots of images**.  It does not produce a useful
+output format, and it does not facilitate checkpointing the results so if it crashes you
+would have to start from scratch. **If you want to run a detector on lots of images, you should 
+check out run_detector_batch.py**.
+
+That said, this script (run_detector.py) is a good way to test our detector on a handful of images 
+and get super-satisfying, graphical results.
+
+If you would like to *not* use the GPU on the machine, set the environment
+variable CUDA_VISIBLE_DEVICES to "-1".
+
+This script will only consider detections with > 0.005 confidence at all times.
+The threshold you provide is only for rendering the results. If you need to
+see lower-confidence detections, you can change DEFAULT_OUTPUT_CONFIDENCE_THRESHOLD.
+
+"""
 
 #%% Constants, imports, environment
 
@@ -163,9 +149,15 @@ device_token_to_mdv5_inference_speed = {
 
 def convert_to_tf_coords(array):
     """
-    From [x1, y1, width, height] to [y1, x1, y2, x2], where x1 is x_min, x2 is x_max
-
-    This is only used to keep the interface of the synchronous API.
+    Converts a bounding box from [x1, y1, width, height] to [y1, x1, y2, x2].  This 
+    is mostly not helpful, this function only exists to maintain backwards compatibility
+    in the synchronous API, which possibly zero people in the world are using.
+    
+    Args:
+        array (list): a bounding box in [x,y,w,h] format
+        
+    Returns:
+        list: a bounding box in [y1,x1,y2,x2] format
     """
     
     x1 = array[0]
@@ -174,13 +166,21 @@ def convert_to_tf_coords(array):
     height = array[3]
     x2 = x1 + width
     y2 = y1 + height
+    
     return [y1, x1, y2, x2]
 
 
 def get_detector_metadata_from_version_string(detector_version):
     """
-    Given a MegaDetector version string (e.g. "v4.1.0"), return the metadata for
+    Given a MegaDetector version string (e.g. "v4.1.0"), returns the metadata for
     the model.  Used for writing standard defaults to batch output files.
+    
+    Args:
+        detector_version (str): a detection version string, e.g. "v4.1.0", which you
+            can extract from a filename using get_detector_version_from_filename()
+    
+    Returns:
+        dict: metadata for this model, suitable for writing to a MD output file
     """
     
     if detector_version not in DETECTOR_METADATA:
@@ -196,20 +196,26 @@ def get_detector_metadata_from_version_string(detector_version):
 
 
 def get_detector_version_from_filename(detector_filename):
-    """
-    Get the version number component of the detector from the model filename.  
+    r"""
+    Gets the version number component of the detector from the model filename.  
     
-    *detector_filename* will almost always end with one of the following:
+    [detector_filename] will almost always end with one of the following:
         
-    megadetector_v2.pb
-    megadetector_v3.pb
-    megadetector_v4.1 (not produed by run_detector_batch.py, only found in Azure Batch API output files)
-    md_v4.1.0.pb
-    md_v5a.0.0.pt
-    md_v5b.0.0.pt
-    
-    ...for which we identify the version number as "v2.0.0", "v3.0.0", "v4.1.0", 
+    * megadetector_v2.pb
+    * megadetector_v3.pb
+    * megadetector_v4.1 (not produed by run_detector_batch.py, only found in output files from the deprecated Azure Batch API)
+    * md_v4.1.0.pb
+    * md_v5a.0.0.pt
+    * md_v5b.0.0.pt
+    
+    This function identifies the version number as "v2.0.0", "v3.0.0", "v4.1.0", 
     "v4.1.0", "v5a.0.0", and "v5b.0.0", respectively.
+    
+    Args:
+        detector_filename (str): model filename, e.g. c:/x/z/md_v5a.0.0.pt
+    
+    Returns:
+        str: a detector version string, e.g. "v5a.0.0", or "multiple" if I'm confused
     """
     
     fn = os.path.basename(detector_filename).lower()
@@ -228,10 +234,20 @@ def get_detector_version_from_filename(detector_filename):
     
 
 def estimate_md_images_per_second(model_file, device_name=None):
-    """
-    Estimate how fast MegaDetector will run based on benchmarks.  Defaults to querying
+    r"""
+    Estimates how fast MegaDetector will run, based on benchmarks.  Defaults to querying
     the current device.  Returns None if no data is available for the current card/model.
-    Estimates only available for a small handful of GPUs.
+    Estimates only available for a small handful of GPUs.  Uses an absurdly simple lookup
+    approach, e.g. if the string "4090" appears in the device name, congratulations,
+    you have an RTX 4090.
+    
+    Args:
+        model_file (str): model filename, e.g. c:/x/z/md_v5a.0.0.pt
+        device_name (str, optional): device name, e.g. blah-blah-4090-blah-blah
+    
+    Returns:
+        float: the approximate number of images this model version can process on this
+        device per second
     """
     
     if device_name is None:
@@ -271,8 +287,14 @@ def estimate_md_images_per_second(model_file, device_name=None):
     
 def get_typical_confidence_threshold_from_results(results):
     """
-    Given the .json data loaded from a MD results file, determine a typical confidence
+    Given the .json data loaded from a MD results file, returns a typical confidence
     threshold based on the detector version.
+    
+    Args:
+        results (dict): a dict of MD results, as it would be loaded from a MD results .json file
+    
+    Returns:
+        float: a sensible default threshold for this model
     """
     
     if 'detector_metadata' in results['info'] and \
@@ -293,10 +315,16 @@ def get_typical_confidence_threshold_from_results(results):
 
     
 def is_gpu_available(model_file):
-    """
-    Decide whether a GPU is available, importing PyTorch or TF depending on the extension
+    r"""
+    Determines whether a GPU is available, importing PyTorch or TF depending on the extension
     of model_file.  Does not actually load model_file, just uses that to determine how to check 
-    for GPU availability.
+    for GPU availability (PT vs. TF).
+    
+    Args:
+        model_file (str): model filename, e.g. c:/x/z/md_v5a.0.0.pt
+    
+    Returns:
+        bool: whether a GPU is available
     """
     
     if model_file.endswith('.pb'):
@@ -323,8 +351,14 @@ def is_gpu_available(model_file):
 
 
 def load_detector(model_file, force_cpu=False):
-    """
-    Load a TF or PT detector, depending on the extension of model_file.
+    r"""
+    Loads a TF or PT detector, depending on the extension of model_file.
+    
+    Args:
+        model_file (str): model filename, e.g. c:/x/z/md_v5a.0.0.pt
+    
+    Returns:
+        object: loaded detector object
     """
     
     # Possibly automatically download the model
@@ -344,19 +378,41 @@ def load_detector(model_file, force_cpu=False):
         raise ValueError('Unrecognized model format: {}'.format(model_file))
     elapsed = time.time() - start_time
     print('Loaded model in {}'.format(humanfriendly.format_timespan(elapsed)))
+    
     return detector
 
 
 #%% Main function
 
-def load_and_run_detector(model_file, image_file_names, output_dir,
+def load_and_run_detector(model_file, 
+                          image_file_names,
+                          output_dir,
                           render_confidence_threshold=DEFAULT_RENDERING_CONFIDENCE_THRESHOLD,
-                          crop_images=False, box_thickness=DEFAULT_BOX_THICKNESS, 
-                          box_expansion=DEFAULT_BOX_EXPANSION, image_size=None,
-                          label_font_size=DEFAULT_LABEL_FONT_SIZE
+                          crop_images=False, 
+                          box_thickness=DEFAULT_BOX_THICKNESS, 
+                          box_expansion=DEFAULT_BOX_EXPANSION,
+                          image_size=None,
+                          label_font_size=DEFAULT_LABEL_FONT_SIZE                          
                           ):
-    """
-    Load and run detector on target images, and visualize the results.
+    r"""
+    Loads and runs a detector on target images, and visualizes the results.
+    
+    Args:
+        model_file (str): model filename, e.g. c:/x/z/md_v5a.0.0.pt, or a known model
+            string, e.g. "MDV5A"
+        image_file_names (list): list of absolute paths to process
+        output_dir (str): folder to write visualized images to
+        render_confidence_threshold (float, optional): only render boxes for detections
+            above this threshold
+        crop_images (bool, optional): whether to crop detected objects to individual images
+            (default is to render images with boxes, rather than cropping)
+        box_thickness (float, optional): thickness in pixels for box rendering
+        box_expansion (float, optional): box expansion in pixels
+        image_size (tuple, optional): image size to use for inference, only mess with this
+            if (a) you're using a model other than MegaDetector or (b) you know what you're
+            doing
+        label_font_size (float, optional): font size to use for displaying class names
+            and confidence values in the rendered images        
     """
     
     if len(image_file_names) == 0:
@@ -507,7 +563,12 @@ def load_and_run_detector(model_file, image_file_names, output_dir,
 
 def download_model(model_name,force_download=False):
     """
-    Download one of the known models to local temp space if it hasn't already been downloaded
+    Downloads one of the known models to local temp space if it hasn't already been downloaded.
+    
+    Args:
+        model_name (str): a known model string, e.g. "MDV5A"
+        force_download (bool, optional): whether download the model even if the local target 
+            file already exists
     """
     
     import tempfile
@@ -536,9 +597,17 @@ def download_model(model_name,force_download=False):
 
 def try_download_known_detector(detector_file):
     """
-    Check whether detector_file is really the name of a known model, in which case we will
+    Checks whether detector_file is really the name of a known model, in which case we will
     either read the actual filename from the corresponding environment variable or download
     (if necessary) to local temp space.  Otherwise just returns the input string.
+    
+    Args:
+        detector_file (str): a known model string (e.g. "MDV5A"), or any other string (in which
+            case this function is a no-op)
+    
+    Returns:
+        str: the local filename to which the model was downloaded, or the same string that
+        was passed in, if it's not recognized as a well-known model name
     """
     
     if detector_file in downloadable_models:
@@ -606,7 +675,7 @@ def main():
     parser.add_argument(
         '--crop',
         default=False,
-        action="store_true",
+        action='store_true',
         help=('If set, produces separate output images for each crop, '
               'rather than adding bounding boxes to the original image'))
     
@@ -630,7 +699,14 @@ def main():
         default=DEFAULT_LABEL_FONT_SIZE,
         help=('Label font size (defaults to {})'.format(
               DEFAULT_LABEL_FONT_SIZE)))
-        
+    
+    parser.add_argument(
+        '--process_likely_output_images',
+        action='store_true',
+        help=('By default, we skip images that end in {}, because they probably came from this script. '\
+              .format(DETECTION_FILENAME_INSERT) + \
+              'This option disables that behavior.'))
+    
     if len(sys.argv[1:]) == 0:
         parser.print_help()
         parser.exit()
@@ -650,6 +726,16 @@ def main():
     else:
         image_file_names = path_utils.find_images(args.image_dir, args.recursive)
 
+    # Optionally skip images that were probably generated by this script
+    if not args.process_likely_output_images:
+        image_file_names_valid = []
+        for fn in image_file_names:
+            if os.path.splitext(fn)[0].endswith(DETECTION_FILENAME_INSERT):
+                print('Skipping likely output image {}'.format(fn))
+            else:
+                image_file_names_valid.append(fn)
+        image_file_names = image_file_names_valid
+                
     print('Running detector on {} images...'.format(len(image_file_names)))
 
     if args.output_dir:
@@ -671,7 +757,6 @@ def main():
                           image_size=args.image_size,
                           label_font_size=args.label_font_size)
 
-
 if __name__ == '__main__':
     main()