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

megadetector/data_management/remap_coco_categories.py

@@ -2,7 +2,9 @@
 
 remap_coco_categories.py
 
-Given a COCO-formatted dataset, remap the categories to a new mapping.
+Given a COCO-formatted dataset, remap the categories to a new mapping.  A common use
+case is to take a fine-grained dataset (e.g. with species-level categories) and
+map them to coarse categories (typically MD categories).
 
 """
 
@@ -10,6 +12,7 @@ Given a COCO-formatted dataset, remap the categories to a new mapping.
 
 import os
 import json
+import argparse
 
 from copy import deepcopy
 from megadetector.utils.ct_utils import invert_dictionary
@@ -25,24 +28,24 @@ def remap_coco_categories(input_data,
     """
     Given a COCO-formatted dataset, remap the categories to a new categories mapping, optionally
     writing the results to a new file.
-    
+
     Args:
-        input_data (str or dict): a COCO-formatted dict or a filename.  If it's a dict, it will 
+        input_data (str or dict): a COCO-formatted dict or a filename.  If it's a dict, it will
             be copied, not modified in place.
-        output_category_name_to_id (dict) a dict mapping strings to ints.  Categories not in 
+        output_category_name_to_id (dict): a dict mapping strings to ints.  Categories not in
             this dict will be ignored or will result in errors, depending on allow_unused_categories.
-        input_category_name_to_output_category_name: a dict mapping strings to strings.  
-            Annotations using categories not in this dict will be omitted or will result in 
+        input_category_name_to_output_category_name (dict): a dict mapping strings to strings.
+            Annotations using categories not in this dict will be omitted or will result in
             errors, depending on allow_unused_categories.
         output_file (str, optional): output file to which we should write remapped COCO data
         allow_unused_categories (bool, optional): should we ignore categories not present in the
             input/output mappings?  If this is False and we encounter an unmapped category, we'll
             error.
