megadetector 5.0.6__py3-none-any.whl → 5.0.7__py3-none-any.whl

detection/process_video.py

@@ -26,12 +26,17 @@ from detection.video_utils import frame_results_to_video_results
 from detection.video_utils import video_folder_to_frames
 from uuid import uuid1
 
+from detection.video_utils import default_fourcc
+
 
 #%% Options classes
 
 class ProcessVideoOptions:
 
-    model_file = ''
+    # 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
     input_video_file = ''
 
     output_json_file = None
@@ -72,9 +77,10 @@ class ProcessVideoOptions:
     
     recursive = False 
     verbose = False
+    
     fourcc = None
 
-    rendering_confidence_threshold = 0.15
+    rendering_confidence_threshold = None
     json_confidence_threshold = 0.005
     frame_sample = None
     
@@ -175,8 +181,14 @@ def process_video(options):
             confidence_threshold=options.rendering_confidence_threshold)
 
         # Combine into a video
-        print('Rendering video to {} at {} fps'.format(options.output_video_file,Fs))
-        frames_to_video(detected_frame_files, Fs, options.output_video_file, codec_spec=options.fourcc)
+        if options.frame_sample is None:
+            rendering_fs = Fs
+        else:
+            rendering_fs = Fs / options.frame_sample
+            
+        print('Rendering video to {} at {} fps (original video {} fps)'.format(
+            options.output_video_file,rendering_fs,Fs))
+        frames_to_video(detected_frame_files, rendering_fs, options.output_video_file, codec_spec=options.fourcc)
         
         # Delete the temporary directory we used for detection images
         if not options.keep_rendered_frames:
@@ -344,11 +356,19 @@ def process_video_folder(options):
             output_video_folder = options.input_video_file
                                 
         # For each video
+        #
+        # TODO: parallelize this loop
+        #
         # i_video=0; input_video_file_abs = video_filenames[i_video]
         for i_video,input_video_file_abs in enumerate(video_filenames):
             
             video_fs = Fs[i_video]
             
+            if options.frame_sample is None:                
+                rendering_fs = video_fs
+            else:
+                rendering_fs = video_fs / options.frame_sample
+            
             input_video_file_relative = os.path.relpath(input_video_file_abs,options.input_video_file)
             video_frame_output_folder = os.path.join(frame_rendering_output_dir,input_video_file_relative)
             assert os.path.isdir(video_frame_output_folder), \
@@ -371,11 +391,10 @@ def process_video_folder(options):
             os.makedirs(os.path.dirname(video_output_file),exist_ok=True)
             
             # Create the output video            
-            print('Rendering detections for video {} to {} at {} fps'.format(input_video_file_relative,
-                                                              video_output_file,video_fs))
-            frames_to_video(video_frame_files, video_fs, video_output_file, codec_spec=options.fourcc)
-            
-    
+            print('Rendering detections for video {} to {} at {} fps (original video {} fps)'.format(
+                input_video_file_relative,video_output_file,rendering_fs,video_fs))
+            frames_to_video(video_frame_files, rendering_fs, video_output_file, codec_spec=options.fourcc)
+                
         # ...for each video
         
         # Possibly clean up rendered frames
@@ -525,12 +544,14 @@ if False:
 
 def main():
 
+    default_options = ProcessVideoOptions()
+    
     parser = argparse.ArgumentParser(description=(
         'Run MegaDetector on each frame in a video (or every Nth frame), optionally '\
         'producing a new video with detections annotated'))
 
     parser.add_argument('model_file', type=str,
-                        help='MegaDetector model file')
+                        help='MegaDetector model file (.pt or .pb) or model name (e.g. "MDV5A")')
 
     parser.add_argument('input_video_file', type=str,
                         help='video file (or folder) to process')
@@ -567,8 +588,8 @@ def main():
     parser.add_argument('--render_output_video', action='store_true',
                         help='enable video output rendering (not rendered by default)')
 
-    parser.add_argument('--fourcc', default=None,
-                        help='fourcc code to use for video encoding, only used if render_output_video is True')
+    parser.add_argument('--fourcc', default=default_fourcc,
+                        help='fourcc code to use for video encoding (default {}), only used if render_output_video is True'.format(default_fourcc))
     
     parser.add_argument('--keep_rendered_frames',
                        action='store_true', help='Disable the deletion of rendered (w/boxes) frames')
