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

megadetector/data_management/databases/subset_json_db.py

@@ -2,14 +2,14 @@
 
 subset_json_db.py
 
-Select a subset of images (and associated annotations) from a .json file in COCO 
+Select a subset of images (and associated annotations) from a .json file in COCO
 Camera Traps format based on a string query.
 
 To subset .json files in the MegaDetector output format, see
 subset_json_detector_output.py.
 
 """
-    
+
 #%% Constants and imports
 
 import os
@@ -18,6 +18,7 @@ import json
 import argparse
 
 from tqdm import tqdm
+from megadetector.utils import ct_utils
 from copy import copy
 
 
@@ -25,22 +26,22 @@ from copy import copy
 
 def subset_json_db(input_json, query, output_json=None, ignore_case=False, verbose=False):
     """
-    Given a json file (or dictionary already loaded from a json file), produce a new 
-    database containing only the images whose filenames contain the string 'query', 
+    Given a json file (or dictionary already loaded from a json file), produce a new
+    database containing only the images whose filenames contain the string 'query',
     optionally writing that DB output to a new json file.
-    
+
     Args:
         input_json (str): COCO Camera Traps .json file to load, or an already-loaded dict
-        query (str or list): string to query for, only include images in the output whose filenames 
+        query (str or list): string to query for, only include images in the output whose filenames
             contain this string.  If this is a list, test for exact matches.
         output_json (str, optional): file to write the resulting .json file to
         ignore_case (bool, optional): whether to perform a case-insensitive search for [query]
         verbose (bool, optional): enable additional debug output
-        
+
     Returns:
         dict: CCT dictionary containing a subset of the images and annotations in the input dict
     """
-        
+
     # Load the input file if necessary
     if isinstance(input_json,str):
         print('Loading input .json...')
@@ -51,26 +52,26 @@ def subset_json_db(input_json, query, output_json=None, ignore_case=False, verbo
 
     # Find images matching the query
     images = []
-    
+
     if isinstance(query,str):
-        
+
         if ignore_case:
             query = query.lower()
-        
+
         for im in tqdm(input_data['images']):
             fn = im['file_name']
             if ignore_case:
                 fn = fn.lower()
             if query in fn:
                 images.append(im)
-                
+
     else:
-        
+
         query = set(query)
-        
+
         if ignore_case:
             query = set([s.lower() for s in query])
-            
+
         for im in input_data['images']:
             fn = im['file_name']
             if ignore_case:
@@ -79,27 +80,26 @@ def subset_json_db(input_json, query, output_json=None, ignore_case=False, verbo
                 images.append(im)
 
     image_ids = set([im['id'] for im in images])
-    
+
     # Find annotations referring to those images
     annotations = []
-    
+
     for ann in input_data['annotations']:
         if ann['image_id'] in image_ids:
             annotations.append(ann)
-    
+
     output_data = copy(input_data)
     output_data['images'] = images
     output_data['annotations'] = annotations
-    
+
     # Write the output file if requested
     if output_json is not None:
         if verbose:
             print('Writing output .json to {}'.format(output_json))
         output_dir = os.path.dirname(output_json)
         os.makedirs(output_dir,exist_ok=True)
-        with open(output_json,'w') as f:
-            json.dump(output_data,f,indent=1)
-    
+        ct_utils.write_json(output_json, output_data)
+
     if verbose:
         print('Keeping {} of {} images, {} of {} annotations'.format(
             len(output_data['images']),len(input_data['images']),
@@ -111,33 +111,33 @@ def subset_json_db(input_json, query, output_json=None, ignore_case=False, verbo
 #%% Interactive driver
 
 if False:
-    
+
     #%%
-    
+
     input_json = r"e:\Statewide_wolf_container\idfg_20190409.json"
     output_json = r"e:\Statewide_wolf_container\idfg_20190409_clearcreek.json"
     query = 'clearcreek'
     ignore_case = True
     db = subset_json_db(input_json, query, output_json, ignore_case)
-    
+
 
 #%% Command-line driver
 
-def main():
-    
+def main(): # noqa
+
     parser = argparse.ArgumentParser()
     parser.add_argument('input_json', type=str, help='Input file (a COCO Camera Traps .json file)')
-    parser.add_argument('output_json', type=str, help='Output file')    
-    parser.add_argument('query', type=str, help='Filename query')    
+    parser.add_argument('output_json', type=str, help='Output file')
+    parser.add_argument('query', type=str, help='Filename query')
     parser.add_argument('--ignore_case', action='store_true')
-    
+
     if len(sys.argv[1:]) == 0:
         parser.print_help()
         parser.exit()
 
-    args = parser.parse_args()    
-    
+    args = parser.parse_args()
+
     subset_json_db(args.input_json,args.query,args.output_json,args.ignore_case)
 
-if __name__ == '__main__':    
+if __name__ == '__main__':
     main()

megadetector/data_management/generate_crops_from_cct.py

@@ -10,6 +10,7 @@ each bounding box.
 #%% Imports and constants
 
 import os
+import argparse
 import json
 
 from tqdm import tqdm
@@ -22,7 +23,7 @@ def generate_crops_from_cct(cct_file,image_dir,output_dir,padding=0,flat_output=
     """
     Given a .json file in COCO Camera Traps format, creates a cropped image for
     each bounding box.
