megadetector 5.0.28__py3-none-any.whl → 10.0.0__py3-none-any.whl

megadetector/postprocessing/md_to_labelme.py

@@ -3,11 +3,11 @@
 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 
+file).  "Convert" is in quotes because this is an opinionated transformation that
 requires a confidence threshold.
 
-TODO:
-   
+TODO: # noqa
+
 * support variable confidence thresholds across classes
 * support classification data
 
@@ -17,6 +17,8 @@ TODO:
 
 import os
 import json
+import sys
+import argparse
 
 from tqdm import tqdm
 
@@ -39,10 +41,10 @@ def get_labelme_dict_for_image(im,image_base_name=None,category_id_to_name=None,
     """
     For the given image struct in MD results format, reformat the detections into
     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; 
+        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
@@ -50,20 +52,20 @@ def get_labelme_dict_for_image(im,image_base_name=None,category_id_to_name=None,
             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:        
+
+    if confidence_threshold is None:
         confidence_threshold = -1.0
-     
+
     output_dict = {}
     if info is not None:
         output_dict['detector_info'] = info
@@ -75,22 +77,22 @@ def get_labelme_dict_for_image(im,image_base_name=None,category_id_to_name=None,
     output_dict['imageWidth'] = im['width']
     output_dict['imageData'] = None
     output_dict['detections'] = im['detections']
-    
+
     # det = im['detections'][1]
     for det in im['detections']:
-        
+
         if det['conf'] < confidence_threshold:
             continue
-        
+
         shape = {}
         shape['conf'] = det['conf']
-        shape['label'] = category_id_to_name[det['category']] 
+        shape['label'] = category_id_to_name[det['category']]
         shape['shape_type'] = 'rectangle'
         shape['description'] = ''
         shape['group_id'] = None
-        
+
         # MD boxes are [x_min, y_min, width_of_box, height_of_box] (relative)
-        # 
+        #
         # labelme boxes are [[x0,y0],[x1,y1]] (absolute)
         x0 = round_float(det['bbox'][0] * im['width'],output_precision)
         y0 = round_float(det['bbox'][1] * im['height'],output_precision)
@@ -98,28 +100,33 @@ def get_labelme_dict_for_image(im,image_base_name=None,category_id_to_name=None,
         y1 = round_float(y0 + det['bbox'][3] * im['height'],output_precision)
         shape['points'] = [[x0,y0],[x1,y1]]
         output_dict['shapes'].append(shape)
-    
+
     # ...for each detection
-    
+
     return output_dict
 
 # ...def get_labelme_dict_for_image()
 
 
-def _write_output_for_image(im,image_base,extension_prefix,info,
-                            confidence_threshold,category_id_to_name,overwrite,
+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))
@@ -130,43 +137,49 @@ def _write_output_for_image(im,image_base,extension_prefix,info,
                                              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,extension_prefix='',n_workers=1,
-                  use_threads=False,bypass_image_size_read=False,
+def md_to_labelme(results_file,
+                  image_base,
+                  confidence_threshold=None,
+                  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
+            will be included in the output dict.  If None, no threshold will be applied.
         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
+        extension_prefix (str, optional): if non-empty, "extension_prefix" will be inserted before the .json
+            extension (typically used to generate multiple copies of labelme files representing different
+            MD thresholds)
         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    
+        verbose (bool, optional): enables additionald ebug output
     """
-    
+
     if extension_prefix is None:
         extension_prefix = ''
-        
+
     # Load MD results if necessary
     if isinstance(results_file,dict):
         md_results = results_file
@@ -174,34 +187,34 @@ def md_to_labelme(results_file,image_base,confidence_threshold=None,
         print('Loading MD results...')
         with open(results_file,'r') as f:
             md_results = json.load(f)
-        
-    # Read image sizes if necessary            
-    if bypass_image_size_read:     
-        
+
+    # Read image sizes if necessary
+    if bypass_image_size_read:
+
         print('Bypassing image size read')
-        
+
     else:
-    
+
         # 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
-            
+
             # 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
@@ -211,88 +224,92 @@ def md_to_labelme(results_file,image_base,confidence_threshold=None,
                         im_full_path))
                     if 'failure' not in im:
                         im['failure'] = FAILURE_IMAGE_OPEN
-    
+
             # ...if we need to read w/h information
-            
+
         # ...for each image
-        
-    # ...if we're not bypassing image size read        
-    
+
+    # ...if we're not bypassing image size read
+
     print('\nGenerating labelme files...')
-        
+
     # Write output
     if n_workers <= 1:
-        for im in tqdm(md_results['images']):    
+        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'])))
-            
+        pool = None
+        try:
+            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'])))
+        finally:
+            pool.close()
+            pool.join()
+            print("Pool closed and joined for labelme file writes")
+
     # ...for each image
-    
+
 # ...def md_to_labelme()
 
 
 #%% Interactive driver
 
 if False:
-    
+
     pass
 
     #%% Configure options
-    
+
     md_results_file = os.path.expanduser('~/data/md-test.json')
     coco_output_file = os.path.expanduser('~/data/md-test-coco.json')
-    image_folder = os.path.expanduser('~/data/md-test')    
+    image_folder = os.path.expanduser('~/data/md-test')
     confidence_threshold = 0.2
-    overwrite = True    
-    
-    
+    overwrite = True
+
+
     #%% Programmatic execution
-    
+
     md_to_labelme(results_file=md_results_file,
                   image_base=image_folder,
                   confidence_threshold=confidence_threshold,
                   overwrite=overwrite)
 
-    
+
     #%% Command-line execution
-    
+
     s = 'python md_to_labelme.py {} {} --confidence_threshold {}'.format(md_results_file,
                                                                          image_folder,
                                                                          confidence_threshold)
     if overwrite:
         s += ' --overwrite'
-        
+
     print(s)
     import clipboard; clipboard.copy(s)
 
 
     #%% Opening labelme
-    
+
     s = 'python labelme {}'.format(image_folder)
     print(s)
     import clipboard; clipboard.copy(s)
-    
 
-#%% Command-line driver
 
-import sys,argparse
+#%% Command-line driver
 
-def main():
+def main(): # noqa
 
     parser = argparse.ArgumentParser(
         description='Convert MD output to labelme annotation format')
@@ -300,19 +317,19 @@ def main():
         'results_file',
         type=str,
         help='Path to MD results file (.json)')