@@ -586,11 +607,12 @@ def main():
                            'whether other files were present in the folder.')
         
     parser.add_argument('--rendering_confidence_threshold', type=float,
-                        default=0.8, help="don't render boxes with confidence below this threshold")
+                        default=None, help="don't render boxes with confidence below this threshold (defaults to choosing based on the MD version)")
 
     parser.add_argument('--json_confidence_threshold', type=float,
                         default=0.0, help="don't include boxes in the .json file with confidence "\
-                            'below this threshold')
+                            'below this threshold (default {})'.format(
+                                default_options.json_confidence_threshold))
 
     parser.add_argument('--n_cores', type=int,
                         default=1, help='number of cores to use for frame separation and detection. '\

detection/pytorch_detector.py

@@ -234,7 +234,7 @@ class PTDetector:
             if self.device == 'mps':
                 # As of v1.13.0.dev20220824, nms is not implemented for MPS.
                 #
-                # Send predication back to the CPU to fix.
+                # Send prediction back to the CPU to fix.
                 pred = non_max_suppression(prediction=pred.cpu(), conf_thres=detection_threshold)
             else: 
                 pred = non_max_suppression(prediction=pred, conf_thres=detection_threshold)

detection/run_detector.py

@@ -10,12 +10,7 @@
 # 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:
-# 
-# 1) run_detector_batch.py (for local execution)
-# 
-# 2) https://github.com/agentmorris/MegaDetector/tree/master/api/batch_processing
-#    (for running large jobs on Azure ML)
+# 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
@@ -136,6 +131,33 @@ downloadable_models = {
     'MDV5B':'https://github.com/agentmorris/MegaDetector/releases/download/v5.0/md_v5b.0.0.pt'
 }
 
+model_string_to_model_version = {
+    'v2':'v2.0.0',
+    'v3':'v3.0.0',
+    'v4.1':'v4.1.0',
+    'v5a.0.0':'v5a.0.0',
+    'v5b.0.0':'v5b.0.0',
+    'mdv5a':'v5a.0.0',
+    'mdv5b':'v5b.0.0',
+    'mdv4':'v4.1.0',
+    'mdv3':'v3.0.0'
+}
+
+# Approximate inference speeds (in images per second) for MDv5 based on 
+# benchmarks, only used for reporting very coarse expectations about inference time.
+device_token_to_mdv5_inference_speed = {
+    '4090':17.6,
+    '3090':11.4,
+    '3080':9.5,
+    '3050':4.2,
+    'P2000':2.1,
+    # These are written this way because they're MDv4 benchmarks, and MDv5
+    # is around 3.5x faster than MDv4.
+    'V100':2.79*3.5,
+    '2080':2.3*3.5,
+    '2060':1.6*3.5
+}
+    
 
 #%% Utility functions
 
@@ -190,18 +212,9 @@ def get_detector_version_from_filename(detector_filename):
     "v4.1.0", "v5a.0.0", and "v5b.0.0", respectively.
     """
     
-    fn = os.path.basename(detector_filename)
-    known_model_versions = {'v2':'v2.0.0',
-                            'v3':'v3.0.0',
-                            'v4.1':'v4.1.0',
-                            'v5a.0.0':'v5a.0.0',
-                            'v5b.0.0':'v5b.0.0',
-                            'MDV5A':'v5a.0.0',
-                            'MDV5B':'v5b.0.0',
-                            'MDV4':'v4.1.0',
-                            'MDV3':'v3.0.0'}
+    fn = os.path.basename(detector_filename).lower()
     matches = []
-    for s in known_model_versions.keys():
+    for s in model_string_to_model_version.keys():
         if s in fn:
             matches.append(s)
     if len(matches) == 0:
@@ -211,9 +224,51 @@ def get_detector_version_from_filename(detector_filename):
         print('Warning: multiple MegaDetector versions for model file {}'.format(detector_filename))
         return 'multiple'
     else:
-        return known_model_versions[matches[0]]
+        return model_string_to_model_version[matches[0]]
     
 
+def estimate_md_images_per_second(model_file, device_name=None):
+    """
+    Estimate 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.
+    """
+    
+    if device_name is None:
+        try:
+            import torch
+            device_name = torch.cuda.get_device_name()
+        except Exception as e:
+            print('Error querying device name: {}'.format(e))
+            return None
+    
+    model_file = model_file.lower().strip()
+    if model_file in model_string_to_model_version.values():
+        model_version = model_file
+    else:
+        model_version = get_detector_version_from_filename(model_file)
+        if model_version not in model_string_to_model_version.values():
+            print('Error determining model version for model file {}'.format(model_file))
+            return None
+    
+    mdv5_inference_speed = None
+    for device_token in device_token_to_mdv5_inference_speed.keys():
+        if device_token in device_name:
+            mdv5_inference_speed = device_token_to_mdv5_inference_speed[device_token]
+            break
+    
+    if mdv5_inference_speed is None:
+        print('No speed estimate available for {}'.format(device_name))
+    
+    if 'v5' in model_version:
+        return mdv5_inference_speed
+    elif 'v2' in model_version or 'v3' in model_version or 'v4' in model_version:
+        return mdv5_inference_speed / 3.5
+    else:
+        print('Could not estimate inference speed for model file {}'.format(model_file))
+        return None
+    
+    
 def get_typical_confidence_threshold_from_results(results):
     """
     Given the .json data loaded from a MD results file, determine a typical confidence

detection/run_detector_batch.py

@@ -751,17 +751,75 @@ if False:
 
     #%%
     
+    model_file = 'MDV5A'
+    image_dir = r'g:\camera_traps\camera_trap_images'
+    output_file = r'g:\temp\md-test.json'
+    
+    recursive = True
+    output_relative_filenames = True
+    include_max_conf = False
+    quiet = True
+    image_size = None
+    use_image_queue = False
+    confidence_threshold = 0.0001
+    checkpoint_frequency = 5   
     checkpoint_path = None
-    model_file = r'G:\temp\models\md_v4.1.0.pb'
-    confidence_threshold = 0.1
-    checkpoint_frequency = -1
-    results = None
+    resume_from_checkpoint = 'auto'
+    allow_checkpoint_overwrite = False
     ncores = 1
-    use_image_queue = False
-    quiet = False
-    image_dir = r'G:\temp\demo_images\ssmini'
-    image_size = None
+    class_mapping_filename = None
+    include_image_size = True
+    include_image_timestamp = True
+    include_exif_data = True
+    overwrite_handling = None
+        
+    # Generate a command line
+    cmd = 'python run_detector_batch.py "{}" "{}" "{}"'.format(
+        model_file,image_dir,output_file)
+    
+    if recursive:
+        cmd += ' --recursive'
+    if output_relative_filenames:
+        cmd += ' --output_relative_filenames'
+    if include_max_conf:
+        cmd += ' --include_max_conf'
+    if quiet:
+        cmd += ' --quiet'
+    if image_size is not None:
+        cmd += ' --image_size {}'.format(image_size)
+    if use_image_queue:
+        cmd += ' --use_image_queue'
+    if confidence_threshold is not None:
+        cmd += ' --threshold {}'.format(confidence_threshold)
+    if checkpoint_frequency is not None:
+        cmd += ' --checkpoint_frequency {}'.format(checkpoint_frequency)
+    if checkpoint_path is not None:
+        cmd += ' --checkpoint_path "{}"'.format(checkpoint_path)
+    if resume_from_checkpoint is not None:
+        cmd += ' --resume_from_checkpoint "{}"'.format(resume_from_checkpoint)
+    if allow_checkpoint_overwrite:
+        cmd += ' --allow_checkpoint_overwrite'
+    if ncores is not None:
+        cmd += ' --ncores {}'.format(ncores)
+    if class_mapping_filename is not None:
+        cmd += ' --class_mapping_filename "{}"'.format(class_mapping_filename)
+    if include_image_size:
+        cmd += ' --include_image_size'
+    if include_image_timestamp:
+        cmd += ' --include_image_timestamp'
+    if include_exif_data:
+        cmd += ' --include_exif_data'
+    if overwrite_handling is not None:
+        cmd += ' --overwrite_handling {}'.format(overwrite_handling)
+    
+    print(cmd)
+    import clipboard; clipboard.copy(cmd)
+    
+    
+    #%% Run inference interactively
+    
     image_file_names = path_utils.find_images(image_dir, recursive=False)    
+    results = None
     
     start_time = time.time()
     
@@ -840,12 +898,15 @@ def main():
         '--checkpoint_path',
         type=str,
         default=None,
-        help='File name to which checkpoints will be written if checkpoint_frequency is > 0')    
+        help='File name to which checkpoints will be written if checkpoint_frequency is > 0, ' + \
+             'defaults to md_checkpoint_[date].json in the same folder as the output file')    
     parser.add_argument(
         '--resume_from_checkpoint',
         type=str,
         default=None,
-        help='Path to a JSON checkpoint file to resume from')
+        help='Path to a JSON checkpoint file to resume from, or "auto" to ' + \
+             'find the most recent checkpoint in the same folder as the output file.  "auto" uses' + \
+             'checkpoint_path (rather than searching the output folder) if checkpoint_path is specified.')
     parser.add_argument(
         '--allow_checkpoint_overwrite',
         action='store_true',
@@ -897,7 +958,7 @@ def main():
     
     assert os.path.exists(args.detector_file), \
         'detector file {} does not exist'.format(args.detector_file)
-    assert 0.0 < args.threshold <= 1.0, 'Confidence threshold needs to be between 0 and 1'
+    assert 0.0 <= args.threshold <= 1.0, 'Confidence threshold needs to be between 0 and 1'
     assert args.output_file.endswith('.json'), 'output_file specified needs to end with .json'
     if args.checkpoint_frequency != -1:
         assert args.checkpoint_frequency > 0, 'Checkpoint_frequency needs to be > 0 or == -1'
@@ -919,19 +980,42 @@ def main():
         else:
             raise ValueError('Illegal overwrite handling string {}'.format(args.overwrite_handling))
 
+    output_dir = os.path.dirname(args.output_file)
+
+    if len(output_dir) > 0:
+        os.makedirs(output_dir,exist_ok=True)
+        
+    assert not os.path.isdir(args.output_file), 'Specified output file is a directory'
+    
     if args.class_mapping_filename is not None:
         load_custom_class_mapping(args.class_mapping_filename)
-        
+    
     # Load the checkpoint if available
     #
     # Relative file names are only output at the end; all file paths in the checkpoint are
-    # still full paths.
+    # still absolute paths.
     if args.resume_from_checkpoint is not None:
-        assert os.path.exists(args.resume_from_checkpoint), \
+        if args.resume_from_checkpoint == 'auto':
+            checkpoint_files = os.listdir(output_dir)
+            checkpoint_files = [fn for fn in checkpoint_files if \
+                                (fn.startswith('md_checkpoint') and fn.endswith('.json'))]
+            if len(checkpoint_files) == 0:
+                raise ValueError('resume_from_checkpoint set to "auto", but no checkpoints found in {}'.format(
+                    output_dir))
+            else:
+                if len(checkpoint_files) > 1:
+                    print('Warning: found {} checkpoints in {}, using the latest'.format(
+                        len(checkpoint_files),output_dir))
+                    checkpoint_files = sorted(checkpoint_files)
+                checkpoint_file_relative = checkpoint_files[-1]
+                checkpoint_file = os.path.join(output_dir,checkpoint_file_relative)                                
+        else:
+            checkpoint_file = args.resume_from_checkpoint
+        assert os.path.exists(checkpoint_file), \
             'File at resume_from_checkpoint specified does not exist'
-        with open(args.resume_from_checkpoint) as f:
+        with open(checkpoint_file) as f:
             print('Loading previous results from checkpoint file {}'.format(
-                args.resume_from_checkpoint))
+                checkpoint_file))
             saved = json.load(f)
         assert 'images' in saved, \
             'The checkpoint file does not have the correct fields; cannot be restored'
@@ -982,13 +1066,6 @@ def main():
     assert os.path.exists(image_file_names[0]), \
         'The first image to be processed does not exist at {}'.format(image_file_names[0])
 
-    output_dir = os.path.dirname(args.output_file)
-
-    if len(output_dir) > 0:
-        os.makedirs(output_dir,exist_ok=True)
-        
-    assert not os.path.isdir(args.output_file), 'Specified output file is a directory'
-
     # Test that we can write to the output_file's dir if checkpointing requested
     if args.checkpoint_frequency != -1:
         
@@ -996,7 +1073,7 @@ def main():
             checkpoint_path = args.checkpoint_path
         else:
             checkpoint_path = os.path.join(output_dir,
-                                           'checkpoint_{}.json'.format(
+                                           'md_checkpoint_{}.json'.format(
                                                datetime.utcnow().strftime("%Y%m%d%H%M%S")))
         
         # Don't overwrite existing checkpoint files, this is a sure-fire way to eventually
@@ -1023,6 +1100,9 @@ def main():
         
     else:
         
+        if args.checkpoint_path is not None:
+            print('Warning: checkpointing disabled because checkpoint_frequency is -1, ' + \
+                  'but a checkpoint path was specified')
         checkpoint_path = None
 
     start_time = time.time()