megadetector 10.0.15__py3-none-any.whl

megadetector/postprocessing/subset_json_detector_output.py

@@ -0,0 +1,964 @@
+r"""
+
+subset_json_detector_output.py
+
+Creates one or more subsets of a detector results file (.json), doing either
+or both of the following (if both are requested, they happen in this order):
+
+1) Retrieve all elements where filenames contain a specified query string,
+   optionally replacing that query with a replacement token. If the query is blank,
+   can also be used to prepend content to all filenames.
+
+   Does not support regex's, but supports a special case of ^string to indicate "must start with
+   to match".
+
+2) Create separate .jsons for each unique path, optionally making the filenames
+   in those .json's relative paths.  In this case, you specify an output directory,
+   rather than an output path.  All images in the folder blah/foo/bar will end up
+   in a .json file called blah_foo_bar.json.
+
+Can also apply a confidence threshold.
+
+Can also subset by categories above a threshold (programmatic invocation only, this is
+not supported at the command line yet).
+
+To subset a COCO Camera Traps .json database, see subset_json_db.py
+
+**Sample invocation (splitting into multiple json's)**
+
+Read from "1800_idfg_statewide_wolf_detections_w_classifications.json", split up into
+individual .jsons in 'd:/temp/idfg/output', making filenames relative to their individual
+folders:
+
+python subset_json_detector_output.py ^
+  "d:/temp/idfg/1800_idfg_statewide_wolf_detections_w_classifications.json" "d:/temp/idfg/output" ^
+  --split_folders --make_folder_relative
+
+Now do the same thing, but instead of writing .json's to d:/temp/idfg/output, write them to *subfolders*
+corresponding to the subfolders for each .json file.
+
+python subset_json_detector_output.py ^
+  "d:/temp/idfg/1800_detections_S2.json" "d:/temp/idfg/output_to_folders" ^
+  --split_folders --make_folder_relative --copy_jsons_to_folders
+
+**Sample invocation (creating a single subset matching a query)**
+
+Read from "1800_detections.json", write to "1800_detections_2017.json"
+
+Include only images matching "2017", and change "2017" to "blah"
+
+python subset_json_detector_output.py "d:/temp/1800_detections.json" "d:/temp/1800_detections_2017_blah.json" ^
+    --query 2017 --replacement blah
+
+Include all images, prepend with "prefix/"
+
+python subset_json_detector_output.py "d:/temp/1800_detections.json" "d:/temp/1800_detections_prefix.json" ^
+    --replacement "prefix/"
+
+"""
+
+#%% Constants and imports
+
+import argparse
+import sys
+import copy
+import json
+import os
+import re
+
+from tqdm import tqdm
+
+from megadetector.utils import ct_utils
+from megadetector.utils.ct_utils import args_to_object, get_max_conf, invert_dictionary
+from megadetector.utils.path_utils import recursive_file_list
+
+
+#%% Helper classes
+
+class SubsetJsonDetectorOutputOptions:
+    """
+    Options used to parameterize subset_json_detector_output()
+    """
+
+    def __init__(self):
+
+        #: Only process files containing the token 'query'
+        #:
+        #: Does not support general regexes, but supports ^ as a special case
+        #: regex-like notation for "starts with"
+        self.query = None
+
+        #: Replace 'query' with 'replacement' if 'replacement' is not None.  If 'query' is None,
+        #: prepend 'replacement'
+        self.replacement = None
+
+        #: Should we split output into individual .json files for each folder?
+        self.split_folders = False
+
+        #: Folder level to use for splitting ['bottom','n_from_bottom','n_from_top','dict']
+        #:
+        #: 'dict' requires 'split_folder_param' to be a dictionary mapping each filename
+        #: to a token.
+        self.split_folder_mode = 'bottom'
+
+        #: When using the 'n_from_bottom' parameter to define folder splitting, this
+        #: defines the number of directories from the bottom.  'n_from_bottom' with
+        #: a parameter of zero is the same as 'bottom'.
+        #:
+        #: Same story with 'n_from_top'.
+        #:
+        #: When 'split_folder_mode' is 'dict', this should be a dictionary mapping each filename
+        #: to a token.
+        self.split_folder_param = 0
+
+        #: Only meaningful if split_folders is True: should we convert pathnames to be relative
+        #: the folder for each .json file?
+        self.make_folder_relative = False
+
+        #: Only meaningful if split_folders and make_folder_relative are True: if not None,
+        #: will copy .json files to their corresponding output directories, relative to
+        #: output_filename
+        self.copy_jsons_to_folders = False
+
+        #: Should we over-write .json files?
+        self.overwrite_json_files = False
+
+        #: If copy_jsons_to_folders is true, do we require that directories already exist?
+        self.copy_jsons_to_folders_directories_must_exist = True
+
+        #: Optional confidence threshold; if not None, detections below this confidence won't be
+        #: included in the output.
+        self.confidence_threshold = None
+
+        #: Should we remove failed images?
+        self.remove_failed_images = False
+
+        #: Either a list of category IDs (as string-ints) (not names), or a dictionary mapping category *IDs*
+        #: (as string-ints) (not names) to thresholds.  Removes non-matching detections, does not
+        #: remove images.  Not technically mutually exclusize with category_names_to_keep, but it's an esoteric
+        #: scenario indeed where you would want to specify both.
+        self.categories_to_keep = None
+
+        #: Either a list of category names (not IDs), or a dictionary mapping category *names* (not IDs) to thresholds.
+        #: Removes non-matching detections, does not remove images.  Not technically mutually exclusize with
+        #: category_ids_to_keep, but it's an esoteric scenario indeed where you would want to specify both.
+        self.category_names_to_keep = None
+
+        #: Set to >0 during testing to limit the number of images that get processed.
+        self.debug_max_images = -1
+
+        #: Keep only files in this list, which can be a list, a .json results file, or a folder.
+        #
+        #: Assumes that the input .json file contains relative paths when comparing to a folder.
+        self.keep_files_in_list = None
+
+        #: Remove classification with <= N instances.  Does not re-map categories
+        #: to be contiguous.  Set to 1 to remove empty categories only.
+        self.remove_classification_categories_below_count = None
+
+        #: Remove detections above a threshold size (as a fraction of the image size)
+        self.maximum_detection_size = None
+
+        #: Remove detections below a threshold size (as a fraction of the image size)
+        self.minimum_detection_size = None
+
+# ...class SubsetJsonDetectorOutputOptions
+
+
+#%% Main function
+
+def _write_detection_results(data, output_filename, options):
+    """
+    Writes the detector-output-formatted dict *data* to *output_filename*.
+    """
+
+    if (not options.overwrite_json_files) and os.path.isfile(output_filename):
+        raise ValueError('File {} exists'.format(output_filename))
+
+    basedir = os.path.dirname(output_filename)
+
+    if options.copy_jsons_to_folders and options.copy_jsons_to_folders_directories_must_exist:
+        if not os.path.isdir(basedir):
+            raise ValueError('Directory {} does not exist'.format(basedir))
+    else:
+        os.makedirs(basedir, exist_ok=True)
+
+    n_images = len(data['images'])
+
+    print('Writing detection output (with {} images) to {}'.format(n_images,output_filename))
+    ct_utils.write_json(output_filename, data)
+
+# ...def _write_detection_results(...)
+
+
+def remove_classification_categories_below_count(data, options):
+    """
+    Removes all classification categories below a threshold count.  Does not re-map
+    classification category IDs.
+
+    Args:
+        data (dict): data loaded from a MD results file
+        options (SubsetJsonDetectorOutputOptions): parameters for subsetting
+
+    Returns:
+        dict: Possibly-modified version of [data] (also modifies in place)
+    """
+
+    if options.remove_classification_categories_below_count is None:
+        return data
+    if 'classification_categories' not in data:
+        return data
+
+    classification_category_id_to_count = {}
+
+    for classification_category_id in data['classification_categories']:
+        classification_category_id_to_count[classification_category_id] = 0
+
+    # Count the number of occurrences of each classification category
+    for im in data['images']:
+        if 'detections' not in im or im['detections'] is None:
+            continue
+        for det in im['detections']:
+            if 'classifications' not in det:
+                continue
+            for classification in det['classifications']:
+                classification_category_id_to_count[classification[0]] = \
+                    classification_category_id_to_count[classification[0]] + 1
+
+
+    # Which categories have above-threshold counts?
+    classification_category_ids_to_keep = set()
+
+    for classification_category_id in classification_category_id_to_count:
+        if classification_category_id_to_count[classification_category_id] >= \
+            options.remove_classification_categories_below_count:
+                classification_category_ids_to_keep.add(classification_category_id)
+
+    n_categories_removed = \
+        len(classification_category_id_to_count) - \
+        len(classification_category_ids_to_keep)
+
+    print('Removing {} of {} classification categories'.format(
+        n_categories_removed,len(classification_category_id_to_count)))
+
+    if n_categories_removed == 0:
+        return data
+
+    # Filter the category list
+    output_classification_categories = {}
+    for category_id in data['classification_categories']:
+        if category_id in classification_category_ids_to_keep:
+            output_classification_categories[category_id] = \
+                data['classification_categories'][category_id]
+    data['classification_categories'] = output_classification_categories
+    assert len(data['classification_categories']) == len(classification_category_ids_to_keep)
+
+    # If necessary, filter the category descriptions
+    if 'classification_category_descriptions' in data:
+        output_classification_category_descriptions = {}
+        for category_id in data['classification_category_descriptions']:
+            if category_id in classification_category_ids_to_keep:
+                output_classification_category_descriptions[category_id] = \
+                    data['classification_category_descriptions'][category_id]
+        data['classification_category_descriptions'] = output_classification_category_descriptions
+
+    # Filter images
+    for im in data['images']:
+        if 'detections' not in im or im['detections'] is None:
+            continue
+        for det in im['detections']:
+            if 'classifications' not in det:
+                continue
+            classifications_to_keep = []
+            for classification in det['classifications']:
+                if classification[0] in classification_category_ids_to_keep:
+                    classifications_to_keep.append(classification)
+            det['classifications'] = classifications_to_keep
+
+    return data
+
+# ...def remove_classification_categories_below_count(...)
+
+
+def subset_json_detector_output_by_size(data, options):
+    """
+    Remove detections above or below threshold sizes (as a fraction
+    of the image size).
+
+    Args:
+        data (dict): data loaded from a MD results file
+        options (SubsetJsonDetectorOutputOptions): parameters for subsetting
+
+    Returns:
+        dict: Possibly-modified version of [data] (also modifies in place)
+    """
+
+    if (options.maximum_detection_size is None) and \
+        (options.minimum_detection_size is None):
+        return data
+
+    if options.maximum_detection_size is None:
+        options.maximum_detection_size = 1000
+
+    if options.minimum_detection_size is None:
+        options.minimum_detection_size = -1000
+
+    print('Subsetting by size ({} <--> {})'.format(
+        options.minimum_detection_size,
+        options.maximum_detection_size))
+
+    images_in = data['images']
+    images_out = []
+
+    # im = images_in[0]
+    for i_image, im in tqdm(enumerate(images_in), total=len(images_in)):
+
+        # Always keep failed images; if the caller wants to remove these, they
+        # will use remove_failed_images
+        if ('detections' not in im) or (im['detections'] is None):
+            images_out.append(im)
+            continue
+
+        detections_to_keep = []
+
+        for det in im['detections']:
+
+            # [x_min, y_min, width_of_box, height_of_box]
+            detection_size = det['bbox'][2] * det['bbox'][3]
+
+            if (detection_size >= options.minimum_detection_size) and \
+               (detection_size <= options.maximum_detection_size):
+                detections_to_keep.append(det)
+
+        im['detections'] = detections_to_keep
+
+        images_out.append(im)
+
+    # ...for each image
+
+    data['images'] = images_out
+    print('done, found {} matches (of {})'.format(
+            len(data['images']),len(images_in)))
+
+    return data
+
+# ...def subset_json_detector_output_by_size(...)
+
+
+def subset_json_detector_output_by_confidence(data, options):
+    """
+    Removes all detections below options.confidence_threshold.
+
+    Args:
+        data (dict): data loaded from a MD results file
+        options (SubsetJsonDetectorOutputOptions): parameters for subsetting
+
+    Returns:
+        dict: Possibly-modified version of [data] (also modifies in place)
+    """
+
+    if options.confidence_threshold is None:
+        return data
+
+    images_in = data['images']
+    images_out = []
+
+    print('Subsetting by confidence >= {}'.format(options.confidence_threshold))
+
+    n_max_changes = 0
+
+    # im = images_in[0]
+    for i_image, im in tqdm(enumerate(images_in), total=len(images_in)):
+
+        # Always keep failed images; if the caller wants to remove these, they
+        # will use remove_failed_images
+        if ('detections' not in im) or (im['detections'] is None):
+            images_out.append(im)
+            continue
+
+        p_orig = get_max_conf(im)
+
+        # Find all detections above threshold for this image
+        detections = [d for d in im['detections'] if d['conf'] >= options.confidence_threshold]
+
+        # If there are no detections above threshold, set the max probability
+        # to -1, unless it already had a negative probability.
+        if len(detections) == 0:
+            if p_orig <= 0:
+                p = p_orig
+            else:
+                p = -1
+
+        # Otherwise find the max confidence
+        else:
+            p = max([d['conf'] for d in detections])
+
+        im['detections'] = detections
+
+        # Did this thresholding result in a max-confidence change?
+        if abs(p_orig - p) > 0.00001:
+
+            # We should only be *lowering* max confidence values (i.e., making them negative)
+            assert (p_orig <= 0) or (p < p_orig), \
+                'Confidence changed from {} to {}'.format(p_orig, p)
+            n_max_changes += 1
+
+        if 'max_detection_conf' in im:
+            im['max_detection_conf'] = p
+
+        images_out.append(im)
+
+    # ...for each image
+
+    data['images'] = images_out
+    print('done, found {} matches (of {}), {} max conf changes'.format(
+            len(data['images']),len(images_in),n_max_changes))
+
+    return data
+
+# ...def subset_json_detector_output_by_confidence(...)
+
+
+def subset_json_detector_output_by_list(data, options):
+    """
+    Keeps only files in options.keep_files_in_list, which can be a .json results file or a folder.
+    Assumes that the input .json file contains relative paths when comparing to a folder.
+
+    Args:
+        data (dict): data loaded from a MD results file
+        options (SubsetJsonDetectorOutputOptions): parameters for subsetting
+
+    Returns:
+        dict: Possibly-modified version of [data] (also modifies in place)
+    """
+
+    if options.keep_files_in_list is None:
+        return data
+
+    files_to_keep = None
+
+    if isinstance(options.keep_files_in_list,list):
+        files_to_keep = options.keep_files_in_list
+    elif os.path.isfile(options.keep_files_in_list):
+        with open(options.keep_files_in_list,'r') as f:
+            d = json.load(f)
+        files_to_keep = [im['file'] for im in d['images']]
+    elif os.path.isdir(options.keep_files_in_list):
+        files_to_keep = \
+            recursive_file_list(options.keep_files_in_list,return_relative_paths=True)
+    else:
+        raise ValueError('Subsetting .json file by list: {} is neither a .json results file nor a folder'.format(
+            options.keep_files_in_list))
+
+    files_to_keep = [fn.replace('\\','/') for fn in files_to_keep]
+    files_to_keep_set = set(files_to_keep)
+
+    images_to_keep = []
+
+    for im in data['images']:
+        fn = im['file'].replace('\\','/')
+        if fn in files_to_keep_set:
+            images_to_keep.append(im)
+
+    print('Subsetting by list kept {} of {} files (expected {})'.format(
+        len(images_to_keep),len(data['images']),len(files_to_keep)))
+
+    data['images'] = images_to_keep
+
+    return data
+
+# ...def subset_json_detector_output_by_list(...)
+
+
+def subset_json_detector_output_by_categories(data, options):
+    """
+    Removes all detections without detections above a threshold for specific categories.
+
+    Args:
+        data (dict): data loaded from a MD results file
+        options (SubsetJsonDetectorOutputOptions): parameters for subsetting
+
+    Returns:
+        dict: Possibly-modified version of [data] (also modifies in place)
+    """
+
+    # If categories_to_keep is supplied as a list, convert to a dict
+    if options.categories_to_keep is not None:
+        if not isinstance(options.categories_to_keep, dict):
+            dict_categories_to_keep = {}
+            for category_id in options.categories_to_keep:
+                # Set unspecified thresholds to a silly negative value
+                dict_categories_to_keep[category_id] = -100000.0
+            options.categories_to_keep = dict_categories_to_keep
+
+    # If category_names_to_keep is supplied as a list, convert to a dict
+    if options.category_names_to_keep is not None:
+        if not isinstance(options.category_names_to_keep, dict):
+            dict_category_names_to_keep = {}
+            for category_name in options.category_names_to_keep:
+                # Set unspecified thresholds to a silly negative value
+                dict_category_names_to_keep[category_name] = -100000.0
+            options.category_names_to_keep = dict_category_names_to_keep
+
+    category_name_to_category_id = invert_dictionary(data['detection_categories'])
+
+    # If some categories are supplied as names, convert all to IDs and add to "categories_to_keep"
+    if options.category_names_to_keep is not None:
+        if options.categories_to_keep is None:
+            options.categories_to_keep = {}
+        for category_name in options.category_names_to_keep:
+            assert category_name in category_name_to_category_id, \
+                'Category {} not in detection categories'.format(category_name)
+            category_id = category_name_to_category_id[category_name]
+            assert category_id not in options.categories_to_keep, \
+                'Category {} ({}) specified as both a name and an ID'.format(
+                    category_name,category_id)
+            options.categories_to_keep[category_id] = options.category_names_to_keep[category_name]
+
+    if options.categories_to_keep is None:
+        return data
+
+    images_in = data['images']
+    images_out = []
+
+    print('Subsetting by categories (keeping {} categories):'.format(
+        len(options.categories_to_keep)))
+
+    for category_id in sorted(list(options.categories_to_keep.keys())):
+        if category_id not in data['detection_categories']:
+            print('Warning: category ID {} not in category map in this file'.format(category_id))
+        else:
+            print('{} ({}) (threshold {})'.format(
+                category_id,
+                data['detection_categories'][category_id],
+                options.categories_to_keep[category_id]))
+
+    n_detections_in = 0
+    n_detections_kept = 0
+
+    # im = images_in[0]
+    for i_image, im in tqdm(enumerate(images_in), total=len(images_in)):
+
+        # Always keep failed images; if the caller wants to remove these, they
+        # will use remove_failed_images
+        if ('detections' not in im) or (im['detections'] is None):
+            images_out.append(im)
+            continue
+
+        n_detections_in += len(im['detections'])
+
+        # Find all matching detections for this image
+        detections = []
+        for d in im['detections']:
+            if (d['category'] in options.categories_to_keep) and \
+               (d['conf'] > options.categories_to_keep[d['category']]):
+               detections.append(d)
+
+        im['detections'] = detections
+
+        if 'max_detection_conf' in im:
+            if len(detections) == 0:
+                p = 0
+            else:
+                p = max([d['conf'] for d in detections])
+            im['max_detection_conf'] = p
+
+        n_detections_kept += len(im['detections'])
+
+        images_out.append(im)
+
+    # ...for each image
+
+    data['images'] = images_out
+    print('done, kept {} detections (of {})'.format(
+        n_detections_kept,n_detections_in))
+
+    return data
+
+# ...def subset_json_detector_output_by_categories(...)
+
+
+def remove_failed_images(data,options):
+    """
+    Removed failed images from [data]
+
+    Args:
+        data (dict): data loaded from a MD results file
+        options (SubsetJsonDetectorOutputOptions): parameters for subsetting
+
+    Returns:
+        dict: Possibly-modified version of [data] (also modifies in place)
+    """
+
+    images_in = data['images']
+    images_out = []
+
+    if not options.remove_failed_images:
+        return data
+
+    print('Removing failed images...', end='')
+
+    # i_image = 0; im = images_in[0]
+    for i_image, im in tqdm(enumerate(images_in), total=len(images_in)):
+
+        if 'failure' in im and isinstance(im['failure'],str):
+            continue
+        else:
+            images_out.append(im)
+
+    # ...for each image
+
+    data['images'] = images_out
+    n_removed = len(images_in) - len(data['images'])
+    print('Done, removed {} of {}'.format(n_removed, len(images_in)))
+
+    return data
+
+# ...def remove_failed_images(...)
+
+
+def subset_json_detector_output_by_query(data, options):
+    """
+    Subsets to images whose filename matches options.query; replace all instances of
+    options.query with options.replacement.  No-op if options.query_string is None or ''.
+
+    Args:
+        data (dict): data loaded from a MD results file
+        options (SubsetJsonDetectorOutputOptions): parameters for subsetting
+
+    Returns:
+        dict: Possibly-modified version of [data] (also modifies in place)
+    """
+
+    images_in = data['images']
+    images_out = []
+
+    print('Subsetting by query {}, replacement {}...'.format(options.query, options.replacement), end='')
+
+    query_string = options.query
+    query_starts_with = False
+
+    # Support a special case regex-like notation for "starts with"
+    if query_string is not None and query_string.startswith('^'):
+        query_string = query_string[1:]
+        query_starts_with = True
+
+    # i_image = 0; im = images_in[0]
+    for i_image, im in tqdm(enumerate(images_in), total=len(images_in)):
+
+        fn = im['file']
+
+        # Only take images that match the query
+        if query_string is not None:
+            if query_starts_with:
+                if (not fn.startswith(query_string)):
+                    continue
+            else:
+                if query_string not in fn:
+                    continue
+
+        if options.replacement is not None:
+            if query_string is not None:
+                fn = fn.replace(query_string, options.replacement)
+            else:
+                fn = options.replacement + fn
+
+        im['file'] = fn
+
+        images_out.append(im)
+
+    # ...for each image
+
+    data['images'] = images_out
+    print('done, found {} matches (of {})'.format(len(data['images']), len(images_in)))
+
+    return data
+
+# ...def subset_json_detector_output_by_query(...)
+
+
+def subset_json_detector_output(input_filename, output_filename, options, data=None):
+    """
+    Main entry point; creates one or more subsets of a detector results file.  See the
+    module header comment for more information about the available subsetting approaches.
+
+    Makes a copy of [data] before modifying if a data dictionary is supplied.
+
+    Args:
+        input_filename (str): filename to load and subset; can be None if [data] is supplied
+        output_filename (str): file or folder name (depending on [options]) to which we should
+            write subset results.
+        options (SubsetJsonDetectorOutputOptions): parameters for .json splitting/subsetting;
+            see SubsetJsonDetectorOutputOptions for details.
+        data (dict, optional): data loaded from a .json file; if this is not None, [input_filename]
+            will be ignored.  If supplied, this will be copied before it's modified.
+
+    Returns:
+        dict: Results that are either loaded from [input_filename] and processed, or copied
+        from [data] and processed.
+    """
+
+    if options is None:
+        options = SubsetJsonDetectorOutputOptions()
+    else:
+        options = copy.deepcopy(options)
+
+    # Input validation
+    if options.copy_jsons_to_folders:
+        assert options.split_folders and options.make_folder_relative, \
+            'copy_jsons_to_folders set without make_folder_relative and split_folders'
+
+    if options.split_folders:
+        if os.path.isfile(output_filename):
+            raise ValueError('When splitting by folders, output must be a valid directory name, ' + \
+                             'you specified an existing file')
+
+    if data is None:
+        print('Reading file {}'.format(input_filename))
+        with open(input_filename) as f:
+            data = json.load(f)
+        print('Read {} images'.format(len(data['images'])))
+        if options.debug_max_images > 0:
+            print('Trimming to {} images'.format(options.debug_max_images))
+            data['images'] = data['images'][:options.debug_max_images]
+    else:
+        print('Copying data')
+        data = copy.deepcopy(data)
+        print('...done')
+
+    if options.query is not None:
+
+        data = subset_json_detector_output_by_query(data, options)
+
+    if options.remove_failed_images:
+
+        data = remove_failed_images(data, options)
+
+    if options.confidence_threshold is not None:
+
+        data = subset_json_detector_output_by_confidence(data, options)
+
+    if (options.categories_to_keep is not None) or (options.category_names_to_keep is not None):
+
+        data = subset_json_detector_output_by_categories(data, options)
+
+    if options.remove_classification_categories_below_count is not None:
+
+        data = remove_classification_categories_below_count(data, options)
+
+    if options.keep_files_in_list is not None:
+
+        data = subset_json_detector_output_by_list(data, options)
+
+    if (options.maximum_detection_size is not None) or \
+        (options.minimum_detection_size is not None):
+
+        data = subset_json_detector_output_by_size(data, options)
+
+    if not options.split_folders:
+
+        _write_detection_results(data, output_filename, options)
+        return data
+
+    else:
+
+        # Map images to unique folders
+        print('Finding unique folders')
+
+        folders_to_images = {}
+
+        # im = data['images'][0]
+        for im in tqdm(data['images']):
+
+            fn = im['file']
+
+            if options.split_folder_mode == 'bottom':
+
+                dirname = os.path.dirname(fn)
+
+            elif options.split_folder_mode == 'n_from_bottom':
+
+                dirname = os.path.dirname(fn)
+                for n in range(0, options.split_folder_param):
+                    dirname = os.path.dirname(dirname)
+
+            elif options.split_folder_mode == 'n_from_top':
+
+                # Split string into folders, keeping delimiters
+
+                # Don't use this, it removes delimiters
+                # tokens = _split_path(fn)
+                tokens = re.split(r'([\\/])',fn)
+
+                n_tokens_to_keep = ((options.split_folder_param + 1) * 2) - 1
+
+                if n_tokens_to_keep > len(tokens):
+                    raise ValueError('Cannot walk {} folders from the top in path {}'.format(
+                                options.split_folder_param, fn))
+                dirname = ''.join(tokens[0:n_tokens_to_keep])
+
+            elif options.split_folder_mode == 'dict':
+
+                assert isinstance(options.split_folder_param, dict)
+                dirname = options.split_folder_param[fn]
+
+            else:
+
+                raise ValueError('Unrecognized folder split mode {}'.format(options.split_folder_mode))
+
+            folders_to_images.setdefault(dirname, []).append(im)
+
+        # ...for each image
+
+        print('Found {} unique folders'.format(len(folders_to_images)))
+
+        # Optionally make paths relative
+        # dirname = list(folders_to_images.keys())[0]
+        if options.make_folder_relative:
+
+            print('Converting database-relative paths to individual-json-relative paths...')
+
+            for dirname in tqdm(folders_to_images):
+                # im = folders_to_images[dirname][0]
+                for im in folders_to_images[dirname]:
+                    fn = im['file']
+                    relfn = os.path.relpath(fn, dirname).replace('\\', '/')
+                    im['file'] = relfn
+
+        # ...if we need to convert paths to be folder-relative
+
+        print('Finished converting to json-relative paths, writing output')
+
+        os.makedirs(output_filename, exist_ok=True)
+        all_images = data['images']
+
+        # dirname = list(folders_to_images.keys())[0]
+        for dirname in tqdm(folders_to_images):
+
+            json_fn = dirname.replace('/', '_').replace('\\', '_') + '.json'
+
+            if options.copy_jsons_to_folders:
+                json_fn = os.path.join(output_filename, dirname, json_fn)
+            else:
+                json_fn = os.path.join(output_filename, json_fn)
+
+            # Recycle the 'data' struct, replacing 'images' every time... medium-hacky, but
+            # forward-compatible in that I don't take dependencies on the other fields
+            dir_data = data
+            dir_data['images'] = folders_to_images[dirname]
+            _write_detection_results(dir_data, json_fn, options)
+            print('Wrote {} images to {}'.format(len(dir_data['images']), json_fn))
+
+        # ...for each directory
+
+        data['images'] = all_images
+
+        return data
+
+    # ...if we're splitting folders
+
+# ...def subset_json_detector_output(...)
+
+
+#%% Interactive driver
+
+if False:
+
+    #%%
+
+    #%% Subset a file without splitting
+
+    input_filename = r"c:\temp\sample.json"
+    output_filename = r"c:\temp\output.json"
+
+    options = SubsetJsonDetectorOutputOptions()
+    options.replacement = None
+    options.query = 'S2'
+
+    data = subset_json_detector_output(input_filename,output_filename,options,None)
+
+
+    #%% Subset and split, but don't copy to individual folders
+
+    input_filename = r"C:\temp\xxx-export.json"
+    output_filename = r"c:\temp\out"
+
+    options = SubsetJsonDetectorOutputOptions()
+    options.split_folders = True
+    options.make_folder_relative = True
+    options.split_folder_mode = 'n_from_top'
+    options.split_folder_param = 1
+
+    data = subset_json_detector_output(input_filename,output_filename,options,None)
+
+
+    #%% Subset and split, copying to individual folders
+
+    input_filename = r"c:\temp\sample.json"
+    output_filename = r"c:\temp\out"
+
+    options = SubsetJsonDetectorOutputOptions()
+    options.split_folders = True
+    options.make_folder_relative = True
+    options.copy_jsons_to_folders = True
+
+    data = subset_json_detector_output(input_filename,output_filename,options,data)
+
+
+#%% Command-line driver
+
+def main(): # noqa
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('input_file', type=str, help='Input .json filename')
+    parser.add_argument('output_file', type=str, help='Output .json filename')
+    parser.add_argument('--query', type=str, default=None,
+                        help='Query string to search for (omitting this matches all)')
+    parser.add_argument('--replacement', type=str, default=None,
+                        help='Replace [query] with this')
+    parser.add_argument('--confidence_threshold', type=float, default=None,
+                        help='Remove detections below this confidence level')
+    parser.add_argument('--maximum_detection_size', type=float, default=None,
+                        help='Remove detections above this size (as a fraction of the image size)')
+    parser.add_argument('--minimum_detection_size', type=float, default=None,
+                        help='Remove detections below this size (as a fraction of the image size)')
+    parser.add_argument('--keep_files_in_list', type=str, default=None,
+                        help='Keep only files in this list, which can be a .json results file or a folder.' + \
+                             ' Assumes that the input .json file contains relative paths when comparing to a folder.')
+    parser.add_argument('--split_folders', action='store_true',
+                        help='Split .json files by leaf-node folder')
+    parser.add_argument('--split_folder_param', type=int,
+                        help='Directory level count for n_from_bottom and n_from_top splitting')
+    parser.add_argument('--split_folder_mode', type=str,
+                        help='Folder level to use for splitting ("bottom", "n_from_bottom", or "n_from_top")')
+    parser.add_argument('--make_folder_relative', action='store_true',
+                        help='Make image paths relative to their containing folder ' + \
+                             '(only meaningful with split_folders)')
+    parser.add_argument('--overwrite_json_files', action='store_true',
+                        help='Overwrite output files')
+    parser.add_argument('--copy_jsons_to_folders', action='store_true',
+                        help='When using split_folders and make_folder_relative, copy jsons to their ' + \
+                             'corresponding folders (relative to output_file)')
+    parser.add_argument('--create_folders', action='store_true',
+                        help='When using copy_jsons_to_folders, create folders that don''t exist')
+    parser.add_argument('--remove_classification_categories_below_count', type=int, default=None,
+                        help='Remove classification categories with less than this many instances ' + \
+                             '(no removal by default)')
+
+    if len(sys.argv[1:]) == 0:
+        parser.print_help()
+        parser.exit()
+
+    args = parser.parse_args()
+
+    # Convert to an options object
+    options = SubsetJsonDetectorOutputOptions()
+    if args.create_folders:
+        options.copy_jsons_to_folders_directories_must_exist = False
+
+    args_to_object(args, options)
+
+    subset_json_detector_output(args.input_file, args.output_file, options)
+
+if __name__ == '__main__':
+    main()