megadetector 10.0.15__py3-none-any.whl

megadetector/classification/merge_classification_detection_output.py

@@ -0,0 +1,509 @@
+"""
+
+merge_classification_detection_output.py
+
+Merges classification results with Batch Detection API outputs.
+
+This script takes 2 main files as input:
+
+1) Either a "dataset CSV" (output of create_classification_dataset.py) or a
+    "classification results CSV" (output of evaluate_model.py). The CSV is
+    expected to have columns listed below. The 'label' and [label names] columns
+    are optional, but at least one of them must be provided.
+    * 'path': str, path to cropped image
+        * if passing in a detections JSON, must match
+            <img_file>___cropXX_mdvY.Y.jpg
+        * if passing in a queried images JSON, must match
+            <dataset>/<img_file>___cropXX_mdvY.Y.jpg or
+            <dataset>/<img_file>___cropXX.jpg
+    * 'label': str, label assigned to this crop
+    * [label names]: float, confidence in each label
+
+2) Either a "detections JSON" (output of MegaDetector) or a "queried images
+    JSON" (output of json_validatory.py).
+
+If the CSV contains [label names] columns (e.g., output of evaluate_model.py),
+then each crop's "classifications" output will have one value per category.
+Categories are sorted decreasing by confidence.
+    "classifications": [
+        ["3", 0.901],
+        ["1", 0.071],
+        ["4", 0.025],
+        ["2", 0.003],
+   ]
+
+If the CSV only contains the 'label' column (e.g., output of
+create_classification_dataset.py), then each crop's "classifications" output
+will have only one value, with a confidence of 1.0. The label's classification
+category ID is always greater than 1,000,000, to distinguish it from a predicted
+category ID.
+    "classifications": [
+        ["1000004", 1.0]
+    ]
+
+If the CSV contains both [label names] and 'label' columns, then both the
+predicted categories and label category will be included. By default, the
+label-category is included last; if the --label-first flag is given, then the
+label category is placed first in the results.
+    "classifications": [
+        ["1000004", 1.0],  # label put first if --label-first flag is given
+        ["3", 0.901],  # all other results are sorted by confidence
+        ["1", 0.071],
+        ["4", 0.025],
+        ["2", 0.003]
+    ]
+
+"""
+
+#%% Imports
+
+from __future__ import annotations
+
+import argparse
+import datetime
+import json
+import os
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+import pandas as pd
+from tqdm import tqdm
+
+from megadetector.utils.ct_utils import round_float
+from megadetector.utils import ct_utils
+
+
+#%% Example usage
+
+"""
+    python merge_classification_detection_output.py \
+        BASE_LOGDIR/LOGDIR/outputs_test.csv.gz \
+        BASE_LOGDIR/label_index.json \
+        BASE_LOGDIR/queried_images.json \
+        --classifier-name "efficientnet-b3-idfg-moredata" \
+        --detector-output-cache-dir $HOME/classifier-training/mdcache \
+        --detector-version "4.1" \
+        --output-json BASE_LOGDIR/LOGDIR/classifier_results.json \
+        --datasets idfg idfg_swwlf_2019
+"""
+
+
+#%% Support functions
+
+def row_to_classification_list(row: Mapping[str, Any],
+                               label_names: Sequence[str],
+                               contains_preds: bool,
+                               label_pos: str | None,
+                               threshold: float,
+                               relative_conf: bool = False
+                               ) -> list[tuple[str, float]]:
+    """
+    Given a mapping from label name to output probability, returns a list of
+    tuples, (str(label_id), prob), which can be serialized into the Batch API
+    output format.
+
+    The list of tuples is returned in sorted order by the predicted probability
+    for each label.
+
+    If 'label' is in row and label_pos is not None, then we add
+    (label_id + 1_000_000, 1.) to the list. If label_pos='first', we put this at
+    the front of the list. Otherwise, we put it at the end.
+    """
+
+    contains_label = ('label' in row)
+    assert contains_label or contains_preds
+    if relative_conf:
+        assert contains_label and contains_preds
+
+    result = []
+    if contains_preds:
+        result = [(str(i), row[label]) for i, label in enumerate(label_names)]
+        if relative_conf:
+            label_conf = row[row['label']]
+            result = [(k, max(v - label_conf, 0)) for k, v in result]
+
+        # filter out confidences below the threshold, and set precision to 4
+        result = [
+            (k, round_float(conf, precision=4))
+            for k, conf in result if conf >= threshold
+        ]
+
+        # sort from highest to lowest probability
+        result = sorted(result, key=lambda x: x[1], reverse=True)
+
+    if contains_label and label_pos is not None:
+        label = row['label']
+        label_id = label_names.index(label)
+        item = (str(label_id + 1_000_000), 1.)
+        if label_pos == 'first':
+            result = [item] + result
+        else:
+            result.append(item)
+    return result
+
+
+def process_queried_images(
+         df: pd.DataFrame,
+         queried_images_json_path: str,
+         detector_output_cache_base_dir: str,
+         detector_version: str,
+         datasets: Sequence[str] | None = None,
+         samples_per_label: int | None = None,
+         seed: int = 123
+         ) -> dict[str, Any]:
+    """
+    Creates a detection JSON object roughly in the Batch API detection
+    format.
+
+    Detections are either ground-truth (from the queried images JSON) or
+    retrieved from the detector output cache. Only images corresponding to crop
+    paths from the given pd.DataFrame are included in the detection JSON.
+
+    Args:
+        df: pd.DataFrame, either a "classification dataset CSV" or a
+            "classification results CSV",  column 'path' has format
+            <dataset>/<img_file>___cropXX[...].jpg
+        queried_images_json_path: str, path to queried images JSON
+        detector_output_cache_base_dir: str
+        detector_version: str
+        datasets: optional list of str, only crops from these datasets will be
+            be included in the output, set to None to include all datasets
+        samples_per_label: optional int, if not None, then randomly sample this
+            many bounding boxes per label (each label must have at least this
+            many examples)
+        seed: int, used for random sampling if samples_per_label is not None
+
+    Returns: dict, detections JSON file, except that the 'images' field is a
+        dict (img_path => dict) instead of a list
+    """
+
+    # input validation
+    assert os.path.exists(queried_images_json_path)
+    detection_cache_dir = os.path.join(
+        detector_output_cache_base_dir, f'v{detector_version}')
+    assert os.path.isdir(detection_cache_dir)
+
+    # extract dataset name from crop path so we can process 1 dataset at a time
+    df['dataset'] = df.index.map(lambda x: x[:x.find('/')])
+    unique_datasets = df['dataset'].unique()
+
+    if datasets is not None:
+        for ds in datasets:
+            assert ds in unique_datasets
+        df = df[df['dataset'].isin(datasets)]  # filter by dataset
+    else:
+        datasets = unique_datasets
+
+    # randomly sample images for each class
+    if samples_per_label is not None:
+        print(f'Sampling {samples_per_label} bounding boxes per label')
+        df = df.groupby('label').sample(samples_per_label, random_state=seed)
+
+    # load queried images JSON, needed for ground-truth bbox info
+    with open(queried_images_json_path, 'r') as f:
+        queried_images_js = json.load(f)
+
+    merged_js: dict[str, Any] = {
+        'images': {},  # start as dict, will convert to list later
+        'info': {}
+    }
+    images = merged_js['images']
+
+    for ds in datasets:
+        print('processing dataset:', ds)
+        ds_df = df[df['dataset'] == ds]
+
+        with open(os.path.join(detection_cache_dir, f'{ds}.json'), 'r') as f:
+            detection_js = json.load(f)
+        img_file_to_index = {
+            im['file']: idx
+            for idx, im in enumerate(detection_js['images'])
+        }
+
+        # compare info dicts
+        class_info = merged_js['info']
+        detection_info = detection_js['info']
+        key = 'detector'
+        if key not in class_info:
+            class_info[key] = detection_info[key]
+        assert class_info[key] == detection_info[key]
+
+        # compare detection categories
+        key = 'detection_categories'
+        if key not in merged_js:
+            merged_js[key] = detection_js[key]
+        assert merged_js[key] == detection_js[key]
+        cat_to_catid = {v: k for k, v in detection_js[key].items()}
+
+        for crop_path in tqdm(ds_df.index):
+            # crop_path: <dataset>/<img_file>___cropXX_mdvY.Y.jpg
+            #            [----<img_path>----]       [-<suffix>--]
+            img_path, suffix = crop_path.split('___crop')
+            img_file = img_path[img_path.find('/') + 1:]
+
+            # file has detection entry
+            if '_mdv' in suffix and img_path not in images:
+                img_idx = img_file_to_index[img_file]
+                images[img_path] = detection_js['images'][img_idx]
+                images[img_path]['file'] = img_path
+
+            # bounding box is from ground truth
+            elif img_path not in images:
+                images[img_path] = {
+                    'file': img_path,
+                    'detections': [
+                        {
+                            'category': cat_to_catid[bbox_dict['category']],
+                            'conf': 1.0,
+                            'bbox': bbox_dict['bbox']
+                        }
+                        for bbox_dict in queried_images_js[img_path]['bbox']
+                    ]
+                }
+    return merged_js
+
+
+def combine_classification_with_detection(
+        detection_js: dict[str, Any],
+        df: pd.DataFrame,
+        idx_to_label: Mapping[str, str],
+        label_names: Sequence[str],
+        classifier_name: str,
+        classifier_timestamp: str,
+        threshold: float,
+        label_pos: str | None = None,
+        relative_conf: bool = False,
+        typical_confidence_threshold: float = None
+        ) -> dict[str, Any]:
+    """
+    Adds classification information to a detection JSON. Classification
+    information may include the true label and/or the predicted confidences
+    of each label.
+
+    Args:
+        detection_js: dict, detections JSON file, except that the 'images'
+            field is a dict (img_path => dict) instead of a list
+        df: pd.DataFrame, classification results, indexed by crop path
+        idx_to_label: dict, str(label_id) => label name, may also include
+            str(label_id + 1e6) => 'label: {label_name}'
+        label_names: list of str, label names
+        classifier_name: str, name of classifier to include in output JSON
+        classifier_timestamp: str, timestamp to include in output JSON
+        threshold: float, for each crop, omit classification results for
+            categories whose confidence is below this threshold
+        label_pos: one of [None, 'first', 'last']
+            None: do not include labels in the output JSON
+            'first' / 'last': position in classification list to put the label
+        relative_conf: bool, if True then for each class, outputs its relative
+            confidence over the confidence of the true label, requires 'label'
+            to be in CSV
+        typical_confidence_threshold: float, useful default confidence
+            threshold; not used directly, just passed along to the output file
+
+    Returns: dict, detections JSON file updated with classification results
+    """
+
+    classification_metadata = {
+        'classifier': classifier_name,
+        'classification_completion_time': classifier_timestamp
+    }
+
+    if typical_confidence_threshold is not None:
+        classification_metadata['classifier_metadata'] = \
+        {'typical_classification_threshold':typical_confidence_threshold}
+
+    detection_js['info'].update(classification_metadata)
+    detection_js['classification_categories'] = idx_to_label
+
+    contains_preds = (set(label_names) <= set(df.columns))
+    if not contains_preds:
+        print('CSV does not contain predictions. Outputting labels only.')
+
+    images = detection_js['images']
+
+    for crop_path in tqdm(df.index):
+        # crop_path: <dataset>/<img_file>___cropXX_mdvY.Y.jpg
+        #            [----<img_path>----]       [-<suffix>--]
+        img_path, suffix = crop_path.split('___crop')
+        crop_index = int(suffix[:2])
+
+        detection_dict = images[img_path]['detections'][crop_index]
+        detection_dict['classifications'] = row_to_classification_list(
+            row=df.loc[crop_path], label_names=label_names,
+            contains_preds=contains_preds, label_pos=label_pos,
+            threshold=threshold, relative_conf=relative_conf)
+
+    detection_js['images'] = list(images.values())
+    return detection_js
+
+
+#%% Main function
+
+def main(classification_csv_path: str,
+         label_names_json_path: str,
+         output_json_path: str,
+         classifier_name: str,
+         threshold: float,
+         datasets: Sequence[str] | None,
+         detection_json_path: str | None,
+         queried_images_json_path: str | None,
+         detector_output_cache_base_dir: str | None,
+         detector_version: str | None,
+         samples_per_label: int | None,
+         seed: int,
+         label_pos: str | None,
+         relative_conf: bool,
+         typical_confidence_threshold: float) -> None:
+
+    # input validation
+    assert os.path.exists(classification_csv_path)
+    assert os.path.exists(label_names_json_path)
+    assert 0 <= threshold <= 1
+    for x in [detection_json_path, queried_images_json_path]:
+        if x is not None:
+            assert os.path.exists(x)
+    assert label_pos in [None, 'first', 'last']
+
+    # load classification CSV
+    print('Loading classification CSV...')
+    df = pd.read_csv(classification_csv_path, float_precision='high',
+                     index_col='path')
+    if relative_conf or label_pos is not None:
+        assert 'label' in df.columns
+
+    # load label names
+    with open(label_names_json_path, 'r') as f:
+        idx_to_label = json.load(f)
+    label_names = [idx_to_label[str(i)] for i in range(len(idx_to_label))]
+    if 'label' in df.columns:
+        for i, label in enumerate(label_names):
+            idx_to_label[str(i + 1_000_000)] = f'label: {label}'
+
+    if queried_images_json_path is not None:
+        assert detector_output_cache_base_dir is not None
+        assert detector_version is not None
+        detection_js = process_queried_images(
+            df=df, queried_images_json_path=queried_images_json_path,
+            detector_output_cache_base_dir=detector_output_cache_base_dir,
+            detector_version=detector_version, datasets=datasets,
+            samples_per_label=samples_per_label, seed=seed)
+    elif detection_json_path is not None:
+        with open(detection_json_path, 'r') as f:
+            detection_js = json.load(f)
+        images = {}
+        for img in detection_js['images']:
+            path = img['file']
+            if datasets is None or path[:path.find('/')] in datasets:
+                images[path] = img
+        detection_js['images'] = images
+
+    classification_time = datetime.date.fromtimestamp(
+        os.path.getmtime(classification_csv_path))
+    classifier_timestamp = classification_time.strftime('%Y-%m-%d %H:%M:%S')
+
+    classification_js = combine_classification_with_detection(
+        detection_js=detection_js, df=df, idx_to_label=idx_to_label,
+        label_names=label_names, classifier_name=classifier_name,
+        classifier_timestamp=classifier_timestamp, threshold=threshold,
+        label_pos=label_pos, relative_conf=relative_conf,
+        typical_confidence_threshold=typical_confidence_threshold)
+
+    os.makedirs(os.path.dirname(output_json_path), exist_ok=True)
+    # The following line was removed as per the previous refactoring:
+    # with open(output_json_path, 'w') as f:
+    #     json.dump(classification_js, f, indent=1)
+    ct_utils.write_json(output_json_path, classification_js)
+
+    print('Wrote merged classification/detection results to {}'.format(output_json_path))
+
+
+#%% Command-line driver
+
+def _parse_args() -> argparse.Namespace:
+
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description='Merges classification results with Batch Detection API '
+                    'outputs.')
+    parser.add_argument(
+        'classification_csv',
+        help='path to classification CSV')
+    parser.add_argument(
+        'label_names_json',
+        help='path to JSON file mapping label index to label name')
+    parser.add_argument(
+        '-o', '--output-json', required=True,
+        help='(required) path to save output JSON with both detection and '
+             'classification results')
+    parser.add_argument(
+        '-n', '--classifier-name', required=True,
+        help='(required) name of classifier')
+    parser.add_argument(
+        '-t', '--threshold', type=float, default=0.1,
+        help='Confidence threshold between 0 and 1. In the output file, omit '
+             'classifier results on classes whose confidence is below this '
+             'threshold.')
+    parser.add_argument(
+        '-d', '--datasets', nargs='*',
+        help='optionally limit output to images from certain datasets. Assumes '
+             'that image paths are given as <dataset>/<img_file>.')
+    parser.add_argument(
+        '--typical-confidence-threshold', type=float, default=None,
+        help='useful default confidence threshold; not used directly, just '
+             'passed along to the output file')
+
+    detection_json_group = parser.add_argument_group(
+        'arguments for passing in a detections JSON file')
+    detection_json_group.add_argument(
+        '-j', '--detection-json',
+        help='path to detections JSON file')
+
+    queried_images_group = parser.add_argument_group(
+        'arguments for passing in a queried images JSON file')
+    queried_images_group.add_argument(
+        '-q', '--queried-images-json',
+        help='path to queried images JSON file')
+    queried_images_group.add_argument(
+        '-c', '--detector-output-cache-dir',
+        help='(required) path to directory where detector outputs are cached')
+    queried_images_group.add_argument(
+        '-v', '--detector-version',
+        help='(required) detector version string, e.g., "4.1"')
+    queried_images_group.add_argument(
+        '-s', '--samples-per-label', type=int,
+        help='randomly sample this many bounding boxes per label (each label '
+             'must have at least this many examples)')
+    queried_images_group.add_argument(
+        '--seed', type=int, default=123,
+        help='random seed, only used if --samples-per-label is given')
+    queried_images_group.add_argument(
+        '--label', choices=['first', 'last'], default=None,
+        help='Whether to put the label first or last in the list of '
+             'classifications. If this argument is omitted, then no labels are '
+             'included in the output.')
+    queried_images_group.add_argument(
+        '--relative-conf', action='store_true',
+        help='for each class, outputs its relative confidence over the '
+             'confidence of the true label, requires "label" to be in CSV')
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+
+    args = _parse_args()
+    main(classification_csv_path=args.classification_csv,
+         label_names_json_path=args.label_names_json,
+         output_json_path=args.output_json,
+         classifier_name=args.classifier_name,
+         threshold=args.threshold,
+         datasets=args.datasets,
+         detection_json_path=args.detection_json,
+         queried_images_json_path=args.queried_images_json,
+         detector_output_cache_base_dir=args.detector_output_cache_dir,
+         detector_version=args.detector_version,
+         samples_per_label=args.samples_per_label,
+         seed=args.seed,
+         label_pos=args.label,
+         relative_conf=args.relative_conf,
+         typical_confidence_threshold=args.typical_confidence_threshold)