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

detection/run_detector_batch.py

@@ -350,26 +350,56 @@ def process_image(im_file, detector, confidence_threshold, image=None,
 # ...def process_image(...)
 
 
+def load_custom_class_mapping(class_mapping_filename):
+    """
+    This is an experimental hack to allow the use of non-MD YOLOv5 models through
+    the same infrastructure; it disables the code that enforces MDv5-like class lists.
+    
+    Should be a .json file that maps int-strings to strings, or a YOLOv5 dataset.yaml file.
+    """
+    
+    if class_mapping_filename is None:
+        return
+    
+    run_detector.USE_MODEL_NATIVE_CLASSES = True
+    if class_mapping_filename.endswith('.json'):
+        with open(class_mapping_filename,'r') as f:
+            class_mapping = json.load(f)
+    elif (class_mapping_filename.endswith('.yml') or class_mapping_filename.endswith('.yaml')):
+        from data_management.yolo_output_to_md_output import read_classes_from_yolo_dataset_file
+        class_mapping = read_classes_from_yolo_dataset_file(class_mapping_filename)
+        # convert from ints to int-strings
+        class_mapping = {str(k):v for k,v in class_mapping.items()}
+    else:
+        raise ValueError('Unrecognized class mapping file {}'.format(class_mapping_filename))
+        
+    print('Loaded custom class mapping:')
+    print(class_mapping)
+    run_detector.DEFAULT_DETECTOR_LABEL_MAP = class_mapping
+    return class_mapping
+    
+    
 #%% Main function
 
 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):