-    
+
     parser.add_argument(
         'image_base',
         type=str,
         help='Path to images (also the output folder)')
-    
+
     parser.add_argument(
         '--confidence_threshold',
         type=float,
         default=default_confidence_threshold,
         help='Confidence threshold (default {})'.format(default_confidence_threshold)
         )
-    
+
     parser.add_argument(
         '--overwrite',
         action='store_true',

megadetector/postprocessing/md_to_wi.py

@@ -1,40 +1,40 @@
-"""
-
-md_to_wi.py
-
-Converts the MD .json format to the WI predictions.json format.
-
-"""
-
-#%% Imports and constants
-
-import sys
-import argparse
-from megadetector.utils.wi_utils import generate_predictions_json_from_md_results
-
-
-#%% Command-line driver
-
-def main():
-    
-    parser = argparse.ArgumentParser()
-    parser.add_argument('md_results_file', action='store', type=str, 
-                        help='output file in MD format to convert')
-    parser.add_argument('predictions_json_file', action='store', type=str, 
-                        help='.json file to write in predictions.json format')
-    parser.add_argument('--base_folder', action='store', type=str, default=None, 
-                        help='folder name to prepend to each path in md_results_file, ' + \
-                             'to convert relative paths to absolute paths.')
-
-    if len(sys.argv[1:]) == 0:
-        parser.print_help()
-        parser.exit()
-            
-    args = parser.parse_args()
-    
-    generate_predictions_json_from_md_results(args.md_results_file,
-                                              args.predictions_json_file,
-                                              base_folder=None)
-
-if __name__ == '__main__':    
-    main()
+"""
+
+md_to_wi.py
+
+Converts the MD .json format to the WI predictions.json format.
+
+"""
+
+#%% Imports and constants
+
+import sys
+import argparse
+from megadetector.utils.wi_utils import generate_predictions_json_from_md_results
+
+
+#%% Command-line driver
+
+def main(): # noqa
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('md_results_file', action='store', type=str,
+                        help='output file in MD format to convert')
+    parser.add_argument('predictions_json_file', action='store', type=str,
+                        help='.json file to write in predictions.json format')
+    parser.add_argument('--base_folder', action='store', type=str, default=None,
+                        help='folder name to prepend to each path in md_results_file, ' + \
+                             'to convert relative paths to absolute paths.')
+
+    if len(sys.argv[1:]) == 0:
+        parser.print_help()
+        parser.exit()
+
+    args = parser.parse_args()
+
+    generate_predictions_json_from_md_results(args.md_results_file,
+                                              args.predictions_json_file,
+                                              base_folder=None)
+
+if __name__ == '__main__':
+    main()