megadetector 10.0.15__py3-none-any.whl

megadetector/visualization/visualize_db.py

@@ -0,0 +1,630 @@
+"""
+
+visualize_db.py
+
+Outputs an HTML page visualizing annotations (class labels and/or bounding boxes)
+on a sample of images in a database in the COCO Camera Traps format.
+
+"""
+
+#%% Imports
+
+import argparse
+import inspect
+import random
+import json
+import math
+import os
+import sys
+import time
+
+import pandas as pd
+import numpy as np
+
+import humanfriendly
+
+from itertools import compress
+from multiprocessing.pool import ThreadPool
+from multiprocessing.pool import Pool
+from tqdm import tqdm
+
+from megadetector.utils.write_html_image_list import write_html_image_list
+from megadetector.data_management.cct_json_utils import IndexedJsonDb
+from megadetector.visualization import visualization_utils as vis_utils
+
+def _isnan(x):
+    return (isinstance(x,float) and np.isnan(x))
+
+
+#%% Settings
+
+class DbVizOptions:
+    """
+    Parameters controlling the behavior of visualize_db().
+    """
+
+    def __init__(self):
+
+        #: Number of images to sample from the database, or None to visualize all images
+        self.num_to_visualize = None
+
+        #: Target size for rendering; set either dimension to -1 to preserve aspect ratio.
+        #:
+        #: If viz_size is None or (-1,-1), the original image size is used.
+        self.viz_size = (1000, -1)
+
+        #: HTML rendering options; see write_html_image_list for details
+        #:
+        #:The most relevant option one might want to set here is:
+        #:
+        #: html_options['maxFiguresPerHtmlFile']
+        #:
+        #: ...which can be used to paginate previews to a number of images that will load well
+        #: in a browser (5000 is a reasonable limit).
+        self.html_options = write_html_image_list()
+
+        #: Whether to sort images by filename (True) or randomly (False)
+        self.sort_by_filename = True
+
+        #: Only show images that contain bounding boxes
+        self.trim_to_images_with_bboxes = False
+
+        #: Random seed to use for sampling images
+        self.random_seed = 0
+
+        #: Should we include Web search links for each category name?
+        self.add_search_links = False
+
+        #: Should each thumbnail image link back to the original image?
+        self.include_image_links = False
+
+        #: Should there be a text link back to each original image?
+        self.include_filename_links = False
+
+        #: Line width in pixels
+        self.box_thickness = 4
+
+        #: Number of pixels to expand each bounding box
+        self.box_expansion = 0
+
+        #: Only include images that contain annotations with these class names (not IDs) (list)
+        #:
+        #: Mutually exclusive with classes_to_exclude
+        self.classes_to_include = None
+
+        #: Exclude images that contain annotations with these class names (not IDs) (list)
+        #:
+        #: Mutually exclusive with classes_to_include
+        self.classes_to_exclude = None
+
+        #: Special tag used to say "show me all images with multiple categories"
+        #:
+        #: :meta private:
+        self.multiple_categories_tag = '*multiple*'
+
+        #: Parallelize rendering across multiple workers
+        self.parallelize_rendering = False
+
+        #: In theory, whether to parallelize with threads (True) or processes (False), but
+        #: process-based parallelization in this function is currently unsupported
+        self.parallelize_rendering_with_threads = True
+
+        #: Number of workers to use for parallelization; ignored if parallelize_rendering
+        #: is False
+        self.parallelize_rendering_n_cores = 16
+
+        #: Should we show absolute (True) or relative (False) paths for each image?
+        self.show_full_paths = False
+
+        #: List of additional fields in the image struct that we should print in image headers
+        self.extra_image_fields_to_print = None
+
+        #: List of additional fields in the annotation struct that we should print in image headers
+        self.extra_annotation_fields_to_print = None
+
+        #: Set to False to skip existing images
+        self.force_rendering = True
+
+        #: Enable additionald debug console output
+        self.verbose = False
+
+        #: COCO files used for evaluation may contain confidence scores, this
+        #: determines the field name used for confidence scores
+        self.confidence_field_name = 'score'
+
+        #: Optionally apply a confidence threshold; this requires that [confidence_field_name]
+        #: be present in all detections.
+        self.confidence_threshold = None
+
+
+#%% Core functions
+
+def visualize_db(db_path, output_dir, image_base_dir, options=None):
+    """
+    Writes images and html to output_dir to visualize the images and annotations in a
+    COCO-formatted .json file.
+
+    Args:
+        db_path (str or dict): the .json filename to load, or a previously-loaded database
+        output_dir (str): the folder to which we should write annotated images
+        image_base_dir (str): the folder where the images live; filenames in [db_path] should
+            be relative to this folder.
+        options (DbVizOptions, optional): See DbVizOptions for details
+
+    Returns:
+        tuple: A length-two tuple containing (the html filename) and (the loaded database).
+    """
+
+    if options is None:
+        options = DbVizOptions()
+
+    # Consistency checking for fields with specific format requirements
+
+    # These should be a lists, but if someone specifies a string, do a reasonable thing
+    if isinstance(options.extra_image_fields_to_print,str):
+        options.extra_image_fields_to_print = [options.extra_image_fields_to_print]
+    if isinstance(options.extra_annotation_fields_to_print,str):
+        options.extra_annotation_fields_to_print = [options.extra_annotation_fields_to_print]
+
+    if not options.parallelize_rendering_with_threads:
+        print('Warning: process-based parallelization is not yet supported by visualize_db')
+        options.parallelize_rendering_with_threads = True
+
+    if image_base_dir.startswith('http'):
+        if not image_base_dir.endswith('/'):
+            image_base_dir += '/'
+    else:
+        assert(os.path.isdir(image_base_dir))
+
+    os.makedirs(os.path.join(output_dir, 'rendered_images'), exist_ok=True)
+
+    if isinstance(db_path,str):
+        assert(os.path.isfile(db_path))
+        print('Loading database from {}...'.format(db_path))
+        image_db = json.load(open(db_path))
+        print('...done, loaded {} images'.format(len(image_db['images'])))
+    elif isinstance(db_path,dict):
+        print('Using previously-loaded DB')
+        image_db = db_path
+    else:
+        raise ValueError('Illegal dictionary or filename')
+
+    annotations = image_db['annotations']
+    images = image_db['images']
+    categories = image_db['categories']
+
+    # Optionally remove all images without bounding boxes, *before* sampling
+    if options.trim_to_images_with_bboxes:
+
+        b_has_bbox = [False] * len(annotations)
+        for i_ann,ann in enumerate(annotations):
+            if 'bbox' in ann or 'bbox_relative' in ann:
+                if 'bbox' in ann:
+                    assert isinstance(ann['bbox'],list)
+                else:
+                    assert isinstance(ann['bbox_relative'],list)
+                b_has_bbox[i_ann] = True
+        annotations_with_boxes = list(compress(annotations, b_has_bbox))
+
+        image_ids_with_boxes = [x['image_id'] for x in annotations_with_boxes]
+        image_ids_with_boxes = set(image_ids_with_boxes)
+
+        image_has_box = [False] * len(images)
+        for i_image,image in enumerate(images):
+            image_id = image['id']
+            if image_id in image_ids_with_boxes:
+                image_has_box[i_image] = True
+        images_with_bboxes = list(compress(images, image_has_box))
+        images = images_with_bboxes
+
+    # Optionally include/remove images with specific labels, *before* sampling
+
+    assert (not ((options.classes_to_exclude is not None) and \
+                 (options.classes_to_include is not None))), \
+                 'Cannot specify an inclusion and exclusion list'
+
+    if options.classes_to_exclude is not None:
+        assert isinstance(options.classes_to_exclude,list), \
+            'If supplied, classes_to_exclude should be a list'
+
+    if options.classes_to_include is not None:
+        assert isinstance(options.classes_to_include,list), \
+            'If supplied, classes_to_include should be a list'
+
+    if (options.classes_to_exclude is not None) or (options.classes_to_include is not None):
+
+        print('Indexing database')
+        indexed_db = IndexedJsonDb(image_db)
+        b_valid_class = [True] * len(images)
+        for i_image,image in enumerate(images):
+            classes = indexed_db.get_classes_for_image(image)
+            if options.classes_to_exclude is not None:
+                for excluded_class in options.classes_to_exclude:
+                    if excluded_class in classes:
+                       b_valid_class[i_image] = False
+                       break
+            elif options.classes_to_include is not None:
+                b_valid_class[i_image] = False
+                if options.multiple_categories_tag in options.classes_to_include:
+                    if len(classes) > 1:
+                        b_valid_class[i_image] = True
+                if not b_valid_class[i_image]:
+                    for c in classes:
+                        if c in options.classes_to_include:
+                            b_valid_class[i_image] = True
+                            break
+            else:
+                raise ValueError('Illegal include/exclude combination')
+
+        images_with_valid_classes = list(compress(images, b_valid_class))
+        images = images_with_valid_classes
+
+    # ...if we need to include/exclude categories
+
+    # Put the annotations in a dataframe so we can select all annotations for a given image
+    print('Creating data frames')
+    df_anno = pd.DataFrame(annotations)
+    df_img = pd.DataFrame(images)
+
+    # Construct label map
+    label_map = {}
+    for cat in categories:
+        label_map[int(cat['id'])] = cat['name']
+
+    # Take a sample of images
+    if options.num_to_visualize is not None:
+        if options.num_to_visualize > len(df_img):
+            print('Warning: asked to visualize {} images, but only {} are available, keeping them all'.\
+                  format(options.num_to_visualize,len(df_img)))
+        else:
+            df_img = df_img.sample(n=options.num_to_visualize,random_state=options.random_seed)
+
+    images_html = []
+
+    # Set of dicts representing inputs to render_db_bounding_boxes:
+    #
+    # bboxes, box_classes, image_path
+    rendering_info = []
+
+    print('Preparing rendering list')
+
+    for i_image,img in tqdm(df_img.iterrows(),total=len(df_img)):
+
+        img_id = img['id']
+        assert img_id is not None
+
+        img_relative_path = img['file_name']
+
+        if image_base_dir.startswith('http'):
+            img_path = image_base_dir + img_relative_path
+        else:
+            img_path = os.path.join(image_base_dir,img_relative_path).replace('\\','/')
+
+        annos_i = df_anno.loc[df_anno['image_id'] == img_id, :] # all annotations on this image
+
+        bboxes = []
+        box_classes = []
+        box_score_strings = []
+
+        # All the class labels we've seen for this image (with or without bboxes)
+        image_categories = set()
+
+        extra_annotation_field_string = ''
+        annotation_level_for_image = ''
+
+        # Did this image come with already-normalized bounding boxes?
+        boxes_are_normalized = None
+
+        # Iterate over annotations for this image
+        # i_ann = 0; anno = annos_i.iloc[i_ann]
+        for i_ann,anno in annos_i.iterrows():
+
+            if options.extra_annotation_fields_to_print is not None:
+                field_names = list(anno.index)
+                for field_name in field_names:
+                    if field_name in options.extra_annotation_fields_to_print:
+                        field_value = anno[field_name]
+                        if (field_value is not None) and (not _isnan(field_value)):
+                            extra_annotation_field_string += ' ({}:{})'.format(
+                                field_name,field_value)
+
+            if options.confidence_threshold is not None:
+                assert options.confidence_field_name in anno, \
+                    'Error: confidence thresholding requested, ' + \
+                        'but at least one annotation does not have the {} field'.format(
+                            options.confidence_field_name)
+                if anno[options.confidence_field_name] < options.confidence_threshold:
+                    continue
+
+            if 'sequence_level_annotation' in anno:
+                b_sequence_level_annotation = anno['sequence_level_annotation']
+                if b_sequence_level_annotation:
+                    annotation_level = 'sequence'
+                else:
+                    annotation_level = 'image'
+                if annotation_level_for_image == '':
+                    annotation_level_for_image = annotation_level
+                elif annotation_level_for_image != annotation_level:
+                    annotation_level_for_image = 'mixed'
+
+            category_id = anno['category_id']
+            category_name = label_map[category_id]
+            if options.add_search_links:
+                category_name = category_name.replace('"','')
+                category_name = '<a href="https://www.google.com/search?tbm=isch&q={}">{}</a>'.format(
+                    category_name,category_name)
+
+            image_categories.add(category_name)
+
+            assert not ('bbox' in anno and 'bbox_relative' in anno), \
+                "An annotation can't have both an absolute and a relative bounding box"
+
+            box_field = 'bbox'
+            if 'bbox_relative' in anno:
+                box_field = 'bbox_relative'
+                assert (boxes_are_normalized is None) or (boxes_are_normalized), \
+                    "An image can't have both absolute and relative bounding boxes"
+                boxes_are_normalized = True
+            elif 'bbox' in anno:
+                assert (boxes_are_normalized is None) or (not boxes_are_normalized), \
+                    "An image can't have both absolute and relative bounding boxes"
+                boxes_are_normalized = False
+
+            if box_field in anno:
+                bbox = anno[box_field]
+                if isinstance(bbox,float):
+                    assert math.isnan(bbox), "I shouldn't see a bbox that's neither a box nor NaN"
+                    continue
+                bboxes.append(bbox)
+                box_classes.append(anno['category_id'])
+
+                box_score_string = ''
+                if options.confidence_field_name is not None and \
+                   options.confidence_field_name in anno:
+                       score = anno[options.confidence_field_name]
+                       box_score_string = '({}%)'.format(round(100 * score))
+                box_score_strings.append(box_score_string)
+
+        # ...for each of this image's annotations
+
+        image_classes = ', '.join(image_categories)
+
+        img_id_string = str(img_id).lower()
+        file_name = '{}_gt.jpg'.format(os.path.splitext(img_id_string)[0])
+
+        # Replace characters that muck up image links, including flattening file
+        # separators.
+        illegal_characters = ['/','\\',':','\t','#',' ','%']
+        for c in illegal_characters:
+            file_name = file_name.replace(c,'~')
+
+        rendering_info_this_image = {'bboxes':bboxes,
+                                     'box_classes':box_classes,
+                                     'tags':box_score_strings,
+                                     'img_path':img_path,
+                                     'output_file_name':file_name,
+                                     'boxes_are_normalized':boxes_are_normalized}
+        rendering_info.append(rendering_info_this_image)
+
+        label_level_string = ''
+        if len(annotation_level_for_image) > 0:
+            label_level_string = ' (annotation level: {})'.format(annotation_level_for_image)
+
+        if 'frame_num' in img and 'seq_num_frames' in img:
+            frame_string = ' frame: {} of {},'.format(img['frame_num'],img['seq_num_frames'])
+        elif 'frame_num' in img:
+            frame_string = ' frame: {},'.format(img['frame_num'])
+        else:
+            frame_string = ''
+
+        if options.show_full_paths:
+            filename_text = img_path
+        else:
+            filename_text = img_relative_path
+        if options.include_filename_links:
+            filename_text = '<a href="{}">{}</a>'.format(img_path,filename_text)
+
+        flag_string = ''
+
+        if ('flags' in img) and (not _isnan(img['flags'])):
+            flag_string = ', flags: {}'.format(str(img['flags']))
+
+        extra_field_string = ''
+
+        if options.extra_image_fields_to_print is not None:
+            for field_name in options.extra_image_fields_to_print:
+                if field_name in img:
+                    # Always include a leading comma; this either separates us from the
+                    # previous field in [extra_fields_to_print] or from the rest of the string
+                    extra_field_string += ', {}: {}'.format(
+                        field_name,str(img[field_name]))
+
+        # We're adding html for an image before we render it, so it's possible this image will
+        # fail to render.  For applications where this script is being used to debua a database
+        # (the common case?), this is useful behavior, for other applications, this is annoying.
+        image_dict = \
+        {
+            'filename': '{}/{}'.format('rendered_images', file_name),
+            'title': '{}<br/>{}, num boxes: {},{} class labels: {}{}{}{}{}'.format(
+                filename_text, img_id, len(bboxes), frame_string, image_classes,
+                label_level_string, flag_string, extra_field_string, extra_annotation_field_string),
+            'textStyle': 'font-family:verdana,arial,calibri;font-size:80%;' + \
+                'text-align:left;margin-top:20;margin-bottom:5'
+        }
+        if options.include_image_links:
+            image_dict['linkTarget'] = img_path
+
+        images_html.append(image_dict)
+
+    # ...for each image
+
+    def render_image_info(rendering_info):
+
+        img_path = rendering_info['img_path']
+        bboxes = rendering_info['bboxes']
+        bbox_classes = rendering_info['box_classes']
+        boxes_are_normalized = rendering_info['boxes_are_normalized']
+        bbox_tags = None
+        if 'tags' in rendering_info:
+            bbox_tags = rendering_info['tags']
+        output_file_name = rendering_info['output_file_name']
+        output_full_path = os.path.join(output_dir, 'rendered_images', output_file_name)
+
+        if (os.path.isfile(output_full_path)) and (not options.force_rendering):
+            if options.verbose:
+                print('Skipping existing image {}'.format(output_full_path))
+            return True
+
+        if not img_path.startswith('http'):
+            if not os.path.exists(img_path):
+                print('Image {} cannot be found'.format(img_path))
+                return False
+
+        try:
+            original_image = vis_utils.open_image(img_path)
+            original_size = original_image.size
+            if (options.viz_size is None) or \
+                (options.viz_size[0] == -1 and options.viz_size[1] == -1):
+                image = original_image
+            else:
+                image = vis_utils.resize_image(original_image,
+                                               options.viz_size[0],
+                                               options.viz_size[1],
+                                               no_enlarge_width=True)
+        except Exception as e:
+            print('Image {} failed to open, error: {}'.format(img_path, e))
+            return False
+
+        vis_utils.render_db_bounding_boxes(boxes=bboxes,
+                                           classes=bbox_classes,
+                                           image=image,
+                                           original_size=original_size,
+                                           label_map=label_map,
+                                           thickness=options.box_thickness,
+                                           expansion=options.box_expansion,
+                                           tags=bbox_tags,
+                                           boxes_are_normalized=boxes_are_normalized)
+
+        image.save(output_full_path)
+
+        return True
+
+    # ...def render_image_info(...)
+
+    print('Rendering images')
+    start_time = time.time()
+
+    if options.parallelize_rendering:
+
+        if options.parallelize_rendering_with_threads:
+            worker_string = 'threads'
+        else:
+            worker_string = 'processes'
+
+        pool = None
+        try:
+            if options.parallelize_rendering_n_cores is None:
+                if options.parallelize_rendering_with_threads:
+                    pool = ThreadPool()
+                else:
+                    pool = Pool()
+            else:
+                if options.parallelize_rendering_with_threads:
+                    pool = ThreadPool(options.parallelize_rendering_n_cores)
+                else:
+                    pool = Pool(options.parallelize_rendering_n_cores)
+                print('Rendering images with {} {}'.format(options.parallelize_rendering_n_cores,
+                                                           worker_string))
+            rendering_success = list(tqdm(pool.imap(render_image_info, rendering_info),
+                                     total=len(rendering_info)))
+        finally:
+            if pool is not None:
+                pool.close()
+                pool.join()
+                print("Pool closed and joined for DB visualization")
+
+    else:
+
+        rendering_success = []
+        for file_info in tqdm(rendering_info):
+            rendering_success.append(render_image_info(file_info))
+
+    elapsed = time.time() - start_time
+
+    print('Rendered {} images in {} ({} successful)'.format(
+        len(rendering_info),humanfriendly.format_timespan(elapsed),sum(rendering_success)))
+
+    if options.sort_by_filename:
+        images_html = sorted(images_html, key=lambda x: x['filename'])
+    else:
+        random.shuffle(images_html)
+
+    html_output_file = os.path.join(output_dir, 'index.html')
+
+    html_options = options.html_options
+    if isinstance(db_path,str):
+        html_options['headerHtml'] = '<h1>Sample annotations from {}</h1>'.format(db_path)
+    else:
+        html_options['headerHtml'] = '<h1>Sample annotations</h1>'
+
+    write_html_image_list(
+            filename=html_output_file,
+            images=images_html,
+            options=html_options)
+
+    print('Visualized {} images, wrote results to {}'.format(len(images_html),html_output_file))
+
+    return html_output_file,image_db
+
+# ...def visualize_db(...)
+
+
+#%% Command-line driver
+
+# Copy all fields from a Namespace (i.e., the output from parse_args) to an object.
+#
+# Skips fields starting with _.  Does not check existence in the target object.
+def _args_to_object(args, obj):
+
+    for n, v in inspect.getmembers(args):
+        if not n.startswith('_'):
+            setattr(obj, n, v)
+
+def main():
+    """
+    Command-line driver for visualize_db
+    """
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('db_path', action='store', type=str,
+                        help='.json file to visualize')
+    parser.add_argument('output_dir', action='store', type=str,
+                        help='Output directory for html and rendered images')
+    parser.add_argument('image_base_dir', action='store', type=str,
+                        help='Base directory (or URL) for input images')
+
+    parser.add_argument('--num_to_visualize', action='store', type=int, default=None,
+                        help='Number of images to visualize (randomly drawn) (defaults to all)')
+    parser.add_argument('--random_sort', action='store_true',
+                        help='Sort randomly (rather than by filename) in output html')
+    parser.add_argument('--trim_to_images_with_bboxes', action='store_true',
+                        help='Only include images with bounding boxes (defaults to false)')
+    parser.add_argument('--random_seed', action='store', type=int, default=None,
+                        help='Random seed for image selection')
+
+    if len(sys.argv[1:]) == 0:
+        parser.print_help()
+        parser.exit()
+
+    args = parser.parse_args()
+
+    # Convert to an options object
+    options = DbVizOptions()
+    _args_to_object(args, options)
+    if options.random_sort:
+        options.sort_by_filename = False
+
+    visualize_db(options.db_path,options.output_dir,options.image_base_dir,options)
+
+if __name__ == '__main__':
+    main()