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

megadetector/data_management/databases/add_width_and_height_to_db.py

@@ -4,30 +4,104 @@ add_width_and_height_to_db.py
 
 Grabs width and height from actual image files for a .json database that is missing w/h.
 
-TODO: this is a one-off script waiting to be cleaned up for more general use.
-
 """
 
 #%% Imports and constants
 
+import os
+import sys
 import json
+import argparse
+
+from tqdm import tqdm
 from PIL import Image
 
-datafile = '/datadrive/snapshotserengeti/databases/snapshotserengeti.json'
-image_base = '/datadrive/snapshotserengeti/images/'
+from megadetector.utils import ct_utils
+
+
+#%% Main resizing function
+
+def add_width_and_height_to_db(input_file,output_file,image_base_folder):
+    """
+    Add width and height to images in the COCO db [input_file]
+    that don't have non-None w/h values.  Does not verify correctness
+    for images that already have non-None w/h values.  Ignores files that
+    fail to open.
+
+    Args:
+        input_file (str): the COCO .json file to process
+        output_file (str): the COCO .json file to write
+        image_base_folder (str): image filenames in [input_file] should be relative
+            to this folder
+
+    Returns:
+        list: the list of image dicts that were modified
+    """
+
+    with open(input_file,'r') as f:
+        d = json.load(f)
+
+    to_return = []
+
+    for im in tqdm(d['images']):
+
+        if ('height' not in im) or ('width' not in im) or \
+           (im['height'] is None) or (im['width'] is None) or \
+           (im['height'] <= 0) or (im['width'] <= 0):
+
+            fn_relative = im['file_name']
+            fn_abs = os.path.join(image_base_folder,fn_relative)
 
-def main():
+            if not os.path.isfile(fn_abs):
+                print('Could not find image file {}'.format(fn_abs))
+                continue
 
-    with open(datafile,'r') as f:
-        data = json.load(f)
+            try:
+                im_w, im_h = Image.open(fn_abs).size
+            except Exception as e:
+                print('Error opening file {}: {}'.format(fn_abs,str(e)))
+                continue
+
+            assert isinstance(im_w,int) and isinstance(im_h,int) and \
+                im_w > 0 and im_h > 0, \
+                'Illegal size retrieved for {}'.format(fn_abs)
 
-    for im in data['images']:
-        if 'height' not in im:
-            im_w, im_h = Image.open(image_base+im['file_name']).size
             im['height'] = im_h
             im['width'] = im_w
+            to_return.append(im)
+
+        # ...if we need to add width and/or height to this image
+
+    # ...for each image
+
+    ct_utils.write_json(output_file, d)
+
+    print('Added size information to {} of {} images'.format(
+        len(to_return), len(d['images'])))
+
+    return to_return
 
-    json.dump(data, open(datafile,'w'))
+# ...def add_width_and_height_to_db(...)
+
+
+#%% Command-line driver
 
 if __name__ == '__main__':
-    main()
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('input_file', type=str,
+                        help='Input COCO-formatted .json file')
+    parser.add_argument('output_file', type=str,
+                        help='Output COCO-formatted .json file')
+    parser.add_argument('image_base_folder', type=str,
+                        help='Base directory for images')
+
+    if len(sys.argv[1:]) == 0:
+        parser.print_help()
+        parser.exit()
+
+    args = parser.parse_args()
+
+    add_width_and_height_to_db(args.input_file,
+                               args.output_file,
+                               args.image_base_folder)

megadetector/data_management/databases/combine_coco_camera_traps_files.py

@@ -10,7 +10,7 @@ writing the results to another .json file.
 - Checks compatibility in info structs, within reason.
 
 *Example command-line invocation*
-    
+
 combine_coco_camera_traps_files input1.json input2.json ... inputN.json output.json
 
 """
@@ -20,12 +20,13 @@ combine_coco_camera_traps_files input1.json input2.json ... inputN.json output.j
 import argparse
 import json
 import sys
+from megadetector.utils import ct_utils
 
 
 #%% Merge functions
 
