megadetector 10.0.13__py3-none-any.whl

megadetector/data_management/databases/subset_json_db.py

@@ -0,0 +1,195 @@
+"""
+
+subset_json_db.py
+
+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 sys
+import json
+import argparse
+
+from tqdm import tqdm
+from copy import copy
+
+from megadetector.utils import ct_utils
+from megadetector.utils.ct_utils import sort_list_of_dicts_by_key
+
+
+#%% Functions
+
+def subset_json_db(input_json,
+                   query,
+                   output_json=None,
+                   ignore_case=False,
+                   remap_categories=True,
+                   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',
+    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
+            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]
+        remap_categories (bool, optional): trim the category list to only the categores used
+            in the subset
+        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...')
+        with open(input_json, 'r') as f:
+            input_data = json.load(f)
+    else:
+        input_data = input_json
+
+    # 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:
+                fn = fn.lower()
+            if fn in query:
+                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
+
+    # Remap categories if necessary
+    if remap_categories:
+
+        category_ids_used = set()
+        for ann in annotations:
+            category_ids_used.add(ann['category_id'])
+
+        if verbose:
+            print('Keeping {} of {} categories'.format(
+                len(category_ids_used),len(input_data['categories'])))
+
+        input_category_id_to_output_category_id = {}
+
+        next_category_id = 0
+
+        # Build mappings from old to new category IDs
+        for input_category_id in category_ids_used:
+            assert isinstance(input_category_id,int), \
+                'Illegal category ID {}'.format(input_category_id)
+            output_category_id = next_category_id
+            next_category_id = next_category_id + 1
+            input_category_id_to_output_category_id[input_category_id] = output_category_id
+
+        # Modify the annotations
+        for ann in annotations:
+            assert ann['category_id'] in input_category_id_to_output_category_id
+            ann['category_id'] = input_category_id_to_output_category_id[ann['category_id']]
+
+        output_categories = []
+
+        # Re-write the category table
+        for cat in input_data['categories']:
+
+            if cat['id'] in input_category_id_to_output_category_id:
+
+                # There may be non-required fields, so don't just create an empty dict
+                # and copy the name/id field, keep the original dict other than "id"
+                output_category = copy(cat)
+                output_category['id'] = input_category_id_to_output_category_id[cat['id']]
+                output_categories.append(output_category)
+
+        output_categories = sort_list_of_dicts_by_key(output_categories,'id')
+        output_data['categories'] = output_categories
+
+    # ...if we need to remap categories
+
+    # Write the output file if requested
+    if output_json is not None:
+        if verbose:
+            print('Writing output .json to {}'.format(output_json))
+        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']),
+            len(output_data['annotations']),len(input_data['annotations'])))
+
+    return output_data
+
+
+#%% 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(): # 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('--ignore_case', action='store_true')
+
+    if len(sys.argv[1:]) == 0:
+        parser.print_help()
+        parser.exit()
+
+    args = parser.parse_args()
+
+    subset_json_db(args.input_json,args.query,args.output_json,args.ignore_case)
+
+if __name__ == '__main__':
+    main()

megadetector/data_management/generate_crops_from_cct.py

@@ -0,0 +1,200 @@
+"""
+
+generate_crops_from_cct.py
+
+Given a .json file in COCO Camera Traps format, creates a cropped image for
+each bounding box.
+
+"""
+
+#%% Imports and constants
+
+import os
+import argparse
+import json
+
+from tqdm import tqdm
+from PIL import Image
+
+
+#%% Functions
+
+def generate_crops_from_cct(cct_file,image_dir,output_dir,padding=0,flat_output=True):
+    """
+    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
+            file should be relative to this folder
+        output_dir (str): the folder where we should write cropped images
+        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 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.
+    """
+
+    ## 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(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)
+            # PIL's crop() method uses exclusive upper bounds for the right and lower
+            # edges, hence "img.width" rather than "img.width-1" here.
+            xmax = min(xmax,img.width)
+            ymax = min(ymax,img.height)
+
+            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_' + str(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
+
+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

@@ -0,0 +1,164 @@
+"""
+
+get_image_sizes.py
+
+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.
+
+"""
+
+#%% Constants and imports
+
+import argparse
+import json
+import os
+import sys
+
+from PIL import Image
+
+from multiprocessing.pool import ThreadPool
+from multiprocessing.pool import Pool
+from functools import partial
+from tqdm import tqdm
+
+from megadetector.utils.path_utils import find_images
+
+image_base = ''
+default_n_threads = 1
+use_threads = False
+
+
+#%% Processing functions
+
+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:
+        pil_im = Image.open(full_path)
+        w = pil_im.width
+        h = pil_im.height
+        return (image_path,w,h)
+    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
+            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
+            base folder for the source images.
+        output_file (str, optional): a .json file to write the image sizes
+        n_workers (int, optional): number of parallel workers to use, set to <=1 to
+            disable parallelization
+        use_threads (bool, optional): whether to use threads (True) or processes (False)
+            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:
+        output_dir = os.path.dirname(output_file)
+        if len(output_dir) > 0:
+            assert os.path.isdir(output_dir), \
+                '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:
+            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)
+
+    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)
+        else:
+            pool = Pool(n_workers)
+        # all_results = list(tqdm(pool.imap(process_image, 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
+
+
+#%% Command-line driver
+
+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')
+    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()