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

api/batch_processing/postprocessing/md_to_labelme.py

@@ -1,17 +1,17 @@
-########
-#
-# md_to_labelme.py
-#
-# "Converts" a MegaDetector output .json file to labelme format (one .json per image
-# file).  "Convert" is in quotes because this is an opinionated transformation that 
-# requires a confidence threshold.
-#
-# TODO:
-#    
-# * support variable confidence thresholds across classes
-# * support classification data
-#
-########
+"""
+
+md_to_labelme.py
+
+"Converts" a MegaDetector output .json file to labelme format (one .json per image
+file).  "Convert" is in quotes because this is an opinionated transformation that 
+requires a confidence threshold.
+
+TODO:
+   
+* support variable confidence thresholds across classes
+* support classification data
+
+"""
 
 #%% Imports and constants
 
@@ -20,8 +20,13 @@ import json
 
 from tqdm import tqdm
 
+from multiprocessing.pool import Pool
+from multiprocessing.pool import ThreadPool
+from functools import partial
+
 from md_visualization.visualization_utils import open_image
 from md_utils.ct_utils import truncate_float
+from detection.run_detector import DEFAULT_DETECTOR_LABEL_MAP
 
 output_precision = 3
 default_confidence_threshold = 0.15
@@ -29,15 +34,36 @@ default_confidence_threshold = 0.15
 
 #%% Functions
 
-def get_labelme_dict_for_image(im,image_base_name,category_id_to_name,info=None,confidence_threshold=None):
+def get_labelme_dict_for_image(im,image_base_name=None,category_id_to_name=None,
+                               info=None,confidence_threshold=None):
     """
     For the given image struct in MD results format, reformat the detections into
-    labelme format.  Returns a dict.
+    labelme format.
+    
+    Args:
+        im (dict): MegaDetector-formatted results dict, must include 'height' and 'width' fields
+        image_base_name (str, optional): written directly to the 'imagePath' field in the output; 
+            defaults to os.path.basename(im['file']).
+        category_id_to_name (dict, optional): maps string-int category IDs to category names, defaults
+            to the standard MD categories
+        info (dict, optional): arbitrary metadata to write to the "detector_info" field in the output
+            dict
+        confidence_threshold (float, optional): only detections at or above this confidence threshold
+            will be included in the output dict
+    
+    Return:
+        dict: labelme-formatted dictionary, suitable for writing directly to a labelme-formatted .json file
     """
     
+    if image_base_name is None:
+        image_base_name = os.path.basename(im['file'])
+        
+    if category_id_to_name:
+        category_id_to_name = DEFAULT_DETECTOR_LABEL_MAP
+        
     if confidence_threshold is None:        
         confidence_threshold = -1.0
-        
+     
     output_dict = {}
     if info is not None:
         output_dict['detector_info'] = info