-    
+
     Returns:
         dict: COCO-formatted dict
     """
-    
+
     if isinstance(input_data,str):
         assert os.path.isfile(input_data), "Can't find file {}".format(input_data)
         with open(input_data,'r') as f:
@@ -51,45 +54,45 @@ def remap_coco_categories(input_data,
     else:
         assert isinstance(input_data,dict), 'Illegal COCO input data'
         input_data = deepcopy(input_data)
-    
+
     # It's safe to modify in-place now
     output_data = input_data
-    
+
     # Read input name --> ID mapping
     input_category_name_to_input_category_id = {}
     for c in input_data['categories']:
         input_category_name_to_input_category_id[c['name']] = c['id']
     input_category_id_to_input_category_name = \
         invert_dictionary(input_category_name_to_input_category_id)
-        
+
     # Map input IDs --> output IDs
     input_category_id_to_output_category_id = {}
     input_category_names = list(input_category_name_to_output_category_name.keys())
-    
+
     # input_name = input_category_names[0]
     for input_name in input_category_names:
-        
+
         output_name = input_category_name_to_output_category_name[input_name]
         assert output_name in output_category_name_to_id, \
             'No output ID for {} --> {}'.format(input_name,output_name)
         input_id = input_category_name_to_input_category_id[input_name]
         output_id = output_category_name_to_id[output_name]
         input_category_id_to_output_category_id[input_id] = output_id
-    
+
     # ...for each category we want to keep
-        
+
     printed_unused_category_warnings = set()
-    
+
     valid_annotations = []
-    
+
     # Map annotations
     for ann in output_data['annotations']:
-        
+
         input_category_id = ann['category_id']
         if input_category_id not in input_category_id_to_output_category_id:
             if allow_unused_categories:
                 if input_category_id not in printed_unused_category_warnings:
-                    printed_unused_category_warnings.add(input_category_id)            
+                    printed_unused_category_warnings.add(input_category_id)
                     input_category_name = \
                         input_category_id_to_input_category_name[input_category_id]
                     s = 'Skipping unmapped category ID {} ({})'.format(
@@ -98,31 +101,31 @@ def remap_coco_categories(input_data,
                 continue
             else:
                 s = 'Unmapped category ID {}'.format(input_category_id)
-                raise ValueError(s)        
+                raise ValueError(s)
         output_category_id = input_category_id_to_output_category_id[input_category_id]
         ann['category_id'] = output_category_id
-        valid_annotations.append(ann)        
-        
+        valid_annotations.append(ann)
+
     # ...for each annotation
-    
+
     # The only reason annotations should get excluded is the case where we allow
     # unused categories
     if not allow_unused_categories:
         assert len(valid_annotations) == len(output_data['annotations'])
-    
+
     output_data['annotations'] = valid_annotations
-    
+
     # Update the category list
     output_categories = []
     for output_name in output_category_name_to_id:
         category = {'name':output_name,'id':output_category_name_to_id[output_name]}
         output_categories.append(category)
     output_data['categories'] = output_categories
-        
+
     if output_file is not None:
         with open(output_file,'w') as f:
             json.dump(output_data,f,indent=1)
-            
+
     return input_data
 
 # ...def remap_coco_categories(...)
@@ -130,4 +133,63 @@ def remap_coco_categories(input_data,
 
 #%% Command-line driver
 
-# TODO
+def main():
+    """
+    Command-line interface to remap COCO categories.
+    """
+
+    parser = argparse.ArgumentParser(
+        description='Remap categories in a COCO-formatted dataset'
+    )
+    parser.add_argument(
+        'input_coco_file',
+        type=str,
+        help='Path to the input COCO .json file'
+    )
+    parser.add_argument(
+        'output_category_map_file',
+        type=str,
+        help="Path to a .json file mapping output category names to integer IDs (e.g., {'cat':0, 'dog':1})"
+    )
+    parser.add_argument(
+        'input_to_output_category_map_file',
+        type=str,
+        help="Path to a .json file mapping input category names to output category names" + \
+            " (e.g., {'old_cat_name':'cat', 'old_dog_name':'dog'})"
+    )
+    parser.add_argument(
+        'output_coco_file',
+        type=str,
+        help='Path to save the remapped COCO .json file'
+    )
+    parser.add_argument(
+        '--allow_unused_categories',
+        action='store_true',
+        help='Allow unmapped categories (by default, errors on unmapped categories)'
+    )
+
+    args = parser.parse_args()
+
+    # Load category mappings
+    with open(args.output_category_map_file, 'r') as f:
+        output_category_name_to_id = json.load(f)
+
+    with open(args.input_to_output_category_map_file, 'r') as f:
+        input_category_name_to_output_category_name = json.load(f)
+
+    # Load COCO data
+    with open(args.input_coco_file, 'r') as f:
+        input_data = json.load(f)
+
+    remap_coco_categories(
+        input_data=input_data,
+        output_category_name_to_id=output_category_name_to_id,
+        input_category_name_to_output_category_name=input_category_name_to_output_category_name,
+        output_file=args.output_coco_file,
+        allow_unused_categories=args.allow_unused_categories
+    )
+
+    print(f'Successfully remapped categories and saved to {args.output_coco_file}')
+
+if __name__ == '__main__':
+    main()

megadetector/data_management/remove_exif.py

@@ -2,7 +2,7 @@
 
 remove_exif.py
 
-Removes all EXIF/IPTC/XMP metadata from a folder of images, without making 
+Removes all EXIF/IPTC/XMP metadata from a folder of images, without making
 backup copies, using pyexiv2.  Ignores non-jpeg images.
 
 This module is rarely used, and pyexiv2 is not thread-safe, so pyexiv2 is not
@@ -14,6 +14,7 @@ included in package-level dependency lists.  YMMV.
 
 import os
 import glob
+import argparse
 
 from multiprocessing.pool import Pool as Pool
 from tqdm import tqdm
@@ -21,22 +22,33 @@ from tqdm import tqdm
 
 #%% Support functions
 
-# Pyexif2 is not thread safe, do not call this function in parallel within a process
-#
-# Parallelizing across processes is fine.
 def remove_exif_from_image(fn):
+    """
+    Remove EXIF information from a single image
+
+    pyexiv2 is not thread safe, do not call this function in parallel within a process.
+
+    Parallelizing across processes is fine.
+
+    Args:
+        fn (str): image file from which we should remove EXIF information
+
+    Returns:
+        bool: whether EXIF removal succeeded
+    """
+
+    import pyexiv2 # type: ignore
 
-    import pyexiv2
-    
     try:
         img = pyexiv2.Image(fn)
         img.clear_exif()
         img.clear_iptc()
         img.clear_xmp()
-        img.close()        
+        img.close()
     except Exception as e:
         print('EXIF error on {}: {}'.format(fn,str(e)))
-    
+        return False
+
     return True
 
 
@@ -44,22 +56,23 @@ def remove_exif_from_image(fn):
 
 def remove_exif(image_base_folder,recursive=True,n_processes=1):
     """
