megadetector 5.0.22__py3-none-any.whl → 5.0.24__py3-none-any.whl

megadetector/visualization/visualization_utils.py

@@ -15,7 +15,7 @@ import os
 import cv2
 
 from io import BytesIO
-from PIL import Image, ImageFile, ImageFont, ImageDraw
+from PIL import Image, ImageFile, ImageFont, ImageDraw, ImageFilter
 from multiprocessing.pool import ThreadPool
 from multiprocessing.pool import Pool
 from tqdm import tqdm
@@ -25,6 +25,7 @@ from megadetector.utils.path_utils import find_images
 from megadetector.data_management.annotations import annotation_constants
 from megadetector.data_management.annotations.annotation_constants import \
     detector_bbox_category_id_to_name
+from megadetector.utils.ct_utils import sort_list_of_dicts_by_key
 
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 
@@ -406,11 +407,46 @@ def crop_image(detections, image, confidence_threshold=0.15, expansion=0):
     return ret_images
 
 
+def blur_detections(image,detections,blur_radius=40):
+    """
+    Blur the regions in [image] corresponding to the MD-formatted list [detections].
+    [image] is modified in place.
+    
+    Args:
+        image (PIL.Image.Image): image in which we should blur specific regions
+        detections (list): list of detections in the MD output format, see render
+            detection_bounding_boxes for more detail.
+    """
+    
+    img_width, img_height = image.size
+    
+    for d in detections:
+        
+        bbox = d['bbox']
+        x_norm, y_norm, width_norm, height_norm = bbox
+
+        # Calculate absolute pixel coordinates
+        x = int(x_norm * img_width)
+        y = int(y_norm * img_height)
+        width = int(width_norm * img_width)
+        height = int(height_norm * img_height)
+        
+        # Calculate box boundaries
+        left = max(0, x)
+        top = max(0, y)
+        right = min(img_width, x + width)
+        bottom = min(img_height, y + height)
+        
+        # Crop the region, blur it, and paste it back
+        region = image.crop((left, top, right, bottom))
+        blurred_region = region.filter(ImageFilter.GaussianBlur(radius=blur_radius))
+        image.paste(blurred_region, (left, top))        
+        
 def render_detection_bounding_boxes(detections, 
                                     image,
                                     label_map='show_categories',
                                     classification_label_map=None, 
-                                    confidence_threshold=0.15, 
+                                    confidence_threshold=0, 
                                     thickness=DEFAULT_BOX_THICKNESS, 
                                     expansion=0,
                                     classification_confidence_threshold=0.3,
@@ -419,7 +455,9 @@ def render_detection_bounding_boxes(detections,
                                     textalign=TEXTALIGN_LEFT,
                                     vtextalign=VTEXTALIGN_TOP,
                                     label_font_size=DEFAULT_LABEL_FONT_SIZE,
-                                    custom_strings=None):
+                                    custom_strings=None,
+                                    box_sort_order='confidence',
+                                    verbose=False):
     """
     Renders bounding boxes (with labels and confidence values) on an image for all
     detections above a threshold.
@@ -497,6 +535,9 @@ def render_detection_bounding_boxes(detections,
         label_font_size (float, optional): font size for labels        
         custom_strings: optional set of strings to append to detection labels, should have the
             same length as [detections].  Appended before any classification labels.
+        box_sort_order (str, optional): sorting scheme for detection boxes, can be None, "confidence", or 
+            "reverse_confidence".
+        verbose (bool, optional): enable additional debug output
     """
 
     # Input validation