@@ -50,6 +76,7 @@ def get_labelme_dict_for_image(im,image_base_name,category_id_to_name,info=None,
     output_dict['imageData'] = None
     output_dict['detections'] = im['detections']
     
+    # det = im['detections'][1]
     for det in im['detections']:
         
         if det['conf'] < confidence_threshold:
@@ -79,69 +106,140 @@ def get_labelme_dict_for_image(im,image_base_name,category_id_to_name,info=None,
 # ...def get_labelme_dict_for_image()
 
 
+def _write_output_for_image(im,image_base,extension_prefix,info,
+                            confidence_threshold,category_id_to_name,overwrite,
+                            verbose=False):
+    
+    if 'failure' in im and im['failure'] is not None:
+        assert 'detections' not in im or im['detections'] is None
+        if verbose:
+            print('Skipping labelme file generation for failed image {}'.format(
+                im['file']))
+        return
+        
+    im_full_path = os.path.join(image_base,im['file'])
+    json_path = os.path.splitext(im_full_path)[0] + extension_prefix + '.json'
+    
+    if (not overwrite) and (os.path.isfile(json_path)):
+        if verbose:
+            print('Skipping existing file {}'.format(json_path))
+        return
+
+    output_dict = get_labelme_dict_for_image(im,
+                                             image_base_name=os.path.basename(im_full_path),
+                                             category_id_to_name=category_id_to_name,
+                                             info=info,
+                                             confidence_threshold=confidence_threshold)
+            
+    with open(json_path,'w') as f:
+        json.dump(output_dict,f,indent=1)
+    
+# ...def write_output_for_image(...)
+
+
+
 def md_to_labelme(results_file,image_base,confidence_threshold=None,
-                  overwrite=False):
+                  overwrite=False,extension_prefix='',n_workers=1,
+                  use_threads=False,bypass_image_size_read=False,
+                  verbose=False):
     """
     For all the images in [results_file], write a .json file in labelme format alongside the
     corresponding relative path within image_base.
+    
+    Args:
+        results_file (str): MD results .json file to convert to Labelme format
+        image_base (str): folder of images; filenames in [results_file] should be relative to
+            this folder
+        confidence_threshold (float, optional): only detections at or above this confidence threshold
+            will be included in the output dict
+        overwrite (bool, optional): whether to overwrite existing output files; if this is False
+            and the output file for an image exists, we'll skip that image
+        extension_prefix (str, optional): if non-empty, "extension_prefix" will be inserted before the .json 
+            extension
+        n_workers (int, optional): enables multiprocessing if > 1
+        use_threads (bool, optional): if [n_workers] > 1, determines whether we parallelize via threads (True)
+            or processes (False)
+        bypass_image_size_read (bool, optional): if True, skips reading image sizes and trusts whatever is in
+            the MD results file (don't set this to "True" if your MD results file doesn't contain image sizes)
+        verbose (bool, optional): enables additionald ebug output    
     """
     
-    # Load MD results
-    with open(results_file,'r') as f:
-        md_results = json.load(f)
+    if extension_prefix is None:
+        extension_prefix = ''
         
-    # Read image sizes
-    #
-    # TODO: parallelize this loop
-    #
-    # im = md_results['images'][0]
-    for im in tqdm(md_results['images']):
+    # Load MD results if necessary
+    if isinstance(results_file,dict):
+        md_results = results_file
+    else:
+        print('Loading MD results...')
+        with open(results_file,'r') as f:
+            md_results = json.load(f)
         
-        # Make sure this file exists
-        im_full_path = os.path.join(image_base,im['file'])
-        assert os.path.isfile(im_full_path), 'Image file {} does not exist'.format(im_full_path)
+    # Read image sizes if necessary            
+    if bypass_image_size_read:     
         
-        # Load w/h information if necessary
-        if 'height' not in im or 'width' not in im:
-            
-            try:
-                pil_im = open_image(im_full_path)
-                im['width'] = pil_im.width
-                im['height'] = pil_im.height
-            except Exception:
-                print('Warning: cannot open image {}, treating as a failure during inference'.format(
-                    im_full_path))
-                if 'failure' not in im:
-                    im['failure'] = 'Failure image access'        
-
-        # ...if we need to read w/h information
+        print('Bypassing image size read')
         
-    # ...for each image
+    else:
     
-    # Write output
-    for im in tqdm(md_results['images']):
-        
-        if 'failure' in im and im['failure'] is not None:
-            assert 'detections' not in im
-            print('Warning: skipping labelme file generation for failed image {}'.format(
-                im['file']))
-            continue
+        # TODO: parallelize this loop
+    
+        print('Reading image sizes...')
+                
+        # im = md_results['images'][0]
+        for im in tqdm(md_results['images']):
+            
+            # Make sure this file exists
+            im_full_path = os.path.join(image_base,im['file'])
+            assert os.path.isfile(im_full_path), 'Image file {} does not exist'.format(im_full_path)
+            
+            json_path = os.path.splitext(im_full_path)[0] + extension_prefix + '.json'
+            
+            # Don't even bother reading sizes for files we're not going to generate
+            if (not overwrite) and (os.path.isfile(json_path)):
+                continue
             
-        im_full_path = os.path.join(image_base,im['file'])
-        json_path = os.path.splitext(im_full_path)[0] + '.json'
+            # Load w/h information if necessary
+            if 'height' not in im or 'width' not in im:
+                
+                try:
+                    pil_im = open_image(im_full_path)
+                    im['width'] = pil_im.width
+                    im['height'] = pil_im.height
+                except Exception:
+                    print('Warning: cannot open image {}, treating as a failure during inference'.format(
+                        im_full_path))
+                    if 'failure' not in im:
+                        im['failure'] = 'Failure image access'        
+    
+            # ...if we need to read w/h information
+            
+        # ...for each image
         
-        if (not overwrite) and (os.path.isfile(json_path)):
-            print('Skipping existing file {}'.format(json_path))
-            continue
+    # ...if we're not bypassing image size read        
     
-        output_dict = get_labelme_dict_for_image(im,
-                                                 image_base_name=os.path.basename(im_full_path),
-                                                 category_id_to_name=md_results['detection_categories'],
-                                                 info=md_results['info'],
-                                                 confidence_threshold=confidence_threshold)
-                
-        with open(json_path,'w') as f:
-            json.dump(output_dict,f,indent=1)
+    print('\nGenerating labelme files...')
+        
+    # Write output
+    if n_workers <= 1:
+        for im in tqdm(md_results['images']):    
+            _write_output_for_image(im,image_base,extension_prefix,md_results['info'],confidence_threshold,
+                                   md_results['detection_categories'],overwrite,verbose)
+    else:
+        if use_threads:
+            print('Starting parallel thread pool with {} workers'.format(n_workers))
+            pool = ThreadPool(n_workers)
+        else:
+            print('Starting parallel process pool with {} workers'.format(n_workers))
+            pool = Pool(n_workers)
+        _ = list(tqdm(pool.imap(
+                partial(_write_output_for_image,
+                        image_base=image_base,extension_prefix=extension_prefix,
+                        info=md_results['info'],confidence_threshold=confidence_threshold,
+                        category_id_to_name=md_results['detection_categories'],
+                        overwrite=overwrite,verbose=verbose),
+                 md_results['images']),
+                 total=len(md_results['images'])))
             
     # ...for each image
     
@@ -227,7 +325,6 @@ def main():
     args = parser.parse_args()
 
     md_to_labelme(args.results_file,args.image_base,args.confidence_threshold,args.overwrite)
-    
-    
+
 if __name__ == '__main__':
     main()

api/batch_processing/postprocessing/merge_detections.py

@@ -1,14 +1,17 @@
-########
-#
-# merge_detections.py
-#
-# Merge high-confidence detections from one or more results files into another 
-# file.   Typically used to combine results from MDv5b and/or MDv4 into a "primary"
-# results file from MDv5a.
-#
-# If you want to literally merge two .json files, see combine_api_outputs.py.
-#
-########
+"""
+
+merge_detections.py
+
+Merge high-confidence detections from one or more results files into another 
+file.  Typically used to combine results from MDv5b and/or MDv4 into a "primary"
+results file from MDv5a.
+
+Detection categories must be the same in both files; if you want to first remap
+one file's category mapping to be the same as another's, see remap_detection_categories.
+
+If you want to literally merge two .json files, see combine_api_outputs.py.
+
+"""
 
 #%% Constants and imports
 
@@ -30,7 +33,7 @@ class MergeDetectionsOptions:
         
         self.max_detection_size = 1.01
         self.min_detection_size = 0
-        self.source_confidence_thresholds = [0.2]
+        self.source_confidence_thresholds = [0.05]
         
         # Don't bother merging into target images if there is a similar detection
         # above this threshold (or if there is *any* detection above this threshold,
@@ -38,7 +41,7 @@ class MergeDetectionsOptions:
         self.target_confidence_threshold = 0.2
         
         # If you want to merge only certain categories, specify one
-        # (but not both) of these.
+        # (but not both) of these.  These are category IDs, not names.
         self.categories_to_include = None
         self.categories_to_exclude = None
 
@@ -47,11 +50,28 @@ class MergeDetectionsOptions:
         self.merge_empty_only = False
         
         self.iou_threshold = 0.65
+        
+        self.overwrite = False
 
 
 #%% Main function
 
 def merge_detections(source_files,target_file,output_file,options=None):
+    """
+    Merge high-confidence detections from one or more results files into another 
+    file.   Typically used to combine results from MDv5b and/or MDv4 into a "primary"
+    results file from MDv5a.
+    
+    [source_files] (a list of files or a single filename) specifies the set of 
+    results files that will be merged into [target_file].  The difference between a 
+    "source file" and the "target file" is that if no merging is necessary, either because
+    two boxes are nearly identical or because merge_only_empty is True and the target
+    file already has above-threshold detection for an image+category, the output file gets
+    the results of the "target" file.  I.e., the "target" file wins all ties.
+    
+    The results are written to [output_file].
+
+    """
     
     if isinstance(source_files,str):
         source_files = [source_files]    
@@ -59,6 +79,10 @@ def merge_detections(source_files,target_file,output_file,options=None):
     if options is None:
         options = MergeDetectionsOptions()    
         
+    if (not options.overwrite) and (os.path.isfile(output_file)):
+        print('File {} exists, bypassing merge'.format(output_file))
+        return
+    
     assert not ((options.categories_to_exclude is not None) and \
                 (options.categories_to_include is not None)), \
                 'categories_to_include and categories_to_exclude are mutually exclusive'
@@ -133,7 +157,8 @@ def merge_detections(source_files,target_file,output_file,options=None):
         output_data['info']['detections_transferred_from'].append(os.path.basename(source_file))
         output_data['info']['detector'] = output_data['info']['detector'] + ' + ' + source_detector_name
         
-        assert source_data['detection_categories'] == output_data['detection_categories']
+        assert source_data['detection_categories'] == output_data['detection_categories'], \
+            'Cannot merge files with different detection category maps'
         
         source_confidence_threshold = options.source_confidence_thresholds[i_source_file]
         
@@ -246,7 +271,7 @@ def merge_detections(source_files,target_file,output_file,options=None):
     # ...for each source file        
     
     with open(output_file,'w') as f:
-        json.dump(output_data,f,indent=2)
+        json.dump(output_data,f,indent=1)
     
     print('Saved merged results to {}'.format(output_file))