megadetector 5.0.11__py3-none-any.whl → 5.0.12__py3-none-any.whl

megadetector/postprocessing/md_to_coco.py

@@ -0,0 +1,310 @@
+"""
+
+md_to_coco.py
+
+"Converts" MegaDetector output files to COCO format.  "Converts" is in quotes because
+this is an opinionated transformation that requires a confidence threshold.
+
+Does not currently handle classification information.
+
+"""
+
+#%% Constants and imports
+
+import os
+import json
+import uuid
+
+from tqdm import tqdm
+
+from megadetector.visualization import visualization_utils as vis_utils
+
+default_confidence_threshold = 0.15
+
+
+#%% Functions
+
+def md_to_coco(md_results_file,
+               coco_output_file=None,
+               image_folder=None,
+               confidence_threshold=default_confidence_threshold,
+               validate_image_sizes=False,
+               info=None,
+               preserve_nonstandard_metadata=True,
+               include_failed_images=True):
+    """
+    "Converts" MegaDetector output files to COCO format.  "Converts" is in quotes because
+    this is an opinionated transformation that requires a confidence threshold.
+    
+    A folder of images is required if width and height information are not available 
+    in the MD results file.
+
+    Args:
+        md_results_file (str): MD results .json file to convert to COCO format
+        coco_output_file (str, optional): COCO .json file to write; if this is None, we'll return 
+            a COCO-formatted dict, but won't write it to disk
+        image_folder (str, optional): folder of images, required if 'width' and 'height' are not
+            present in the MD results file (they are not required by the format)
+        confidence_threshold (float, optional): boxes below this confidence threshold will not be
+            included in the output data
+        validate_image_sizes (bool, optional): if this is True, we'll check the image sizes 
+            regardless of whether "width" and "height" are present in the MD results file.
+        info (dict, optional): arbitrary metadata to include in an "info" field in the COCO-formatted
+            output
+        preserve_nonstandard_metadata (bool, optional): if this is True, confidence will be preserved in a 
+            non-standard "conf" field in each annotation, and any random fields present in each image's data
+            (e.g. EXIF metadata) will be propagated to COCO output    
+        include_failed_images (boo, optional): if this is True, failed images will be propagated to COCO output 
+            with a non-empty "failure" field and no other fields, otherwise failed images will be skipped.
+    
+    Returns:
+        dict: the COCO data dict, identical to what's written to [coco_output_file] if [coco_output_file]
+        is not None.
+    """
+
+    with open(md_results_file,'r') as f:    
+        md_results = json.load(f)    
+        
+    coco_images = []
+    coco_annotations = []
+    
+    # im = md_results['images'][0]
+    for im in tqdm(md_results['images']):
+        
+        coco_im = {}
+        coco_im['id'] = im['file']
+        coco_im['file_name'] = im['file']
+        
+        # There is no concept of this in the COCO standard        
+        if 'failure' in im and im['failure'] is not None:
+            if include_failed_images:
+                coco_im['failure'] = im['failure']
+                coco_images.append(coco_im)
+            continue
+        
+        # Read/validate image size
+        w = None
+        h = None
+        
+        if ('width' not in im) or ('height' not in im) or validate_image_sizes:
+            if image_folder is None:
+                raise ValueError('Must provide an image folder when height/width need to be read from images')
+            image_file_abs = os.path.join(image_folder,im['file'])
+            pil_im = vis_utils.open_image(image_file_abs)
+            w = pil_im.width
+            h = pil_im.height
+            if validate_image_sizes:
+                if 'width' in im:
+                    assert im['width'] == w, 'Width mismatch for image {}'.format(im['file'])
+                if 'height' in im:
+                    assert im['height'] == h, 'Height mismatch for image {}'.format(im['file'])
+        else:
+            w = im['width']
+            h = im['height']
+            
+        coco_im['width'] = w
+        coco_im['height'] = h
+        
+        # Add other, non-standard fields to the output dict
+        if preserve_nonstandard_metadata:
+            for k in im.keys():
+                if k not in ('file','detections','width','height'):
+                    coco_im[k] = im[k]
+                
+        coco_images.append(coco_im)
+        
+        # detection = im['detections'][0]
+        for detection in im['detections']:
+            
+            # Skip below-threshold detections
+            if confidence_threshold is not None and detection['conf'] < confidence_threshold:
+                continue
+        
+            # Create an annotation
+            ann = {}        
+            ann['id'] = str(uuid.uuid1())
+            ann['image_id'] = coco_im['id']                
+                        
+            md_category_id = detection['category']
+            coco_category_id = int(md_category_id)
+            ann['category_id'] = coco_category_id
+            
+            # In very esoteric cases, we use the empty category (0) in MD-formatted output files
+            if md_category_id != '0':
+                
+                assert 'bbox' in detection,\
+                    'Oops: non-empty category with no bbox in {}'.format(im['file'])
+            
+                ann['bbox'] = detection['bbox']
+                # MegaDetector: [x,y,width,height] (normalized, origin upper-left)
+                # COCO: [x,y,width,height] (absolute, origin upper-left)
+                ann['bbox'][0] = ann['bbox'][0] * coco_im['width']
+                ann['bbox'][1] = ann['bbox'][1] * coco_im['height']
+                ann['bbox'][2] = ann['bbox'][2] * coco_im['width']
+                ann['bbox'][3] = ann['bbox'][3] * coco_im['height']                
+            
+            else:
+                
+                print('Warning: empty category annotation in file {}'.format(im['file']))
+                      
+            if preserve_nonstandard_metadata:
+                ann['conf'] = detection['conf']
+                
+            coco_annotations.append(ann)            
+        
+        # ...for each detection
+    
+    # ...for each image
+
+    output_dict = {}
+    
+    if info is not None:
+        output_dict['info'] = info
+    else:
+        output_dict['info'] = {'description':'Converted from MD results file {}'.format(md_results_file)}
+    output_dict['info']['confidence_threshold'] = confidence_threshold
+    
+    output_dict['images'] = coco_images
+    output_dict['annotations'] = coco_annotations
+    
+    output_dict['categories'] = []
+    
+    for md_category_id in md_results['detection_categories'].keys():
+        
+        coco_category_id = int(md_category_id)
+        coco_category = {'id':coco_category_id,
+                         'name':md_results['detection_categories'][md_category_id]}
+        output_dict['categories'].append(coco_category)
+        
+    if coco_output_file is not None:
+        with open(coco_output_file,'w') as f:
+            json.dump(output_dict,f,indent=1)
+    
+    return output_dict
+
+# def md_to_coco(...)
+
+
+#%% 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')
+    validate_image_sizes = True
+    confidence_threshold = 0.2
+    validate_image_sizes=True
+    info=None
+    preserve_nonstandard_metadata=True
+    include_failed_images=False
+    
+    
+    #%% Programmatic execution
+    
+    output_dict = md_to_coco(md_results_file,
+                   coco_output_file=coco_output_file,
+                   image_folder=image_folder,
+                   confidence_threshold=confidence_threshold,
+                   validate_image_sizes=validate_image_sizes,
+                   info=info,
+                   preserve_nonstandard_metadata=preserve_nonstandard_metadata,
+                   include_failed_images=include_failed_images)
+    
+    
+    #%% Command-line example
+        
+    s = f'python md_to_coco.py {md_results_file} {coco_output_file} {confidence_threshold} '
+    if image_folder is not None:
+        s += f' --image_folder {image_folder}'
+    if preserve_nonstandard_metadata:
+        s += ' --preserve_nonstandard_metadata'
+    if include_failed_images:
+        s += ' --include_failed_images'
+
+    print(s); import clipboard; clipboard.copy(s)
+
+    
+    #%% Preview the resulting file
+    
+    from megadetector.visualization import visualize_db
+    options = visualize_db.DbVizOptions()
+    options.parallelize_rendering = True
+    options.viz_size = (900, -1)
+    options.num_to_visualize = 5000
+
+    html_file,_ = visualize_db.visualize_db(coco_output_file,
+                                              os.path.expanduser('~/tmp/md_to_coco_preview'),
+                                              image_folder,options)    
+
+    from megadetector.utils import path_utils # noqa
+    path_utils.open_file(html_file)
+    
+    
+#%% Command-line driver
+
+import sys,argparse
+
+def main():
+
+    parser = argparse.ArgumentParser(
+        description='"Convert" MD output to COCO format, in quotes because this is an opinionated transformation that requires a confidence threshold')
+    
+    parser.add_argument(
+        'md_results_file',
+        type=str,
+        help='Path to MD results file (.json)')
+    
+    parser.add_argument(
+        'coco_output_file',
+        type=str,
+        help='Output filename (.json)')
+    
+    parser.add_argument(
+        'confidence_threshold',
+        type=float,
+        default=default_confidence_threshold,
+        help='Confidence threshold (default {})'.format(default_confidence_threshold)
+        )
+    
+    parser.add_argument(
+        '--image_folder',
+        type=str,
+        default=None,
+        help='Image folder, only required if we will need to access image sizes'
+        )
+    
+    parser.add_argument(
+        '--preserve_nonstandard_metadata',
+        action='store_true',
+        help='Preserve metadata that isn\'t normally included in ' + 
+             'COCO-formatted data (e.g. EXIF metadata, confidence values)'
+        )
+
+    parser.add_argument(
+        '--include_failed_images',
+        action='store_true',
+        help='Keep a record of corrupted images in the output; may not be completely COCO-compliant'
+        )
+    
+    if len(sys.argv[1:]) == 0:
+        parser.print_help()
+        parser.exit()
+
+    args = parser.parse_args()
+
+    md_to_coco(args.md_results_file,
+               args.coco_output_file,
+               args.image_folder,
+               args.confidence_threshold,
+               validate_image_sizes=False,
+               info=None,
+               preserve_nonstandard_metadata=args.preserve_nonstandard_metadata,
+               include_failed_images=args.include_failed_images)    
+    
+if __name__ == '__main__':
+    main()

megadetector/postprocessing/md_to_labelme.py

@@ -0,0 +1,330 @@
+"""
+
+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
+
+import os
+import json
+
+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 truncate_float
+from megadetector.detection.run_detector import DEFAULT_DETECTOR_LABEL_MAP
+
+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:
+        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 = truncate_float(det['bbox'][0] * im['width'],output_precision)
+        y0 = truncate_float(det['bbox'][1] * im['height'],output_precision)
+        x1 = truncate_float(x0 + det['bbox'][2] * im['width'],output_precision)
+        y1 = truncate_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)
+            
+    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,
+                  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    
+    """
+    
+    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 access'        
+    
+            # ...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:
+        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
+    
+# ...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
+
+import sys,argparse
+
+def main():
+
+    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()