@@ -516,6 +557,15 @@ def render_detection_bounding_boxes(detections,
     # for color selection
     classes = []  
 
+    if box_sort_order is not None:
+        
+        if box_sort_order == 'confidence':            
+            detections = sort_list_of_dicts_by_key(detections,k='conf',reverse=False)
+        elif box_sort_order == 'reverse_confidence':
+            detections = sort_list_of_dicts_by_key(detections,k='conf',reverse=True)
+        else:
+            raise ValueError('Unrecognized sorting scheme {}'.format(box_sort_order))
+            
     for i_detection,detection in enumerate(detections):
 
         score = detection['conf']
@@ -536,8 +586,9 @@ def render_detection_bounding_boxes(detections,
             # category or on the most confident classification category.
             clss = detection['category']
             
-            # {} is the default, which means "show labels with no mapping", so don't use "if label_map" here
-            # if label_map:                
+            # This will be a list of strings that should be rendered above/below this box            
+            displayed_label = []
+            
             if label_map is not None:
                 label = label_map[clss] if clss in label_map else clss
                 if score is not None:
@@ -545,17 +596,14 @@ def render_detection_bounding_boxes(detections,
                 else:
                     displayed_label = ['{}'.format(label)]
             else:
-                displayed_label = ''
+                displayed_label = ['']
 
             if custom_strings is not None:
                 custom_string = custom_strings[i_detection]
                 if custom_string is not None and len(custom_string) > 0:
-                    if isinstance(displayed_label,str):
-                        displayed_label += ' ' + custom_string
-                    else:
-                        assert len(displayed_label) == 1
-                        displayed_label[0] += ' ' + custom_string
-                                    
+                    assert len(displayed_label) == 1
+                    displayed_label[0] += ' ' + custom_string
+
             if ('classifications' in detection) and len(detection['classifications']) > 0:
 
                 classifications = detection['classifications']
@@ -597,6 +645,7 @@ def render_detection_bounding_boxes(detections,
                 
             # ...if we have classification results
                         
+            # display_strs is a list of labels for each box
             display_strs.append(displayed_label)
             classes.append(clss)
 
@@ -606,6 +655,9 @@ def render_detection_bounding_boxes(detections,
     
     display_boxes = np.array(display_boxes)
 
+    if verbose:
+        print('Rendering {} of {} detections'.format(len(display_boxes),len(detections)))
+        
     draw_bounding_boxes_on_image(image, display_boxes, classes,
                                  display_strs=display_strs, thickness=thickness, 
                                  expansion=expansion, colormap=colormap, 
@@ -953,7 +1005,8 @@ def render_db_bounding_boxes(boxes,
                              textalign=TEXTALIGN_LEFT,
                              vtextalign=VTEXTALIGN_TOP,
                              text_rotation=None,
-                             label_font_size=DEFAULT_LABEL_FONT_SIZE):
+                             label_font_size=DEFAULT_LABEL_FONT_SIZE,
+                             tags=None):
     """
     Render bounding boxes (with class labels) on an image.  This is a wrapper for
     draw_bounding_boxes_on_image, allowing the caller to operate on a resized image
@@ -981,6 +1034,8 @@ def render_db_bounding_boxes(boxes,
         vtextalign (int, optional): VTEXTALIGN_TOP or VTEXTALIGN_BOTTOM
         text_rotation (float, optional): rotation to apply to text
         label_font_size (float, optional): font size for labels
+        tags (list, optional): list of strings of length len(boxes) that should be appended
+            after each class name (e.g. to show scores)
     """
 
     display_boxes = []
@@ -993,8 +1048,11 @@ def render_db_bounding_boxes(boxes,
 
     img_width, img_height = image_size
 
-    for box, clss in zip(boxes, classes):
+    for i_box in range(0,len(boxes)):
 
+        box = boxes[i_box]
+        clss = classes[i_box]
+        
         x_min_abs, y_min_abs, width_abs, height_abs = box[0:4]
 
         ymin = y_min_abs / img_height
@@ -1008,9 +1066,17 @@ def render_db_bounding_boxes(boxes,
         if label_map:
             clss = label_map[int(clss)]
             
+        display_str = str(clss)
+    
+        # Do we have a tag to append to the class string?
+        if tags is not None and tags[i_box] is not None and len(tags[i_box]) > 0:
+            display_str += ' ' + tags[i_box]
+            
         # need to be a string here because PIL needs to iterate through chars
-        display_strs.append([str(clss)])  
+        display_strs.append([display_str])
 
+    # ...for each box
+    
     display_boxes = np.array(display_boxes)
     
     draw_bounding_boxes_on_image(image, 
@@ -1706,7 +1772,7 @@ if False:
         TEXTALIGN_LEFT,TEXTALIGN_RIGHT,VTEXTALIGN_BOTTOM,VTEXTALIGN_TOP, \
         DEFAULT_LABEL_FONT_SIZE
         
-    fn = os.path.expanduser('~\AppData\Local\Temp\md-tests\md-test-images\ena24_7904.jpg')
+    fn = os.path.expanduser('~/AppData/Local/Temp/md-tests/md-test-images/ena24_7904.jpg')
     output_fn = r'g:\temp\test.jpg'
     
     image = load_image(fn)

megadetector/visualization/visualize_db.py

@@ -167,7 +167,7 @@ def visualize_db(db_path, output_dir, image_base_dir, options=None):
         options (DbVizOptions, optional): See DbVizOptions for details
     
     Returns:
-        tuple: A length-two tuple containing (the html filename) and (the loaded database).        
+        tuple: A length-two tuple containing (the html filename) and (the loaded database).
     """    
     
     if options is None:
@@ -315,6 +315,7 @@ def visualize_db(db_path, output_dir, image_base_dir, options=None):
     
         bboxes = []
         box_classes = []
+        box_score_strings = []
         
         # All the class labels we've seen for this image (with out without bboxes)
         image_categories = set()
@@ -363,12 +364,19 @@ def visualize_db(db_path, output_dir, image_base_dir, options=None):
             image_categories.add(category_name)
             
             if 'bbox' in anno:
-                bbox = anno['bbox']        
+                bbox = anno['bbox']                
                 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
         
@@ -382,8 +390,12 @@ def visualize_db(db_path, output_dir, image_base_dir, options=None):
         for c in illegal_characters:
             file_name = file_name.replace(c,'~')
         
-        rendering_info.append({'bboxes':bboxes, 'box_classes':box_classes, 'img_path':img_path,
-                               'output_file_name':file_name})
+        rendering_info_this_image = {'bboxes':bboxes,
+                                     'box_classes':box_classes,
+                                     'tags':box_score_strings,
+                                     'img_path':img_path,
+                                     'output_file_name':file_name}
+        rendering_info.append(rendering_info_this_image)
                 
         label_level_string = ''
         if len(annotation_level_for_image) > 0:
@@ -442,6 +454,9 @@ def visualize_db(db_path, output_dir, image_base_dir, options=None):
         img_path = rendering_info['img_path']
         bboxes = rendering_info['bboxes']
         bbox_classes = rendering_info['box_classes']
+        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)
         
@@ -467,11 +482,14 @@ def visualize_db(db_path, output_dir, image_base_dir, options=None):
             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,
+        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)
+                                           expansion=options.box_expansion,
+                                           tags=bbox_tags)
         
         image.save(output_full_path)
                 

megadetector/visualization/visualize_detector_output.py

@@ -22,6 +22,7 @@ from tqdm import tqdm
 
 from megadetector.data_management.annotations.annotation_constants import detector_bbox_category_id_to_name
 from megadetector.visualization import visualization_utils as vis_utils
+from megadetector.visualization.visualization_utils import blur_detections
 from megadetector.utils.ct_utils import get_max_conf
 from megadetector.utils import write_html_image_list
 from megadetector.detection.run_detector import get_typical_confidence_threshold_from_results
@@ -38,10 +39,17 @@ DEFAULT_DETECTOR_LABEL_MAP = {
 #%% Support functions
 
 def _render_image(entry,
-                 detector_label_map,classification_label_map,
-                 confidence_threshold,classification_confidence_threshold,
-                 render_detections_only,preserve_path_structure,out_dir,images_dir,
-                 output_image_width):
+                  detector_label_map,
+                  classification_label_map,
+                  confidence_threshold,
+                  classification_confidence_threshold,
+                  render_detections_only,
+                  preserve_path_structure,
+                  out_dir,
+                  images_dir,
+                  output_image_width,
+                  box_sort_order=None,
+                  category_names_to_blur=None):
     """
     Internal function for rendering a single image.
     """
@@ -65,22 +73,43 @@ def _render_image(entry,
         rendering_result['skipped_image'] = True
         return rendering_result
     
-    image_obj = os.path.join(images_dir, image_id)
-    if not os.path.exists(image_obj):
+    image_filename_in_abs = os.path.join(images_dir, image_id)
+    if not os.path.exists(image_filename_in_abs):
         print(f'Image {image_id} not found in images_dir')
         rendering_result['missing_image'] = True
         return rendering_result
 
+    # Load the image
+    image = vis_utils.open_image(image_filename_in_abs)
+    
+    # Find categories we're supposed to blur
+    category_ids_to_blur = []
+    if category_names_to_blur is not None:
+        if isinstance(category_names_to_blur,str):
+            category_names_to_blur = [category_names_to_blur]
+        for category_id in detector_label_map:
+            if detector_label_map[category_id] in category_names_to_blur:
+                category_ids_to_blur.append(category_id)
+    
+    detections_to_blur = []
+    for d in entry['detections']:
+        if d['conf'] >= confidence_threshold and d['category'] in category_ids_to_blur:
+            detections_to_blur.append(d)
+    if len(detections_to_blur) > 0:
+        blur_detections(image,detections_to_blur)
+        
+    # Resize if necessary
+    #
     # If output_image_width is -1 or None, this will just return the original image
-    image = vis_utils.resize_image(
-        vis_utils.open_image(image_obj), output_image_width)
+    image = vis_utils.resize_image(image, output_image_width)
 
     vis_utils.render_detection_bounding_boxes(
         entry['detections'], image, 
         label_map=detector_label_map,
         classification_label_map=classification_label_map,
         confidence_threshold=confidence_threshold,
-        classification_confidence_threshold=classification_confidence_threshold)
+        classification_confidence_threshold=classification_confidence_threshold,
+        box_sort_order=box_sort_order)
 
     if not preserve_path_structure:
         for char in ['/', '\\', ':']:
@@ -113,7 +142,9 @@ def visualize_detector_output(detector_output_path,
                               preserve_path_structure=False,
                               parallelize_rendering=False,
                               parallelize_rendering_n_cores=10,
-                              parallelize_rendering_with_threads=True):
+                              parallelize_rendering_with_threads=True,
+                              box_sort_order=None,
+                              category_names_to_blur=None):
     
     """
     Draws bounding boxes on images given the output of a detector.
@@ -145,6 +176,10 @@ def visualize_detector_output(detector_output_path,
         parallelize_rendering_with_threads (bool, optional): determines whether we use
             threads (True) or processes (False) for parallelization (ignored if parallelize_rendering 
             is False)
+        box_sort_order (str, optional): sorting scheme for detection boxes, can be None, "confidence", or
+            "reverse_confidence"
+        category_names_to_blur (list of str, optional): category names for which we should blur detections,
+            most commonly ['person']
 
     Returns:
         list: list of paths to annotated images
@@ -235,7 +270,9 @@ def visualize_detector_output(detector_output_path,
                                          preserve_path_structure=preserve_path_structure,
                                          out_dir=out_dir,
                                          images_dir=images_dir,
-                                         output_image_width=output_image_width),
+                                         output_image_width=output_image_width,
+                                         box_sort_order=box_sort_order,
+                                         category_names_to_blur=category_names_to_blur),
                                  images), total=len(images)))
         
     else:
@@ -245,7 +282,8 @@ def visualize_detector_output(detector_output_path,
             rendering_result = _render_image(entry,detector_label_map,classification_label_map,
                                             confidence_threshold,classification_confidence_threshold,
                                             render_detections_only,preserve_path_structure,out_dir,
-                                            images_dir,output_image_width)
+                                            images_dir,output_image_width,box_sort_order,
+                                            category_names_to_blur=category_names_to_blur)
             rendering_results.append(rendering_result)
         
     # ...for each image
@@ -334,12 +372,20 @@ def main():
     parser.add_argument(
         '-pps', '--preserve_path_structure', action='store_true',
         help='Preserve relative image paths (otherwise flattens and assigns unique file names)')
+    parser.add_argument(
+        '--category_names_to_blur', default=None, type=str,
+        help='Comma-separated list of category names to blur (or a single category name, typically "person")')
 
     if len(sys.argv[1:]) == 0:
         parser.print_help()
         parser.exit()
 
     args = parser.parse_args()
+    
+    category_names_to_blur = args.category_names_to_blur
+    if category_names_to_blur is not None:
+        category_names_to_blur = category_names_to_blur.split(',')
+    
     visualize_detector_output(
         detector_output_path=args.detector_output_path,
         out_dir=args.out_dir,
@@ -350,7 +396,8 @@ def main():
         random_seed=args.random_seed,
         render_detections_only=args.detections_only,
         preserve_path_structure=args.preserve_path_structure,
-        html_output_file=args.html_output_file)
+        html_output_file=args.html_output_file,
+        category_names_to_blur=category_names_to_blur)
 
     if args.html_output_file is not None and args.open_html_output_file:
         from megadetector.utils.path_utils import open_file