-def combine_cct_files(input_files, 
-                      output_file=None, 
+def combine_cct_files(input_files,
+                      output_file=None,
                       require_uniqueness=True,
                       filename_prefixes=None):
     """
@@ -38,12 +39,12 @@ def combine_cct_files(input_files,
         require_uniqueness (bool, optional): whether to require that the images in
             each input_dict be unique
         filename_prefixes (dict, optional): dict mapping input filenames to strings
-            that should be prepended to image filenames from that source            
-    
+            that should be prepended to image filenames from that source
+
     Returns:
         dict: the merged COCO-formatted .json dict
     """
-    
+
     input_dicts = []
     print('Loading input files')
     for fn in input_files:
@@ -60,34 +61,33 @@ def combine_cct_files(input_files,
 
     print('Writing output')
     if output_file is not None:
-        with open(output_file, 'w') as f:
-            json.dump(merged_dict, f, indent=1)
+        ct_utils.write_json(output_file, merged_dict)
 
     return merged_dict
 
 
-def combine_cct_dictionaries(input_dicts, require_uniqueness=True):                                    
+def combine_cct_dictionaries(input_dicts, require_uniqueness=True):
     """
     Merges the list of COCO Camera Traps dictionaries [input_dicts].  See module header
     comment for details on merge rules.
 
     Args:
-        input_dicts: list of CCT dicts
-        require_uniqueness: bool, whether to require that the images in
+        input_dicts (list of dict): list of CCT dicts
+        require_uniqueness (bool, optional): whether to require that the images in
             each input_dict be unique
 
-    Returns: 
+    Returns:
         dict: the merged COCO-formatted .json dict
     """
-    
+
     filename_to_image = {}
     all_annotations = []
     info = None
-    
+
     category_name_to_id = {}
     category_name_to_id['empty'] = 0
     next_category_id = 1
-    
+
     known_fields = ['info', 'categories', 'annotations','images','filename_prefix']
 
     # i_input_dict = 0; input_dict = input_dicts[i_input_dict]
@@ -96,19 +96,19 @@ def combine_cct_dictionaries(input_dicts, require_uniqueness=True):
         filename_prefix = ''
         if ('filename_prefix' in input_dict.keys()):
             filename_prefix = input_dict['filename_prefix']
-        
+
         for k in input_dict.keys():
             if k not in known_fields:
                 raise ValueError(f'Unrecognized CCT field: {k}')
 
         # We will prepend an index to every ID to guarantee uniqueness
         index_string = 'ds' + str(i_input_dict).zfill(3) + '_'
-        
+
         old_cat_id_to_new_cat_id = {}
-        
+
         # Map detection categories from the original data set into the merged data set
         for original_category in input_dict['categories']:
-        
+
             original_cat_id = original_category['id']
             cat_name = original_category['name']
             if cat_name in category_name_to_id:
@@ -117,58 +117,58 @@ def combine_cct_dictionaries(input_dicts, require_uniqueness=True):
                 new_cat_id = next_category_id
                 next_category_id += 1
                 category_name_to_id[cat_name] = new_cat_id
-                
+
             if original_cat_id in old_cat_id_to_new_cat_id:
                 assert old_cat_id_to_new_cat_id[original_cat_id] == new_cat_id
             else:
                 old_cat_id_to_new_cat_id[original_cat_id] = new_cat_id
-            
+
         # ...for each category
-        
-        
+
+
         # Merge original image list into the merged data set
         for im in input_dict['images']:
-            
+
             if 'seq_id' in im:
                 im['seq_id'] = index_string + im['seq_id']
             if 'location' in im:
                 im['location'] = index_string + im['location']
-                
+
             im_file = filename_prefix + im['file_name']
             im['file_name'] = im_file
             if require_uniqueness:
-                assert im_file not in filename_to_image, f'Duplicate image: {im_file}'                
+                assert im_file not in filename_to_image, f'Duplicate image: {im_file}'
             else:
                 if im_file in filename_to_image:
                     print('Redundant image {}'.format(im_file))
-            
+
             # Create a unique ID
             im['id'] = index_string + im['id']
             filename_to_image[im_file] = im
 
         # ...for each image
-                  
-        
+
+
         # Same for annotations
         for ann in input_dict['annotations']:
-            
+
             ann['image_id'] = index_string + ann['image_id']
             ann['id'] = index_string + ann['id']
             assert ann['category_id'] in old_cat_id_to_new_cat_id
             ann['category_id'] = old_cat_id_to_new_cat_id[ann['category_id']]
-            
+
         # ...for each annotation
-        
+
         all_annotations.extend(input_dict['annotations'])
-        
+
         # Merge info dicts, don't check completion time fields
         if info is None:
-            import copy            
+            import copy
             info = copy.deepcopy(input_dict['info'])
             info['original_info'] = [input_dict['info']]
         else:
             info['original_info'].append(input_dict['info'])
-            
+
     # ...for each dictionary
 
     # Convert merged image dictionaries to a sorted list
@@ -176,12 +176,12 @@ def combine_cct_dictionaries(input_dicts, require_uniqueness=True):
 
     all_categories = [{'id':category_name_to_id[cat_name],'name':cat_name} for\
                       cat_name in category_name_to_id.keys()]
-    
+
     merged_dict = {'info': info,
                    'categories': all_categories,
                    'images': sorted_images,
                    'annotations': all_annotations}
-    
+
     return merged_dict
 
 # ...combine_cct_dictionaries(...)
@@ -189,8 +189,8 @@ def combine_cct_dictionaries(input_dicts, require_uniqueness=True):
 
 #%% Command-line driver
 
-def main():
-    
+def main(): # noqa
+
     parser = argparse.ArgumentParser()
     parser.add_argument(
         'input_paths', nargs='+',
@@ -202,7 +202,7 @@ def main():
     if len(sys.argv[1:]) == 0:
         parser.print_help()
         parser.exit()
-        
+
     args = parser.parse_args()
     combine_cct_files(args.input_paths, args.output_path)