-    Removes all EXIF/IPTC/XMP metadata from a folder of images, without making 
+    Removes all EXIF/IPTC/XMP metadata from a folder of images, without making
     backup copies, using pyexiv2.  Ignores non-jpeg images.
-    
+
     Args:
         image_base_folder (str): the folder from which we should remove EXIF data
         recursive (bool, optional): whether to process [image_base_folder] recursively
         n_processes (int, optional): number of concurrent workers.  Because pyexiv2 is not
-            thread-safe, only process-based parallelism is supported.        
+            thread-safe, only process-based parallelism is supported.
     """
+
     try:
-        import pyexiv2 #noqa
+        import pyexiv2 # type: ignore #noqa
     except:
         print('pyexiv2 not available; try "pip install pyexiv2"')
         raise
 
-        
+
     ##%% List files
 
     assert os.path.isdir(image_base_folder), \
@@ -67,25 +80,72 @@ def remove_exif(image_base_folder,recursive=True,n_processes=1):
     all_files = [f for f in glob.glob(image_base_folder+ "*/**", recursive=recursive)]
     image_files = [s for s in all_files if \
                    (s.lower().endswith('.jpg') or s.lower().endswith('.jpeg'))]
-        
+
 
     ##%% Remove EXIF data (execution)
 
     if n_processes == 1:
-        
+
         # fn = image_files[0]
         for fn in tqdm(image_files):
             remove_exif_from_image(fn)
-            
+
     else:
         # pyexiv2 is not thread-safe, so we need to use processes
-        print('Starting parallel process pool with {} workers'.format(n_processes))
-        pool = Pool(n_processes)
-        _ = list(tqdm(pool.imap(remove_exif_from_image,image_files),total=len(image_files)))
-            
+        pool = None
+        try:
+            print('Starting parallel process pool with {} workers'.format(n_processes))
+            pool = Pool(n_processes)
+            _ = list(tqdm(pool.imap(remove_exif_from_image,image_files),total=len(image_files)))
+        finally:
+            pool.close()
+            pool.join()
+            print("Pool closed and joined for EXIF removal")
+
 # ...remove_exif(...)
 
 
 #%% Command-line driver
 
-## TODO
+def main():
+    """
+    Command-line interface to remove EXIF data from images.
+    """
+
+    parser = argparse.ArgumentParser(
+        description='Removes EXIF/IPTC/XMP metadata from images in a folder'
+    )
+    parser.add_argument(
+        'image_base_folder',
+        type=str,
+        help='Folder to process for EXIF removal'
+    )
+    parser.add_argument(
+        '--nonrecursive',
+        action='store_true',
+        help="Don't recurse into [image_base_folder] (default is recursive)"
+    )
+    parser.add_argument(
+        '--n_processes',
+        type=int,
+        default=1,
+        help='Number of concurrent processes for EXIF removal (default: 1)'
+    )
+
+    args = parser.parse_args()
+
+    recursive = (not args.nonrecursive)
+
+    print('Processing folder: {}'.format(args.image_base_folder))
+    if not os.path.isdir(args.image_base_folder):
+        raise ValueError('Folder not found at {}'.format(args.image_base_folder))
+
+    remove_exif(
+        image_base_folder=args.image_base_folder,
+        recursive=recursive,
+        n_processes=args.n_processes
+    )
+    print('Finished removing EXIF data')
+
+if __name__ == '__main__':
+    main()