megadetector 5.0.12__py3-none-any.whl → 5.0.14__py3-none-any.whl

megadetector/detection/pytorch_detector.py

@@ -8,11 +8,9 @@ Module to run MegaDetector v5, a PyTorch YOLOv5 animal detection model.
 
 #%% Imports and constants
 
-import sys
 import torch
 import numpy as np
 import traceback
-import builtins
 
 from megadetector.detection.run_detector import CONF_DIGITS, COORD_DIGITS, FAILURE_INFER
 from megadetector.utils import ct_utils
@@ -130,7 +128,20 @@ class PTDetector:
                     self.device = 'mps'
             except AttributeError:
                 pass
-        self.model = PTDetector._load_model(model_path, self.device)
+        try:
+            self.model = PTDetector._load_model(model_path, self.device)
+        except Exception as e:
+            # In a very estoeric scenario where an old version of YOLOv5 is used to run
+            # newer models, we run into an issue because the "Model" class became
+            # "DetectionModel".  New YOLOv5 code handles this case by just setting them
+            # to be the same, so doing that via monkey-patch doesn't seem *that* rude.
+            if "Can't get attribute 'DetectionModel'" in str(e):
+                print('Forward-compatibility issue detected, patching')
+                from models import yolo
+                yolo.DetectionModel = yolo.Model                
+                self.model = PTDetector._load_model(model_path, self.device)                
+            else:
+                raise
         if (self.device != 'cpu'):
             print('Sending model to GPU')
             self.model.to(self.device)
@@ -169,9 +180,12 @@ class PTDetector:
             
         return model
 
-    def generate_detections_one_image(self, img_original, image_id='unknown', 
-                                      detection_threshold=0.00001, image_size=None,
-                                      skip_image_resizing=False):
+    def generate_detections_one_image(self, img_original, 
+                                      image_id='unknown', 
+                                      detection_threshold=0.00001, 
+                                      image_size=None,
+                                      skip_image_resizing=False,
+                                      augment=False):
         """
         Applies the detector to an image.
 
@@ -181,11 +195,11 @@ class PTDetector:
                 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)
+            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 getting into
+            skip_image_resizing (bool, optional): whether to skip internal image resizing (and rely on 
+                external resizing)
+            augment (bool, optional): enable (implementation-specific) image augmentation
 
         Returns:
             dict: a dictionary with the following fields:
@@ -201,6 +215,9 @@ class PTDetector:
         detections = []
         max_conf = 0.0
 
+        if detection_threshold is None:
+            detection_threshold = 0
+            
         try:
             
             img_original = np.asarray(img_original)
@@ -228,8 +245,10 @@ class PTDetector:
             if skip_image_resizing:
                 img = img_original
             else:
-                letterbox_result = letterbox(img_original, new_shape=target_size,
-                                stride=PTDetector.STRIDE, auto=True)
+                letterbox_result = letterbox(img_original, 
+                                             new_shape=target_size,
+                                             stride=PTDetector.STRIDE, 
+                                             auto=True)
                 img = letterbox_result[0]                
             
             # HWC to CHW; PIL Image is RGB already
@@ -244,7 +263,7 @@ class PTDetector:
             if len(img.shape) == 3:  
                 img = torch.unsqueeze(img, 0)
 
-            pred: list = self.model(img)[0]
+            pred = self.model(img,augment=augment)[0]
 
             # NMS
             if self.device == 'mps':

megadetector/detection/run_detector.py

@@ -394,7 +394,8 @@ def load_and_run_detector(model_file,
                           box_thickness=DEFAULT_BOX_THICKNESS, 
                           box_expansion=DEFAULT_BOX_EXPANSION,
                           image_size=None,
-                          label_font_size=DEFAULT_LABEL_FONT_SIZE                          
+                          label_font_size=DEFAULT_LABEL_FONT_SIZE,
+                          augment=False
                           ):
     r"""
     Loads and runs a detector on target images, and visualizes the results.
