megadetector 10.0.13__py3-none-any.whl

megadetector/postprocessing/md_to_labelme.py

@@ -0,0 +1,351 @@
+"""
+
+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: # noqa
+
+* support variable confidence thresholds across classes
+* support classification data
+
+"""
+
+#%% Imports and constants
+
+import os
+import json
+import sys
+import argparse
+
+from tqdm import tqdm
+
+from multiprocessing.pool import Pool
+from multiprocessing.pool import ThreadPool
+from functools import partial
+
+from megadetector.visualization.visualization_utils import open_image
+from megadetector.utils.ct_utils import round_float
+from megadetector.utils.ct_utils import write_json
+from megadetector.detection.run_detector import DEFAULT_DETECTOR_LABEL_MAP, FAILURE_IMAGE_OPEN
+
+output_precision = 3
+default_confidence_threshold = 0.15
+
+
+#%% Functions
+
+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.
+
+    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 is None:
+        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
+    output_dict['version'] = '5.3.0a0'
+    output_dict['flags'] = {}
+    output_dict['shapes'] = []
+    output_dict['imagePath'] = image_base_name
+    output_dict['imageHeight'] = im['height']
+    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['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)
+        x1 = round_float(x0 + det['bbox'][2] * im['width'],output_precision)
+        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,
+                            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)
+
+    write_json(json_path,output_dict)
+
+# ...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,
+                  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.  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 (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
+    """
+
+    if extension_prefix is None:
+        extension_prefix = ''
+
+    # 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)
+
+    # 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
+                    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_OPEN
+
+            # ...if we need to read w/h information
+
+        # ...for each image
+
+    # ...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']):
+            _write_output_for_image(im,image_base,extension_prefix,md_results['info'],confidence_threshold,
+                                   md_results['detection_categories'],overwrite,verbose)
+    else:
+        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:
+            if pool is not None:
+                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')
+    confidence_threshold = 0.2
+    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
+
+def main(): # noqa
+
+    parser = argparse.ArgumentParser(
+        description='Convert MD output to labelme annotation format')
+    parser.add_argument(
+        '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',
+        help='Overwrite existing labelme .json files')
+
+    if len(sys.argv[1:]) == 0:
+        parser.print_help()
+        parser.exit()
+
+    args = parser.parse_args()
+
+    md_to_labelme(args.results_file,args.image_base,args.confidence_threshold,args.overwrite)
+
+if __name__ == '__main__':
+    main()

megadetector/postprocessing/md_to_wi.py

@@ -0,0 +1,41 @@
+"""
+
+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_taxonomy_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=args.base_folder)
+
+if __name__ == '__main__':
+    main()