-    
+
     Args:
         cct_file (str): the COCO .json file from which we should load data
         image_dir (str): the folder where the images live; filenames in the .json
@@ -31,119 +32,167 @@ def generate_crops_from_cct(cct_file,image_dir,output_dir,padding=0,flat_output=
         padding (float, optional): number of pixels we should expand each box before
             cropping
         flat_output (bool, optional): if False, folder structure will be preserved
-            in the output, e.g. the image a/b/c/d.jpg will result in image files 
+            in the output, e.g. the image a/b/c/d.jpg will result in image files
             in the output folder called, e.g., a/b/c/d_crop_000_id_12345.jpg.  If
-            [flat_output] is True, the corresponding output image will be 
-            a_b_c_d_crop_000_id_12345.jpg.            
+            [flat_output] is True, the corresponding output image will be
+            a_b_c_d_crop_000_id_12345.jpg.
     """
-    
+
     ## Read and validate input
-    
+
     assert os.path.isfile(cct_file)
     assert os.path.isdir(image_dir)
     os.makedirs(output_dir,exist_ok=True)
 
     with open(cct_file,'r') as f:
         d = json.load(f)
-   
-    
+
+
     ## Find annotations for each image
-    
+
     from collections import defaultdict
-    
+
     # This actually maps image IDs to annotations, but only to annotations
     # containing boxes
     image_id_to_boxes = defaultdict(list)
-    
+
     n_boxes = 0
-    
+
     for ann in d['annotations']:
         if 'bbox' in ann:
             image_id_to_boxes[ann['image_id']].append(ann)
             n_boxes += 1
-            
+
     print('Found {} boxes in {} annotations for {} images'.format(
         n_boxes,len(d['annotations']),len(d['images'])))
-    
-    
+
+
     ## Generate crops
-        
+
     # im = d['images'][0]
     for im in tqdm(d['images']):
-        
+
         input_image_fn = os.path.join(os.path.join(image_dir,im['file_name']))
         assert os.path.isfile(input_image_fn), 'Could not find image {}'.format(input_image_fn)
-        
+
         if im['id'] not in image_id_to_boxes:
             continue
-        
+
         annotations_this_image = image_id_to_boxes[im['id']]
-        
+
         # Load the image
         img = Image.open(input_image_fn)
-        
+
         # Generate crops
         # i_ann = 0; ann = annotations_this_image[i_ann]
         for i_ann,ann in enumerate(annotations_this_image):
-            
+
             # x/y/w/h, origin at the upper-left
             bbox = ann['bbox']
-            
+
             xmin = bbox[0]
             ymin = bbox[1]
             xmax = xmin + bbox[2]
             ymax = ymin + bbox[3]
-            
+
             xmin -= padding / 2
             ymin -= padding / 2
             xmax += padding / 2
             ymax += padding / 2
-            
+
             xmin = max(xmin,0)
             ymin = max(ymin,0)
             xmax = min(xmax,img.width-1)
             ymax = min(ymax,img.height-1)
-                        
+
             crop = img.crop(box=[xmin, ymin, xmax, ymax])
-            
+
             output_fn = os.path.splitext(im['file_name'])[0].replace('\\','/')
             if flat_output:
                 output_fn = output_fn.replace('/','_')
             output_fn = output_fn + '_crop' + str(i_ann).zfill(3) + '_id_' + ann['id']
             output_fn = output_fn + '.jpg'
-                
+
             output_full_path = os.path.join(output_dir,output_fn)
-            
+
             if not flat_output:
                 os.makedirs(os.path.dirname(output_full_path),exist_ok=True)
-                
+
             crop.save(output_full_path)
-            
+
         # ...for each box
-        
+
     # ...for each image
-    
+
 # ...generate_crops_from_cct()
 
 
 #%% Interactive driver
 
 if False:
-    
+
     pass
 
     #%%
-    
+
     cct_file = os.path.expanduser('~/data/noaa/noaa_estuary_fish.json')
     image_dir = os.path.expanduser('~/data/noaa/JPEGImages')
     padding = 50
     flat_output = True
     output_dir = '/home/user/tmp/noaa-fish-crops'
-    
+
     generate_crops_from_cct(cct_file,image_dir,output_dir,padding,flat_output=True)
     files = os.listdir(output_dir)
-    
-    
+
+
 #%% Command-line driver
 
-# TODO
+def main():
+    """
+    Command-line interface to generate crops from a COCO Camera Traps .json file.
+    """
+
+    parser = argparse.ArgumentParser(
+        description='Generate cropped images from a COCO Camera Traps .json file'
+    )
+    parser.add_argument(
+        'cct_file',
+        type=str,
+        help='COCO .json file to load data from'
+    )
+    parser.add_argument(
+        'image_dir',
+        type=str,
+        help='Folder where images are located'
+    )
+    parser.add_argument(
+        'output_dir',
+        type=str,
+        help='Folder to which we should write cropped images'
+    )
+    parser.add_argument(
+        '--padding',
+        type=int,
+        default=0,
+        help='Pixels to expand each box before cropping'
+    )
+    parser.add_argument(
+        '--flat_output',
+        action='store_true',
+        help='Flatten folder structure in output (preserves folder structure by default)'
+    )
+
+    args = parser.parse_args()
+
+    generate_crops_from_cct(
+        cct_file=args.cct_file,
+        image_dir=args.image_dir,
+        output_dir=args.output_dir,
+        padding=args.padding,
+        flat_output=args.flat_output
+    )
+
+    print(f'Generated crops in {args.output_dir}')
+
+if __name__ == '__main__':
+    main()

megadetector/data_management/get_image_sizes.py

@@ -2,7 +2,7 @@
 
 get_image_sizes.py
 
-Given a json-formatted list of image filenames, retrieves the width and height of 
+Given a json-formatted list of image filenames, retrieves the width and height of
 every image, optionally writing the results to a new .json file.
 
 """
@@ -35,45 +35,45 @@ def _get_image_size(image_path,image_prefix=None):
     Support function to get the size of a single image.  Returns a (path,w,h) tuple.
     w and h will be -1 if the image fails to load.
     """
-    
+
     if image_prefix is not None:
         full_path = os.path.join(image_prefix,image_path)
     else:
         full_path = image_path
-    
+
     # Is this image on disk?
     if not os.path.isfile(full_path):
         print('Could not find image {}'.format(full_path))
         return (image_path,-1,-1)
 
-    try:        
+    try:
         pil_im = Image.open(full_path)
-        w = pil_im.width            
+        w = pil_im.width
         h = pil_im.height
         return (image_path,w,h)
-    except Exception as e:    
+    except Exception as e:
         print('Error reading image {}: {}'.format(full_path,str(e)))
         return (image_path,-1,-1)
-    
-    
+
+
 def get_image_sizes(filenames,image_prefix=None,output_file=None,
                     n_workers=default_n_threads,use_threads=True,
                     recursive=True):
     """
     Gets the width and height of all images in [filenames], which can be:
-        
+
     * A .json-formatted file containing list of strings
     * A folder
     * A list of files
 
     ...returning a list of (path,w,h) tuples, and optionally writing the results to [output_file].
-    
+
     Args:
-        filenames (str or list): the image filenames for which we should retrieve sizes, 
-            can be the name of a .json-formatted file containing list of strings, a folder 
+        filenames (str or list): the image filenames for which we should retrieve sizes,
+            can be the name of a .json-formatted file containing list of strings, a folder
             in which we should enumerate images, or a list of files.
         image_prefix (str, optional): optional prefix to add to images to get to full paths;
-            useful when [filenames] contains relative files, in which case [image_prefix] is the 
+            useful when [filenames] contains relative files, in which case [image_prefix] is the
             base folder for the source images.
         output_file (str, optional): a .json file to write the imgae sizes
         n_workers (int, optional): number of parallel workers to use, set to <=1 to
@@ -82,57 +82,62 @@ def get_image_sizes(filenames,image_prefix=None,output_file=None,
             for parallelization; not relevant if [n_workers] <= 1
         recursive (bool, optional): only relevant if [filenames] is actually a folder,
             determines whether image enumeration within that folder will be recursive
-        
+
     Returns:
         list: list of (path,w,h) tuples
-    """        
-    
+    """
+
     if output_file is not None:
         assert os.path.isdir(os.path.dirname(output_file)), \
             'Illegal output file {}, parent folder does not exist'.format(output_file)
-        
+
     if isinstance(filenames,str) and os.path.isfile(filenames):
-        with open(filenames,'r') as f:        
+        with open(filenames,'r') as f:
             filenames = json.load(f)
         filenames = [s.strip() for s in filenames]
     elif isinstance(filenames,str) and os.path.isdir(filenames):
         filenames = find_images(filenames,recursive=recursive,
                                 return_relative_paths=False,convert_slashes=True)
     else:
-        assert isinstance(filenames,list)        
-    
+        assert isinstance(filenames,list)
+
     if n_workers <= 1:
-        
+
         all_results = []
         for i_file,fn in tqdm(enumerate(filenames),total=len(filenames)):
             all_results.append(_get_image_size(fn,image_prefix=image_prefix))
-    
+
     else:
-        
+
         print('Creating a pool with {} workers'.format(n_workers))
         if use_threads:
-            pool = ThreadPool(n_workers)        
+            pool = ThreadPool(n_workers)
         else:
             pool = Pool(n_workers)
         # all_results = list(tqdm(pool.imap(process_image, filenames), total=len(filenames)))
-        all_results = list(tqdm(pool.imap(
-            partial(_get_image_size,image_prefix=image_prefix), filenames), total=len(filenames)))
-    
+        try:
+            all_results = list(tqdm(pool.imap(
+                partial(_get_image_size,image_prefix=image_prefix), filenames), total=len(filenames)))
+        finally:
+            pool.close()
+            pool.join()
+            print("Pool closed and joined for image size reads")
+
     if output_file is not None:
         with open(output_file,'w') as f:
             json.dump(all_results,f,indent=1)
-            
+
     return all_results
 
-    
+
 #%% Interactive driver
 
 if False:
 
-    pass    
+    pass
 
     #%%
-    
+
     # List images in a test folder
     base_dir = r'c:\temp\test_images'
     image_list_file = os.path.join(base_dir,'images.json')
@@ -140,50 +145,50 @@ if False:
     image_size_file = os.path.join(base_dir,'image_sizes.json')
     from megadetector.utils import path_utils
     image_names = path_utils.find_images(base_dir,recursive=True)
-    
+
     with open(image_list_file,'w') as f:
         json.dump(image_names,f,indent=1)
-        
+
     relative_image_names = []
     for s in image_names:
         relative_image_names.append(os.path.relpath(s,base_dir))
-    
+
     with open(relative_image_list_file,'w') as f:
         json.dump(relative_image_names,f,indent=1)
-    
-    
+
+
     #%%
-    
+
     get_image_sizes(relative_image_list_file,image_size_file,image_prefix=base_dir,n_threads=4)
-    
-    
+
+
 #%% Command-line driver
-    
-def main():
-    
+
+def main(): # noqa
+
     parser = argparse.ArgumentParser()
     parser.add_argument('filenames',type=str,
                         help='Folder from which we should fetch image sizes, or .json file with a list of filenames')
     parser.add_argument('output_file',type=str,
                         help='Output file (.json) to which we should write image size information')
     parser.add_argument('--image_prefix', type=str, default=None,
-                        help='Prefix to append to image filenames, only relevant if [filenames] points to a list of ' + \
-                             'relative paths')        
+                        help='Prefix to append to image filenames, only relevant if [filenames] points to a ' + \
+                             'list of relative paths')
     parser.add_argument('--n_threads', type=int, default=default_n_threads,
                         help='Number of concurrent workers, set to <=1 to disable parallelization (default {})'.format(
                             default_n_threads))
-                        
+
     if len(sys.argv[1:])==0:
         parser.print_help()
         parser.exit()
-        
+
     args = parser.parse_args()
-    
+
     _ = get_image_sizes(filenames=args.filenames,
                         output_file=args.output_file,
                         image_prefix=args.image_prefix,
                         n_workers=args.n_threads)
-        
+
 if __name__ == '__main__':
-    
+
     main()