diff --git a/README.ipynb b/README.ipynb index 8af4116..322e181 100644 --- a/README.ipynb +++ b/README.ipynb @@ -17,7 +17,38 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## Features" + "## Command Line Interface" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Batch Detection\n", + "\n", + "CamTrapML includes a Command Line Interface (CLI) for running batch detection. The interface is designed to follow that of the MegaDetector run_batch_detection.py script, currently with some caveats. Currently the CLI assumes reccursive directory search for images and does not support image_queue which means the MegaDetector run_batch_detection.py will likely process faster. \n", + "\n", + " $ python -m camtrapml.scripts.batch_detection model_name input_dir output_json_path\n", + "\n", + "The options available for `model_name` are:\n", + "\n", + "* `md2`: MegaDetector v2\n", + "* `md3`: MegaDetector v3\n", + "* `md4`: MegaDetector v4.1\n", + "\n", + "Add `--output_relative_filenames` to output filenames relative to the `input_dir` in the output JSON.\n", + "\n", + "E.g.\n", + "\n", + " $ python -m camtrapml.scripts.batch_detection --output_relative_filenames md4 ~/Survey2022 ~/Survey2022/md.4.1.0.json \n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Programatic Usage" ] }, { diff --git a/README.md b/README.md index 64b9510..3b67365 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,29 @@ $ pip install camtrapml -## Features +## Command Line Interface + +### Batch Detection + +CamTrapML includes a Command Line Interface (CLI) for running batch detection. The interface is designed to follow that of the MegaDetector run_batch_detection.py script, currently with some caveats. Currently the CLI assumes reccursive directory search for images and does not support image_queue which means the MegaDetector run_batch_detection.py will likely process faster. + + $ python -m camtrapml.scripts.batch_detection model_name input_dir output_json_path + +The options available for `model_name` are: + +* `md2`: MegaDetector v2 +* `md3`: MegaDetector v3 +* `md4`: MegaDetector v4.1 + +Add `--output_relative_filenames` to output filenames relative to the `input_dir` in the output JSON. + +E.g. + + $ python -m camtrapml.scripts.batch_detection --output_relative_filenames md4 ~/Survey2022 ~/Survey2022/md.4.1.0.json + + + +## Programatic Usage ### Loading Data diff --git a/camtrapml/scripts/batch_detection.py b/camtrapml/scripts/batch_detection.py new file mode 100644 index 0000000..e0d6015 --- /dev/null +++ b/camtrapml/scripts/batch_detection.py @@ -0,0 +1,112 @@ +""" +Command Line Utility for batch detection. +""" + +from argparse import ArgumentParser +from json import dump +from datetime import datetime +from tqdm import tqdm +from camtrapml.dataset import ImageDataset +from camtrapml.detection.models.megadetector import ( + MegaDetectorV2, + MegaDetectorV3, + MegaDetectorV4_1, +) + + +def parse_args(): + """ + Parse command line arguments. + """ + + parser = ArgumentParser() + parser.add_argument( + "model", + type=str, + help="Detection model to utilise [md2, md3, md4]", + ) + parser.add_argument( + "dataset_path", + type=str, + help="Path to directory containing the images", + ) + parser.add_argument("output_path", type=str, help="Path to store the JSON output") + + parser.add_argument("--output_relative_filenames", action="store_true") + + return parser.parse_args() + + +def get_model(model_name): + """ + Load a detection model based on a short name. + """ + + if model_name == "md2": + return MegaDetectorV2() + + if model_name == "md3": + return MegaDetectorV3() + + if model_name == "md4": + return MegaDetectorV4_1() + + raise ValueError(f"Unknown model {model_name}") + + +def detection_to_json_types(detection): + """ + Convert a detection to a JSON-compatible dictionary. + """ + detection["conf"] = float(detection["conf"]) + detection["bbox"] = [float(x) for x in detection["bbox"]] + detection["category"] = str(detection["category"]) + return detection + + +def batch_detection(): + """ + Run detection on a batch of images. + """ + + args = parse_args() + + model = get_model(args.model) + + print("Enumerating images...", end="") + dataset = ImageDataset(args.dataset_path) + image_paths = list(tqdm(dataset.enumerate_images())) + print(" Done") + + results = {"images": []} + + for image_path in tqdm(image_paths): + if args.output_relative_filenames: + output_image_path = str(image_path.relative_to(args.dataset_path)) + else: + output_image_path = str(image_path) + + prediction = model.detect(image_path) + + results["images"].append( + { + "file": output_image_path, + "detections": [ + detection_to_json_types(detection) for detection in prediction + ], + } + ) + + results["detection_categories"] = {"1": "animal", "2": "person", "3": "vehicle"} + + results["info"] = { + "detection_completion_time": str(datetime.now()), + "format_version": "1.0", + } + + with open(args.output_path, "w", encoding="utf-8") as file_handle: + dump(results, file_handle, indent=2) + + +if __name__ == "__main__": + batch_detection()