+                                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):
     """
     Args
     - model_file: path to model file, or supported model string (e.g. "MDV5A")
     - image_file_names: list of strings (image filenames), a single image filename, 
-                        a folder to recursively search for images in, or a .json file containing
-                        a list of images.
+                        a folder to recursively search for images in, or a .json or .txt file
+                        containing a list of images.
     - checkpoint_path: str, path to JSON checkpoint file
     - confidence_threshold: float, only detections above this threshold are returned
     - checkpoint_frequency: int, write results to JSON checkpoint file every N images
     - results: list of dict, existing results loaded from checkpoint
     - n_cores: int, # of CPU cores to use
     - class_mapping_filename: str, use a non-default class mapping supplied in a .json file
+      or YOLOv5 dataset.yaml file.
 
     Returns
     - results: list of dicts; each dict represents detections on one image
@@ -384,15 +414,8 @@ def load_and_run_detector_batch(model_file, image_file_names, checkpoint_path=No
     if checkpoint_frequency is None:
         checkpoint_frequency = -1
 
-    # This is an experimental hack to allow the use of non-MD YOLOv5 models through
-    # the same infrastructure; it disables the code that enforces MDv5-like class lists.
     if class_mapping_filename is not None:
-        run_detector.USE_MODEL_NATIVE_CLASSES = True
-        with open(class_mapping_filename,'r') as f:
-            class_mapping = json.load(f)
-        print('Loaded custom class mapping:')
-        print(class_mapping)
-        run_detector.DEFAULT_DETECTOR_LABEL_MAP = class_mapping
+        load_custom_class_mapping(class_mapping_filename)
         
     # Handle the case where image_file_names is not yet actually a list
     if isinstance(image_file_names,str):
@@ -403,22 +426,32 @@ def load_and_run_detector_batch(model_file, image_file_names, checkpoint_path=No
             image_file_names = path_utils.find_images(image_dir, True)
             print('{} image files found in folder {}'.format(len(image_file_names),image_dir))
             
-        # A json list of image paths
-        elif os.path.isfile(image_file_names) and image_file_names.endswith('.json'):
+        # A single file, or a list of image paths
+        elif os.path.isfile(image_file_names):
             list_file = image_file_names
-            with open(list_file) as f:
-                image_file_names = json.load(f)
-            print('Loaded {} image filenames from list file {}'.format(len(image_file_names),list_file))
-            
-        # A single image file
-        elif os.path.isfile(image_file_names) and path_utils.is_image_file(image_file_names):
-            image_file_names = [image_file_names]
-            print('Processing image {}'.format(image_file_names[0]))
+            if image_file_names.endswith('.json'):
+                with open(list_file,'r') as f:
+                    image_file_names = json.load(f)
+                print('Loaded {} image filenames from .json list file {}'.format(
+                    len(image_file_names),list_file))
+            elif image_file_names.endswith('.txt'):
+                with open(list_file,'r') as f:
+                    image_file_names = f.readlines()
+                    image_file_names = [s.strip() for s in image_file_names if len(s.strip()) > 0]
+                print('Loaded {} image filenames from .txt list file {}'.format(
+                    len(image_file_names),list_file))
+            elif path_utils.is_image_file(image_file_names):
+                image_file_names = [image_file_names]
+                print('Processing image {}'.format(image_file_names[0]))
+            else:
+                raise ValueError(
+                    'File {} supplied as [image_file_names] argument, but extension is neither .json nor .txt'\
+                        .format(
+                        list_file))
+        else:            
+            raise ValueError(
+                '{} supplied as [image_file_names] argument, but it does not appear to be a file or folder')
             
-        else:        
-            raise ValueError('image_file_names is a string, but is not a directory, a json ' + \
-                             'list (.json), or an image file')
-    
     if results is None:
         results = []
 
@@ -718,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()
     
@@ -759,7 +850,8 @@ def main():
         help='Path to detector model file (.pb or .pt).  Can also be the strings "MDV4", "MDV5A", or "MDV5B" to request automatic download.')
     parser.add_argument(
         'image_file',
-        help='Path to a single image file, a JSON file containing a list of paths to images, or a directory')
+        help=\
+        'Path to a single image file, a .json or .txt file containing a list of paths to images, or a directory')
     parser.add_argument(
         'output_file',
         help='Path to output JSON results file, should end with a .json extension')
@@ -806,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',
@@ -828,7 +923,7 @@ def main():
         default=None,
         help='Use a non-default class mapping, supplied in a .json file with a dictionary mapping' + \
             'int-strings to strings.  This will also disable the addition of "1" to all category ' + \
-            'IDs, so your class mapping should start at zero.')
+            'IDs, so your class mapping should start at zero.  Can also be a YOLOv5 dataset.yaml file.')
     parser.add_argument(
         '--include_image_size',
         action='store_true',
@@ -863,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'
@@ -885,26 +980,42 @@ def main():
         else:
             raise ValueError('Illegal overwrite handling string {}'.format(args.overwrite_handling))
 
-    # This is an experimental hack to allow the use of non-MD YOLOv5 models through
-    # the same infrastructure; it disables the code that enforces MDv5-like class lists.
-    if args.class_mapping_filename is not None:
-        run_detector.USE_MODEL_NATIVE_CLASSES = True
-        with open(args.class_mapping_filename,'r') as f:
-            class_mapping = json.load(f)
-        print('Loaded custom class mapping:')
-        print(class_mapping)
-        run_detector.DEFAULT_DETECTOR_LABEL_MAP = class_mapping
+    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'
@@ -931,7 +1042,15 @@ def main():
     elif os.path.isfile(args.image_file) and args.image_file.endswith('.json'):
         with open(args.image_file) as f:
             image_file_names = json.load(f)
-        print('Loaded {} image filenames from list file {}'.format(
+        print('Loaded {} image filenames from .json list file {}'.format(
+            len(image_file_names),args.image_file))
+    
+    # A text list of image paths
+    elif os.path.isfile(args.image_file) and args.image_file.endswith('.txt'):
+        with open(args.image_file) as f:
+            image_file_names = f.readlines()
+            image_file_names = [fn.strip() for fn in image_file_names if len(fn.strip()) > 0]
+        print('Loaded {} image filenames from .txt list file {}'.format(
             len(image_file_names),args.image_file))
         
     # A single image file
@@ -945,14 +1064,7 @@ def main():
 
     assert len(image_file_names) > 0, 'Specified image_file does not point to valid image files'
     assert os.path.exists(image_file_names[0]), \
-        'The first image to be scored 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'
+        'The first image to be processed does not exist at {}'.format(image_file_names[0])
 
     # Test that we can write to the output_file's dir if checkpointing requested
     if args.checkpoint_frequency != -1:
@@ -961,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
@@ -988,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()