megadetector 5.0.13__py3-none-any.whl → 5.0.15__py3-none-any.whl

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))
             
@@ -645,8 +748,13 @@ def load_and_run_detector_batch(model_file, image_file_names, checkpoint_path=No
             results += new_results
 
         # ...if checkpointing is/isn't enabled
-        
-    # ...if we're running (1) with image queue, (2) on one core, (3) on multiple cores
+    
+        try:
+            pool.close()
+        except Exception as e:
+            print('Warning: error closing multiprocessing pool:\n{}'.format(str(e)))
+                
+    # ...if we're running (1) with image queue, (2) on one core, or (3) on multiple cores
     
     # 'results' may have been modified in place, but we also return it for
     # backwards-compatibility.
@@ -724,9 +832,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 +1064,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 +1314,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

megadetector/detection/tf_detector.py

@@ -125,8 +125,13 @@ class TFDetector:
         return box_tensor_out, score_tensor_out, class_tensor_out
 
 
-    def generate_detections_one_image(self, image, image_id, detection_threshold, image_size=None,
-                                      skip_image_resizing=False):
+    def generate_detections_one_image(self, 
+                                      image, 
+                                      image_id, 
+                                      detection_threshold, 
+                                      image_size=None,
+                                      skip_image_resizing=False,
+                                      augment=False):
         """
         Runs the detector on an image.
 
@@ -139,7 +144,9 @@ class TFDetector:
                 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)... not currently supported, but included here for compatibility with PTDetector.
+                resizing).  Not currently supported, but included here for compatibility with PTDetector.
+            augment (bool, optional): enable image augmentation.  Not currently  supported, but included 
+                here for compatibility with PTDetector.
 
         Returns:
             dict: a dictionary with the following fields:
@@ -151,6 +158,7 @@ class TFDetector:
         
         assert image_size is None, 'Image sizing not supported for TF detectors'
         assert not skip_image_resizing, 'Image sizing not supported for TF detectors'
+        assert not augment, 'Image augmentation is not supported for TF detectors'
         
         if detection_threshold is None:
             detection_threshold = 0