@@ -415,6 +416,7 @@ def load_and_run_detector(model_file,
             doing
         label_font_size (float, optional): font size to use for displaying class names
             and confidence values in the rendered images        
+        augment (bool, optional): enable (implementation-specific) image augmentation
     """
     
     if len(image_file_names) == 0:
@@ -506,9 +508,12 @@ def load_and_run_detector(model_file,
         try:
             start_time = time.time()
 
-            result = detector.generate_detections_one_image(image, im_file,
-                       detection_threshold=DEFAULT_OUTPUT_CONFIDENCE_THRESHOLD,
-                       image_size=image_size)
+            result = detector.generate_detections_one_image(
+                image, 
+                im_file,
+                detection_threshold=DEFAULT_OUTPUT_CONFIDENCE_THRESHOLD,
+                image_size=image_size,
+                augment=augment)
             detection_results.append(result)
 
             elapsed = time.time() - start_time
@@ -679,6 +684,12 @@ def main():
         help=('If set, produces separate output images for each crop, '
               'rather than adding bounding boxes to the original image'))
     
+    parser.add_argument(
+        '--augment',
+        default=False,
+        action='store_true',
+        help=('Enable image augmentation'))
+    
     parser.add_argument(
         '--box_thickness',
         type=int,
@@ -755,7 +766,8 @@ def main():
                           box_expansion=args.box_expansion,                          
                           crop_images=args.crop,
                           image_size=args.image_size,
-                          label_font_size=args.label_font_size)
+                          label_font_size=args.label_font_size,
+                          augment=args.augment)
 
 if __name__ == '__main__':
     main()

megadetector/detection/run_detector_batch.py

@@ -126,7 +126,15 @@ def _producer_func(q,image_files):
     print('Finished image loading'); sys.stdout.flush()
     
     
-def _consumer_func(q,return_queue,model_file,confidence_threshold,image_size=None):
+def _consumer_func(q,
+                   return_queue,
+                   model_file,
+                   confidence_threshold,
+                   image_size=None,
+                   include_image_size=False,
+                   include_image_timestamp=False, 
+                   include_exif_data=False,
+                   augment=False):
     """ 
     Consumer function; only used when using the (optional) image queue.
     
@@ -168,16 +176,30 @@ def _consumer_func(q,return_queue,model_file,confidence_threshold,image_size=Non
             results.append({'file': im_file,
                             'failure': image})
         else:
-            results.append(process_image(im_file=im_file,detector=detector,
+            results.append(process_image(im_file=im_file,
+                                         detector=detector,
                                          confidence_threshold=confidence_threshold,
-                                         image=image,quiet=True,image_size=image_size))
+                                         image=image,
+                                         quiet=True,
+                                         image_size=image_size,
+                                         include_image_size=include_image_size,
+                                         include_image_timestamp=include_image_timestamp, 
+                                         include_exif_data=include_exif_data,
+                                         augment=augment))
         if verbose:
             print('Processed image {}'.format(im_file)); sys.stdout.flush()
         q.task_done()
             
 
-def run_detector_with_image_queue(image_files,model_file,confidence_threshold,
-                                  quiet=False,image_size=None):
+def run_detector_with_image_queue(image_files,
+                                  model_file,
+                                  confidence_threshold,
+                                  quiet=False,
+                                  image_size=None,
+                                  include_image_size=False, 
+                                  include_image_timestamp=False,
+                                  include_exif_data=False,
+                                  augment=False):
     """
     Driver function for the (optional) multiprocessing-based image queue; only used 
     when --use_image_queue is specified.  Starts a reader process to read images from disk, but 
@@ -218,15 +240,37 @@ def run_detector_with_image_queue(image_files,model_file,confidence_threshold,
 
     if run_separate_consumer_process:
         if use_threads_for_queue:
-            consumer = Thread(target=_consumer_func,args=(q,return_queue,model_file,
-                                                         confidence_threshold,image_size,))
+            consumer = Thread(target=_consumer_func,args=(q,
+                                                          return_queue,
+                                                          model_file,
+                                                          confidence_threshold,
+                                                          image_size,
+                                                          include_image_size,
+                                                          include_image_timestamp, 
+                                                          include_exif_data,
+                                                          augment))
         else:
-            consumer = Process(target=_consumer_func,args=(q,return_queue,model_file,
-                                                          confidence_threshold,image_size,))
+            consumer = Process(target=_consumer_func,args=(q,
+                                                           return_queue,
+                                                           model_file,
+                                                           confidence_threshold,
+                                                           image_size,
+                                                           include_image_size,
+                                                           include_image_timestamp, 
+                                                           include_exif_data,
+                                                           augment))
         consumer.daemon = True
         consumer.start()
     else:
-        _consumer_func(q,return_queue,model_file,confidence_threshold,image_size)
+        _consumer_func(q,
+                       return_queue,
+                       model_file,
+                       confidence_threshold,
+                       image_size,
+                       include_image_size,
+                       include_image_timestamp, 
+                       include_exif_data,
+                       augment)
 
     producer.join()
     print('Producer finished')
@@ -262,15 +306,20 @@ def _chunks_by_number_of_chunks(ls, n):
 
 #%% Image processing functions
 
-def process_images(im_files, detector, confidence_threshold, use_image_queue=False, 
-                   quiet=False, image_size=None, checkpoint_queue=None, 
-                   include_image_size=False, include_image_timestamp=False, 
-                   include_exif_data=False):
+def process_images(im_files, 
+                   detector, 
+                   confidence_threshold, 
+                   use_image_queue=False, 
+                   quiet=False, 
+                   image_size=None, 
+                   checkpoint_queue=None, 
+                   include_image_size=False, 
+                   include_image_timestamp=False, 
+                   include_exif_data=False,
+                   augment=False):
     """
-    Runs a detector (typically MegaDetector) over a list of image files.  
-    As of 3/2024, this entry point is used when the image queue is enabled, but not in the 
-    standard inference path (which instead loops over process_image()).
-
+    Runs a detector (typically MegaDetector) over a list of image files on a single thread.
+    
     Args:
         im_files (list: paths to image files                                   
         detector (str or detector object): loaded model or str; if this is a string, it can be a
@@ -284,7 +333,8 @@ def process_images(im_files, detector, confidence_threshold, use_image_queue=Fal
         checkpoint_queue (Queue, optional): internal parameter used to pass image queues around
         include_image_size (bool, optional): should we include image size in the output for each image?
         include_image_timestamp (bool, optional): should we include image timestamps in the output for each image?
-        include_exif_data (bool, optional): should we include EXIF data in the output for each image?        
+        include_exif_data (bool, optional): should we include EXIF data in the output for each image?
+        augment (bool, optional): enable image augmentation
 
     Returns:
         list: list of dicts, in which each dict represents detections on one image,
@@ -292,25 +342,37 @@ def process_images(im_files, detector, confidence_threshold, use_image_queue=Fal
     """
     
     if isinstance(detector, str):
+        
         start_time = time.time()
         detector = load_detector(detector)
         elapsed = time.time() - start_time
         print('Loaded model (batch level) in {}'.format(humanfriendly.format_timespan(elapsed)))
 
     if use_image_queue:
-        run_detector_with_image_queue(im_files, detector, confidence_threshold, 
-                                      quiet=quiet, image_size=image_size,
+        
+        run_detector_with_image_queue(im_files, 
+                                      detector, 
+                                      confidence_threshold, 
+                                      quiet=quiet, 
+                                      image_size=image_size,
                                       include_image_size=include_image_size, 
                                       include_image_timestamp=include_image_timestamp,
-                                      include_exif_data=include_exif_data)
+                                      include_exif_data=include_exif_data,
+                                      augment=augment)
+        
     else:            
+        
         results = []
         for im_file in im_files:
-            result = process_image(im_file, detector, confidence_threshold,
-                                         quiet=quiet, image_size=image_size, 
-                                         include_image_size=include_image_size, 
-                                         include_image_timestamp=include_image_timestamp,
-                                         include_exif_data=include_exif_data)
+            result = process_image(im_file, 
+                                   detector, 
+                                   confidence_threshold,
+                                   quiet=quiet, 
+                                   image_size=image_size, 
+                                   include_image_size=include_image_size, 
+                                   include_image_timestamp=include_image_timestamp,
+                                   include_exif_data=include_exif_data,
+                                   augment=augment)
 
             if checkpoint_queue is not None:
                 checkpoint_queue.put(result)
@@ -321,10 +383,16 @@ def process_images(im_files, detector, confidence_threshold, use_image_queue=Fal
 # ...def process_images(...)
 
 
-def process_image(im_file, detector, confidence_threshold, image=None, 
-                  quiet=False, image_size=None, include_image_size=False,
-                  include_image_timestamp=False, include_exif_data=False,
-                  skip_image_resizing=False):
+def process_image(im_file, detector, 
+                  confidence_threshold, 
+                  image=None, 
+                  quiet=False, 
+                  image_size=None, 
+                  include_image_size=False,
+                  include_image_timestamp=False, 
+                  include_exif_data=False,
+                  skip_image_resizing=False,
+                  augment=False):
     """
     Runs a detector (typically MegaDetector) on a single image file.
 
@@ -343,6 +411,7 @@ def process_image(im_file, detector, confidence_threshold, image=None,
         include_image_timestamp (bool, optional): should we include image timestamps in the output for each image?
         include_exif_data (bool, optional): should we include EXIF data in the output for each image?                
         skip_image_resizing (bool, optional): whether to skip internal image resizing and rely on external resizing
+        augment (bool, optional): enable image augmentation
 
     Returns:
         dict: dict representing detections on one image,
@@ -367,8 +436,12 @@ def process_image(im_file, detector, confidence_threshold, image=None,
 
     try:
         result = detector.generate_detections_one_image(
-            image, im_file, detection_threshold=confidence_threshold, image_size=image_size,
-            skip_image_resizing=skip_image_resizing)
+                    image, 
+                    im_file, 
+                    detection_threshold=confidence_threshold, 
+                    image_size=image_size,
+                    skip_image_resizing=skip_image_resizing,
+                    augment=augment)
     except Exception as e:
         if not quiet:
             print('Image {} cannot be processed. Exception: {}'.format(im_file, e))
@@ -423,12 +496,21 @@ def _load_custom_class_mapping(class_mapping_filename):
     
 #%% Main function
 
-def load_and_run_detector_batch(model_file, image_file_names, checkpoint_path=None,
+def load_and_run_detector_batch(model_file, 
+                                image_file_names, 
+                                checkpoint_path=None,
                                 confidence_threshold=run_detector.DEFAULT_OUTPUT_CONFIDENCE_THRESHOLD,
-                                checkpoint_frequency=-1, results=None, n_cores=1,
-                                use_image_queue=False, quiet=False, image_size=None, 
-                                class_mapping_filename=None, include_image_size=False, 
-                                include_image_timestamp=False, include_exif_data=False):
+                                checkpoint_frequency=-1, 
+                                results=None, 
+                                n_cores=1,
+                                use_image_queue=False, 
+                                quiet=False, 
+                                image_size=None, 
+                                class_mapping_filename=None, 
+                                include_image_size=False, 
+                                include_image_timestamp=False, 
+                                include_exif_data=False,
+                                augment=False):
     """
     Load a model file and run it on a list of images.
     
@@ -455,7 +537,8 @@ def load_and_run_detector_batch(model_file, image_file_names, checkpoint_path=No
             file or YOLOv5 dataset.yaml file
         include_image_size (bool, optional): should we include image size in the output for each image?
         include_image_timestamp (bool, optional): should we include image timestamps in the output for each image?
-        include_exif_data (bool, optional): should we include EXIF data in the output for each image?        
+        include_exif_data (bool, optional): should we include EXIF data in the output for each image?
+        augment (bool, optional): enable image augmentation
         
     Returns:
         results: list of dicts; each dict represents detections on one image
@@ -539,9 +622,15 @@ def load_and_run_detector_batch(model_file, image_file_names, checkpoint_path=No
         assert len(results) == 0, \
             'Using an image queue with results loaded from a checkpoint is not currently supported'
         assert n_cores <= 1
-        results = run_detector_with_image_queue(image_file_names, model_file, 
-                                                confidence_threshold, quiet, 
-                                                image_size=image_size)
+        results = run_detector_with_image_queue(image_file_names, 
+                                                model_file, 
+                                                confidence_threshold, 
+                                                quiet, 
+                                                image_size=image_size,
+                                                include_image_size=include_image_size,
+                                                include_image_timestamp=include_image_timestamp,
+                                                include_exif_data=include_exif_data,
+                                                augment=augment)
         
     elif n_cores <= 1:
 
@@ -565,11 +654,15 @@ def load_and_run_detector_batch(model_file, image_file_names, checkpoint_path=No
 
             count += 1
 
-            result = process_image(im_file, detector, 
-                                   confidence_threshold, quiet=quiet, 
-                                   image_size=image_size, include_image_size=include_image_size,
+            result = process_image(im_file, 
+                                   detector, 
+                                   confidence_threshold, 
+                                   quiet=quiet, 
+                                   image_size=image_size, 
+                                   include_image_size=include_image_size,
                                    include_image_timestamp=include_image_timestamp,
-                                   include_exif_data=include_exif_data)
+                                   include_exif_data=include_exif_data,
+                                   augment=augment)
             results.append(result)
 
             # Write a checkpoint if necessary
@@ -616,14 +709,18 @@ def load_and_run_detector_batch(model_file, image_file_names, checkpoint_path=No
                                              checkpoint_queue, results), daemon=True)
             checkpoint_thread.start()
 
-            pool.map(partial(process_images, detector=detector,
-                                    confidence_threshold=confidence_threshold,
-                                    image_size=image_size, 
-                                    include_image_size=include_image_size,
-                                    include_image_timestamp=include_image_timestamp,
-                                    include_exif_data=include_exif_data,
-                                    checkpoint_queue=checkpoint_queue), 
-                                    image_batches)
+            pool.map(partial(process_images, 
+                             detector=detector,
+                             confidence_threshold=confidence_threshold,
+                             use_image_queue=False,
+                             quiet=quiet,
+                             image_size=image_size, 
+                             checkpoint_queue=checkpoint_queue,
+                             include_image_size=include_image_size,
+                             include_image_timestamp=include_image_timestamp,
+                             include_exif_data=include_exif_data,
+                             augment=augment), 
+                             image_batches)
 
             checkpoint_queue.put(None)
 
@@ -631,12 +728,18 @@ def load_and_run_detector_batch(model_file, image_file_names, checkpoint_path=No
             
             # Multprocessing is enabled, but checkpointing is not
             
-            new_results = pool.map(partial(process_images, detector=detector,
-                                   confidence_threshold=confidence_threshold,image_size=image_size,
-                                   include_image_size=include_image_size,
-                                   include_image_timestamp=include_image_timestamp,
-                                   include_exif_data=include_exif_data), 
-                                   image_batches)
+            new_results = pool.map(partial(process_images, 
+                                           detector=detector,
+                                           confidence_threshold=confidence_threshold,
+                                           use_image_queue=False,
+                                           quiet=quiet,
+                                           checkpoint_queue=None,
+                                           image_size=image_size,
+                                           include_image_size=include_image_size,
+                                           include_image_timestamp=include_image_timestamp,
+                                           include_exif_data=include_exif_data,
+                                           augment=augment), 
+                                           image_batches)
 
             new_results = list(itertools.chain.from_iterable(new_results))
             
@@ -724,9 +827,14 @@ def get_image_datetime(image):
         return None        
 
 
-def write_results_to_file(results, output_file, relative_path_base=None, 
-                          detector_file=None, info=None, include_max_conf=False,
-                          custom_metadata=None, force_forward_slashes=True):
+def write_results_to_file(results, 
+                          output_file, 
+                          relative_path_base=None, 
+                          detector_file=None, 
+                          info=None, 
+                          include_max_conf=False,
+                          custom_metadata=None, 
+                          force_forward_slashes=True):
     """
     Writes list of detection results to JSON output file. Format matches:
 
@@ -951,7 +1059,12 @@ def main():
         '--image_size',
         type=int,
         default=None,
-        help=('Force image resizing to a (square) integer size (not recommended to change this)'))    
+        help=('Force image resizing to a specific integer size on the long axis (not recommended to change this)'))    
+    parser.add_argument(
+        '--augment',
+        action='store_true',
+        help='Enable image augmentation'
+    )
     parser.add_argument(
         '--use_image_queue',
         action='store_true',
@@ -1196,7 +1309,8 @@ def main():
                                           class_mapping_filename=args.class_mapping_filename,
                                           include_image_size=args.include_image_size,
                                           include_image_timestamp=args.include_image_timestamp,
-                                          include_exif_data=args.include_exif_data)
+                                          include_exif_data=args.include_exif_data,
+                                          augment=args.augment)
 
     elapsed = time.time() - start_time
     images_per_second = len(results) / elapsed