PyPI - megadetector - Versions diffs - 10.0.13__py3-none-any.whl - Mend

megadetector 10.0.13__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.

This version of megadetector might be problematic. Click here for more details.

Files changed (147) hide show

megadetector/__init__.py +0 -0
megadetector/api/__init__.py +0 -0
megadetector/api/batch_processing/integration/digiKam/setup.py +6 -0
megadetector/api/batch_processing/integration/digiKam/xmp_integration.py +465 -0
megadetector/api/batch_processing/integration/eMammal/test_scripts/config_template.py +5 -0
megadetector/api/batch_processing/integration/eMammal/test_scripts/push_annotations_to_emammal.py +125 -0
megadetector/api/batch_processing/integration/eMammal/test_scripts/select_images_for_testing.py +55 -0
megadetector/classification/__init__.py +0 -0
megadetector/classification/aggregate_classifier_probs.py +108 -0
megadetector/classification/analyze_failed_images.py +227 -0
megadetector/classification/cache_batchapi_outputs.py +198 -0
megadetector/classification/create_classification_dataset.py +626 -0
megadetector/classification/crop_detections.py +516 -0
megadetector/classification/csv_to_json.py +226 -0
megadetector/classification/detect_and_crop.py +853 -0
megadetector/classification/efficientnet/__init__.py +9 -0
megadetector/classification/efficientnet/model.py +415 -0
megadetector/classification/efficientnet/utils.py +608 -0
megadetector/classification/evaluate_model.py +520 -0
megadetector/classification/identify_mislabeled_candidates.py +152 -0
megadetector/classification/json_to_azcopy_list.py +63 -0
megadetector/classification/json_validator.py +696 -0
megadetector/classification/map_classification_categories.py +276 -0
megadetector/classification/merge_classification_detection_output.py +509 -0
megadetector/classification/prepare_classification_script.py +194 -0
megadetector/classification/prepare_classification_script_mc.py +228 -0
megadetector/classification/run_classifier.py +287 -0
megadetector/classification/save_mislabeled.py +110 -0
megadetector/classification/train_classifier.py +827 -0
megadetector/classification/train_classifier_tf.py +725 -0
megadetector/classification/train_utils.py +323 -0
megadetector/data_management/__init__.py +0 -0
megadetector/data_management/animl_to_md.py +161 -0
megadetector/data_management/annotations/__init__.py +0 -0
megadetector/data_management/annotations/annotation_constants.py +33 -0
megadetector/data_management/camtrap_dp_to_coco.py +270 -0
megadetector/data_management/cct_json_utils.py +566 -0
megadetector/data_management/cct_to_md.py +184 -0
megadetector/data_management/cct_to_wi.py +293 -0
megadetector/data_management/coco_to_labelme.py +284 -0
megadetector/data_management/coco_to_yolo.py +702 -0
megadetector/data_management/databases/__init__.py +0 -0
megadetector/data_management/databases/add_width_and_height_to_db.py +107 -0
megadetector/data_management/databases/combine_coco_camera_traps_files.py +210 -0
megadetector/data_management/databases/integrity_check_json_db.py +528 -0
megadetector/data_management/databases/subset_json_db.py +195 -0
megadetector/data_management/generate_crops_from_cct.py +200 -0
megadetector/data_management/get_image_sizes.py +164 -0
megadetector/data_management/labelme_to_coco.py +559 -0
megadetector/data_management/labelme_to_yolo.py +349 -0
megadetector/data_management/lila/__init__.py +0 -0
megadetector/data_management/lila/create_lila_blank_set.py +556 -0
megadetector/data_management/lila/create_lila_test_set.py +187 -0
megadetector/data_management/lila/create_links_to_md_results_files.py +106 -0
megadetector/data_management/lila/download_lila_subset.py +182 -0
megadetector/data_management/lila/generate_lila_per_image_labels.py +777 -0
megadetector/data_management/lila/get_lila_annotation_counts.py +174 -0
megadetector/data_management/lila/get_lila_image_counts.py +112 -0
megadetector/data_management/lila/lila_common.py +319 -0
megadetector/data_management/lila/test_lila_metadata_urls.py +164 -0
megadetector/data_management/mewc_to_md.py +344 -0
megadetector/data_management/ocr_tools.py +873 -0
megadetector/data_management/read_exif.py +964 -0
megadetector/data_management/remap_coco_categories.py +195 -0
megadetector/data_management/remove_exif.py +156 -0
megadetector/data_management/rename_images.py +194 -0
megadetector/data_management/resize_coco_dataset.py +663 -0
megadetector/data_management/speciesnet_to_md.py +41 -0
megadetector/data_management/wi_download_csv_to_coco.py +247 -0
megadetector/data_management/yolo_output_to_md_output.py +594 -0
megadetector/data_management/yolo_to_coco.py +876 -0
megadetector/data_management/zamba_to_md.py +188 -0
megadetector/detection/__init__.py +0 -0
megadetector/detection/change_detection.py +840 -0
megadetector/detection/process_video.py +479 -0
megadetector/detection/pytorch_detector.py +1451 -0
megadetector/detection/run_detector.py +1267 -0
megadetector/detection/run_detector_batch.py +2159 -0
megadetector/detection/run_inference_with_yolov5_val.py +1314 -0
megadetector/detection/run_md_and_speciesnet.py +1494 -0
megadetector/detection/run_tiled_inference.py +1038 -0
megadetector/detection/tf_detector.py +209 -0
megadetector/detection/video_utils.py +1379 -0
megadetector/postprocessing/__init__.py +0 -0
megadetector/postprocessing/add_max_conf.py +72 -0
megadetector/postprocessing/categorize_detections_by_size.py +166 -0
megadetector/postprocessing/classification_postprocessing.py +1752 -0
megadetector/postprocessing/combine_batch_outputs.py +249 -0
megadetector/postprocessing/compare_batch_results.py +2110 -0
megadetector/postprocessing/convert_output_format.py +403 -0
megadetector/postprocessing/create_crop_folder.py +629 -0
megadetector/postprocessing/detector_calibration.py +570 -0
megadetector/postprocessing/generate_csv_report.py +522 -0
megadetector/postprocessing/load_api_results.py +223 -0
megadetector/postprocessing/md_to_coco.py +428 -0
megadetector/postprocessing/md_to_labelme.py +351 -0
megadetector/postprocessing/md_to_wi.py +41 -0
megadetector/postprocessing/merge_detections.py +392 -0
megadetector/postprocessing/postprocess_batch_results.py +2077 -0
megadetector/postprocessing/remap_detection_categories.py +226 -0
megadetector/postprocessing/render_detection_confusion_matrix.py +677 -0
megadetector/postprocessing/repeat_detection_elimination/find_repeat_detections.py +206 -0
megadetector/postprocessing/repeat_detection_elimination/remove_repeat_detections.py +82 -0
megadetector/postprocessing/repeat_detection_elimination/repeat_detections_core.py +1665 -0
megadetector/postprocessing/separate_detections_into_folders.py +795 -0
megadetector/postprocessing/subset_json_detector_output.py +964 -0
megadetector/postprocessing/top_folders_to_bottom.py +238 -0
megadetector/postprocessing/validate_batch_results.py +332 -0
megadetector/taxonomy_mapping/__init__.py +0 -0
megadetector/taxonomy_mapping/map_lila_taxonomy_to_wi_taxonomy.py +491 -0
megadetector/taxonomy_mapping/map_new_lila_datasets.py +213 -0
megadetector/taxonomy_mapping/prepare_lila_taxonomy_release.py +165 -0
megadetector/taxonomy_mapping/preview_lila_taxonomy.py +543 -0
megadetector/taxonomy_mapping/retrieve_sample_image.py +71 -0
megadetector/taxonomy_mapping/simple_image_download.py +224 -0
megadetector/taxonomy_mapping/species_lookup.py +1008 -0
megadetector/taxonomy_mapping/taxonomy_csv_checker.py +159 -0
megadetector/taxonomy_mapping/taxonomy_graph.py +346 -0
megadetector/taxonomy_mapping/validate_lila_category_mappings.py +83 -0
megadetector/tests/__init__.py +0 -0
megadetector/tests/test_nms_synthetic.py +335 -0
megadetector/utils/__init__.py +0 -0
megadetector/utils/ct_utils.py +1857 -0
megadetector/utils/directory_listing.py +199 -0
megadetector/utils/extract_frames_from_video.py +307 -0
megadetector/utils/gpu_test.py +125 -0
megadetector/utils/md_tests.py +2072 -0
megadetector/utils/path_utils.py +2832 -0
megadetector/utils/process_utils.py +172 -0
megadetector/utils/split_locations_into_train_val.py +237 -0
megadetector/utils/string_utils.py +234 -0
megadetector/utils/url_utils.py +825 -0
megadetector/utils/wi_platform_utils.py +968 -0
megadetector/utils/wi_taxonomy_utils.py +1759 -0
megadetector/utils/write_html_image_list.py +239 -0
megadetector/visualization/__init__.py +0 -0
megadetector/visualization/plot_utils.py +309 -0
megadetector/visualization/render_images_with_thumbnails.py +243 -0
megadetector/visualization/visualization_utils.py +1940 -0
megadetector/visualization/visualize_db.py +630 -0
megadetector/visualization/visualize_detector_output.py +479 -0
megadetector/visualization/visualize_video_output.py +705 -0
megadetector-10.0.13.dist-info/METADATA +134 -0
megadetector-10.0.13.dist-info/RECORD +147 -0
megadetector-10.0.13.dist-info/WHEEL +5 -0
megadetector-10.0.13.dist-info/licenses/LICENSE +19 -0
megadetector-10.0.13.dist-info/top_level.txt +1 -0

megadetector/utils/directory_listing.py ADDED Viewed

@@ -0,0 +1,199 @@
+"""
+directory_listing.py
+Script for creating Apache-style HTML directory listings for a local directory
+and all its subdirectories.
+Also includes a preview of a jpg file (the first in an alphabetical list),
+if present.
+"""
+#%% Imports
+import os
+import sys
+import argparse
+from megadetector.utils.path_utils import is_image_file
+#%% Directory enumeration functions
+def _create_plain_index(root, dirs, files, dirname=None):
+    """
+    Creates the fairly plain HTML folder index including a preview of a single image file,
+    if any is present.
+    Args:
+        root (str): path to the root directory, all paths in [dirs] and
+            [files] are relative to this root folder
+        dirs (list): list of strings, the directories in [root]
+        files (list): list of strings, the files in [root]
+        dirname (str, optional): name to print in the html,
+            which may be different than [root]
+    Returns:
+        str: HTML source of the directory listing
+    """
+    if dirname is None:
+        dirname = root or '/'
+    dirname = dirname.replace('\\','/')
+    html = "<!DOCTYPE html>\n"
+    html += "<html lang='en'><head>"
+    html += "<title>Index of {}</title>\n".format(dirname)
+    html += "<meta charset='UTF-8'>\n"
+    html += "<style>\n"
+    html += "body { font-family: Segoe UI, Helvetica, Arial, sans-serif; }\n"
+    html += "a { text-decoration:none; }\n"
+    html += "</style>\n"
+    html += "</head><body>\n"
+    html += "<h1>Index of {}</h1>\n".format(dirname)
+    # Insert preview image
+    jpg_files = [f for f in files if is_image_file(f)]
+    if len(jpg_files) > 0:
+        # This is slow, so defaulting to just the first image:
+        #
+        # Use the largest image file as this is most likely to contain
+        # useful content.
+        #
+        # jpg_file_sizes = [os.path.getsize(f) for f in jpg_files]
+        # largest_file_index = max(range(len(jpg_files)), key=lambda x: jpg_file_sizes[x])
+        html += "<a href='{0}'><img style='height:200px; float:right;' src='{0}' alt='Preview image'></a>\n".\
+            format(jpg_files[0])
+    else:
+        html += "\n"
+        # html += "<p style='width:15em; float:right; margin:0;'>[No preview available]</p>\n"
+    if root:
+        html += "<p><a href='../index.html'>To parent directory</a></p>\n"
+    else:
+        html += "\n"
+        # html += "<p>This is the root directory.</p>\n"
+    html += "<h2>Folders</h2>\n"
+    if len(dirs) > 0:
+        html += "<ul style='list-style-type: none; padding-left:1em;'>\n"
+        for dir in sorted(dirs):
+            html += "<li>&#128193; <a href='{0}/index.html'>{0}</a></li>\n".format(dir)
+        html += "</ul>\n"
+    else:
+        html += "<p style='padding-left:1em;'>No folders</p>\n"
+    html += "<h2>Files</h2>\n"
+    if len(files) > 0:
+        html += "<ul style='list-style-type: none; padding-left:1em;'>\n"
+        for fname in sorted(files):
+            if is_image_file(fname):
+                html += "<li>&#128443; <a href='{0}'>{0}</a></li>\n".format(fname)
+            else:
+                html += "<li>&#128442; <a href='{0}'>{0}</a></li>\n".format(fname)
+        html += "</ul>\n"
+    else:
+        html += "<p style = 'padding-left:1em;'>No files</p>\n"
+    # Add some space at the bottom because the browser's status bar might hide stuff
+    html += "<p style='margin:2em;'>&nbsp;</p>\n"
+    html += "</body></html>\n"
+    return html
+# ...def _create_plain_index(...)
+def create_html_index(dir,
+                      overwrite=False,
+                      template_fun=_create_plain_index,
+                      basepath=None,
+                      recursive=True):
+    """
+    Recursively traverses the local directory [dir] and generates a index
+    file for each folder using [template_fun] to generate the HTML output.
+    Excludes hidden files.
+    Args:
+        dir (str): directory to process
+        overwrite (bool, optional): whether to over-write existing index file
+        template_fun (func, optional): function taking three arguments (string,
+            list of string, list of string) representing the current root, the list of folders,
+            and the list of files.  Should return the HTML source of the index file.
+        basepath (str, optional): if not None, the name used for each subfolder in [dir]
+            in the output files will be relative to [basepath]
+        recursive (bool, optional): recurse into subfolders
+    """
+    if template_fun is None:
+        template_fun = _create_plain_index
+    print('Traversing {}'.format(dir))
+    # Make sure we remove the trailing /
+    dir = os.path.normpath(dir)
+    # Traverse directory and all sub directories, excluding hidden files
+    for root, dirs, files in os.walk(dir):
+        # Exclude files and folders that are hidden
+        files = [f for f in files if not f[0] == '.']
+        dirs[:] = [d for d in dirs if not d[0] == '.']
+        # Output is written to file *root*/index.html
+        output_file = os.path.join(root, "index.html")
+        if (not overwrite) and os.path.isfile(output_file):
+            print('Skipping {}, file exists'.format(output_file))
+            continue
+        print("Generating {}".format(output_file))
+        # Generate HTML with template function
+        dirname = None
+        if basepath is not None:
+            dirname = os.path.relpath(root,basepath)
+        html = template_fun(root[len(dir):], dirs, files, dirname)
+        # Write to file
+        with open(output_file, 'wt') as fi:
+            fi.write(html)
+        if not recursive:
+            break
+# ...def create_html_index(...)
+#%% Command-line driver
+def main(): # noqa
+    parser = argparse.ArgumentParser()
+    parser.add_argument("directory", type=str,
+                        help='Path to directory which should be traversed.')
+    parser.add_argument("--basepath", type=str,
+                        help='Folder names will be printed relative to basepath, if specified',
+                        default=None)
+    parser.add_argument("--overwrite", action='store_true', default=False,
+                        help='If set, the script will overwrite existing index.html files.')
+    if len(sys.argv[1:]) == 0:
+        parser.print_help()
+        parser.exit()
+    args = parser.parse_args()
+    assert os.path.isdir(args.directory), "{} is not a valid directory".format(args.directory)
+    create_html_index(args.directory,
+                      overwrite=args.overwrite,
+                      basepath=args.basepath)
+if __name__ == '__main__':
+    main()

megadetector/utils/extract_frames_from_video.py ADDED Viewed

@@ -0,0 +1,307 @@
+"""
+extract_frames_from_video.py
+Extracts frames from a source video or folder of videos and writes those frames to jpeg files.
+For single videos, writes frame images to the destination folder.  For folders of videos, creates
+subfolders in the destination folder (one per video) and writes frame images to those subfolders.
+"""
+#%% Constants and imports
+import argparse
+import inspect
+import json
+import os
+import sys
+from megadetector.detection.video_utils import \
+    video_to_frames, video_folder_to_frames, is_video_file
+#%% Options class
+class FrameExtractionOptions:
+    """
+    Parameters controlling the behavior of extract_frames().
+    """
+    def __init__(self):
+        #: Number of workers to use for parallel processing
+        self.n_workers = 1
+        #: Use threads for parallel processing
+        self.parallelize_with_threads = False
+        #: JPEG quality for extracted frames
+        self.quality = 80
+        #: Maximum width for extracted frames (defaults to None)
+        self.max_width = None
+        #: Enable additional debug output
+        self.verbose = False
+        #: Sample every Nth frame starting from the first frame; if this is None
+        #: or 1, every frame is extracted.  If this is a negative value, it's interpreted
+        #: as a sampling rate in seconds, which is rounded to the nearest frame sampling
+        #: rate. Mutually exclusive with detector_output_file.
+        self.frame_sample = None
+        #: Path to MegaDetector .json output file. When specified, extracts frames
+        #: referenced in this file.  Mutually exclusive with frame_sample. [source]
+        #: must be a folder when this is specified.
+        self.detector_output_file = None
+    # ...def __init__(...)
+# ...class FrameExtractionOptions
+#%% Core functions
+def extract_frames(source, destination, options=None):
+    """
+    Extracts frames from a video or folder of videos.
+    Args:
+        source (str): path to a single video file or folder of videos
+        destination (str): folder to write frame images to (will be created if it doesn't exist)
+        options (FrameExtractionOptions, optional): parameters controlling frame extraction
+    Returns:
+        tuple: for single videos, returns (list of frame filenames, frame rate).
+               for folders, returns (list of lists of frame filenames, list of frame rates, list
+               of video filenames)
+    """
+    if options is None:
+        options = FrameExtractionOptions()
+    # Validate inputs
+    if not os.path.exists(source):
+        raise ValueError('Source path {} does not exist'.format(source))
+    if os.path.abspath(source) == os.path.abspath(destination):
+        raise ValueError('Source and destination cannot be the same')
+    # Create destination folder if it doesn't exist
+    os.makedirs(destination, exist_ok=True)
+    # Determine whether source is a file or folder
+    source_is_file = os.path.isfile(source)
+    if source_is_file:
+        # Validate that source is a video file
+        if not is_video_file(source):
+            raise ValueError('Source file {} is not a video file'.format(source))
+        # detector_output_file requires source to be a folder
+        if options.detector_output_file is not None:
+            raise ValueError('detector_output_file option requires source to be a folder, not a file')
+        # Extract frames from single video
+        return video_to_frames(input_video_file=source,
+                               output_folder=destination,
+                               overwrite=True,
+                               every_n_frames=options.frame_sample,
+                               verbose=options.verbose,
+                               quality=options.quality,
+                               max_width=options.max_width,
+                               allow_empty_videos=True)
+    else:
+        frames_to_extract = None
+        relative_paths_to_process = None
+        # Handle detector output file
+        if options.detector_output_file is not None:
+            frames_to_extract, relative_paths_to_process = _parse_detector_output(
+                options.detector_output_file, source, options.verbose)
+            options.frame_sample = None
+        return video_folder_to_frames(input_folder=source,
+                                      output_folder_base=destination,
+                                      recursive=True,
+                                      overwrite=True,
+                                      n_threads=options.n_workers,
+                                      every_n_frames=options.frame_sample,
+                                      verbose=options.verbose,
+                                      parallelization_uses_threads=options.parallelize_with_threads,
+                                      quality=options.quality,
+                                      max_width=options.max_width,
+                                      frames_to_extract=frames_to_extract,
+                                      relative_paths_to_process=relative_paths_to_process,
+                                      allow_empty_videos=True)
+# ...def extract_frames(...)
+def _parse_detector_output(detector_output_file, source_folder, verbose=False):
+    """
+    Parses a MegaDetector .json output file and returns frame extraction information.
+    Args:
+        detector_output_file (str): path to MegaDetector .json output file
+        source_folder (str): folder containing the source videos
+        verbose (bool, optional): enable additional debug output
+    Returns:
+        tuple: (frames_to_extract_dict, relative_paths_to_process) where:
+               - frames_to_extract_dict maps relative video paths to lists of frame numbers
+               - relative_paths_to_process is a list of relative video paths to process
+    """
+    print('Parsing detector output file: {}'.format(detector_output_file))
+    # Load the detector results
+    with open(detector_output_file, 'r') as f:
+        detector_results = json.load(f)
+    if 'images' not in detector_results:
+        raise ValueError('Detector output file does not contain "images" field')
+    images = detector_results['images']
+    frames_to_extract_dict = {}
+    video_files_in_results = set()
+    for image_entry in images:
+        file_path = image_entry['file']
+        # Skip non-video files
+        if not is_video_file(file_path):
+            if verbose:
+                print('Skipping non-video file {}'.format(file_path))
+            continue
+        # Check whether video file exists in source folder
+        full_video_path = os.path.join(source_folder, file_path)
+        if not os.path.isfile(full_video_path):
+            print('Warning: video file {} not found in source folder, skipping'.format(file_path))
+            continue
+        video_files_in_results.add(file_path)
+        # Determine which frames to extract for this video
+        frames_for_this_video = []
+        if 'frames_processed' in image_entry:
+            # Use the frames_processed field if available
+            frames_for_this_video = image_entry['frames_processed']
+            if verbose:
+                print('Video {}: using frames_processed field with {} frames'.format(
+                    file_path, len(frames_for_this_video)))
+        else:
+            # Extract frames from detections
+            if ('detections' in image_entry) and (image_entry['detections'] is not None):
+                frame_numbers = set()
+                for detection in image_entry['detections']:
+                    if 'frame_number' in detection:
+                        frame_numbers.add(detection['frame_number'])
+                frames_for_this_video = sorted(list(frame_numbers))
+                if verbose:
+                    print('Video {}: extracted {} unique frame numbers from detections'.format(
+                        file_path, len(frames_for_this_video)))
+        if len(frames_for_this_video) > 0:
+            frames_to_extract_dict[file_path] = frames_for_this_video
+    # ...for each image/video in this file
+    relative_paths_to_process = sorted(list(video_files_in_results))
+    print('Found {} videos with frames to extract'.format(len(frames_to_extract_dict)))
+    return frames_to_extract_dict, relative_paths_to_process
+# ...def _parse_detector_output(...)
+#%% Command-line driver
+def _args_to_object(args, obj):
+    """
+    Copy all fields from a Namespace (i.e., the output from parse_args) to an object.
+    Skips fields starting with _.  Does not check existence in the target object.
+    """
+    for n, v in inspect.getmembers(args):
+        if not n.startswith('_'):
+            setattr(obj, n, v)
+def main():
+    """
+    Command-line driver for extract_frames_from_video
+    """
+    parser = argparse.ArgumentParser(
+        description='Extract frames from videos and save as JPEG files')
+    parser.add_argument('source', type=str,
+                        help='Path to a single video file or folder containing videos')
+    parser.add_argument('destination', type=str,
+                        help='Output folder for extracted frames (will be created if it does not exist)')
+    parser.add_argument('--n_workers', type=int, default=1,
+                        help='Number of workers to use for parallel processing (default: %(default)s)')
+    parser.add_argument('--parallelize_with_threads', action='store_true',
+                        help='Use threads for parallel processing (default: use processes)')
+    parser.add_argument('--quality', type=int, default=80,
+                        help='JPEG quality for extracted frames (default: %(default)s)')
+    parser.add_argument('--max_width', type=int, default=None,
+                        help='Maximum width for extracted frames (default: no resizing)')
+    parser.add_argument('--verbose', action='store_true',
+                        help='Enable additional debug output')
+    # Mutually exclusive group for frame sampling options
+    frame_group = parser.add_mutually_exclusive_group()
+    frame_group.add_argument('--frame_sample', type=float, default=None,
+                             help='Sample every Nth frame starting from the first frame; if this is None or 1, ' +
+                                  'every frame is extracted. If this is a negative value, it\'s interpreted as a ' +
+                                  'sampling rate in seconds, which is rounded to the nearest frame sampling rate')
+    frame_group.add_argument('--detector_output_file', type=str, default=None,
+                             help='Path to MegaDetector .json output file. When specified, extracts frames ' +
+                                  'referenced in this file. Source must be a folder when this is specified.')
+    if len(sys.argv[1:]) == 0:
+        parser.print_help()
+        parser.exit()
+    args = parser.parse_args()
+    # Convert to an options object
+    options = FrameExtractionOptions()
+    _args_to_object(args, options)
+    # Additional validation
+    if options.detector_output_file is not None:
+        if not os.path.isfile(options.detector_output_file):
+            print('Error: detector_output_file {} does not exist'.format(options.detector_output_file))
+            sys.exit(1)
+    try:
+        result = extract_frames(args.source, args.destination, options)
+        if os.path.isfile(args.source):
+            frame_filenames, frame_rate = result
+            print('Extracted {} frames from {} (frame rate: {:.2f} fps)'.format(
+                len(frame_filenames), args.source, frame_rate))
+        else:
+            frame_filenames_by_video, fs_by_video, video_filenames = result
+            total_frames = sum(len(frames) for frames in frame_filenames_by_video)
+            print('Processed {} videos, extracted {} total frames'.format(
+                len(video_filenames), total_frames))
+    except Exception as e:
+        print('Error: {}'.format(str(e)))
+        sys.exit(1)
+if __name__ == '__main__':
+    main()

megadetector/utils/gpu_test.py ADDED Viewed

@@ -0,0 +1,125 @@
+"""
+gpu_test.py
+Simple script to verify CUDA availability, used to verify a CUDA environment
+for TF or PyTorch
+"""
+# Minimize TF printouts
+import os
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
+try:
+    import logging
+    logging.getLogger('tensorflow').setLevel(logging.ERROR)
+except Exception:
+    pass
+#%% Torch/TF test functions
+def torch_test():
+    """
+    Print diagnostic information about Torch/CUDA status, including Torch/CUDA versions
+    and all available CUDA device names.
+    Returns:
+        int: The number of CUDA devices reported by PyTorch.
+    """
+    try:
+        import torch
+    except Exception as e: #noqa
+        print('PyTorch unavailable, not running PyTorch tests.  PyTorch import error was:\n{}'.format(
+            str(e)))
+        return 0
+    print('Torch version: {}'.format(str(torch.__version__)))
+    print('CUDA available (according to PyTorch): {}'.format(torch.cuda.is_available()))
+    if torch.cuda.is_available():
+        print('CUDA version (according to PyTorch): {}'.format(torch.version.cuda))
+        print('CuDNN version (according to PyTorch): {}'.format(torch.backends.cudnn.version()))
+    device_ids = list(range(torch.cuda.device_count()))
+    if len(device_ids) > 0:
+        cuda_str = 'Found {} CUDA devices:'.format(len(device_ids))
+        print(cuda_str)
+        for device_id in device_ids:
+            device_name = 'unknown'
+            try:
+                device_name = torch.cuda.get_device_name(device=device_id)
+            except Exception as e: #noqa
+                pass
+            print('{}: {}'.format(device_id,device_name))
+    else:
+        print('No GPUs reported by PyTorch')
+    try:
+        if torch.backends.mps.is_built and torch.backends.mps.is_available():
+            print('PyTorch reports that Metal Performance Shaders are available')
+    except Exception:
+        pass
+    return len(device_ids)
+def tf_test():
+    """
+    Print diagnostic information about TF/CUDA status.
+    Returns:
+        int: The number of CUDA devices reported by TensorFlow.
+    """
+    try:
+        import tensorflow as tf # type: ignore
+    except Exception as e: #noqa
+        print('TensorFlow unavailable, not running TF tests.  TF import error was:\n{}'.format(
+            str(e)))
+        return 0
+    from tensorflow.python.platform import build_info as build # type: ignore
+    print(f"TF version: {tf.__version__}")
+    if 'cuda_version' not in build.build_info:
+        print('TF does not appear to be built with CUDA')
+    else:
+        print(f"CUDA build version reported by TensorFlow: {build.build_info['cuda_version']}")
+    if 'cudnn_version' not in build.build_info:
+        print('TF does not appear to be built with CuDNN')
+    else:
+        print(f"CuDNN build version reported by TensorFlow: {build.build_info['cudnn_version']}")
+    try:
+        from tensorflow.python.compiler.tensorrt import trt_convert as trt # type: ignore
+        print("Linked TensorRT version: {}".format(trt.trt_utils._pywrap_py_utils.get_linked_tensorrt_version()))
+    except Exception:
+        print('Could not probe TensorRT version')
+    gpus = tf.config.list_physical_devices('GPU')
+    if gpus is None:
+        gpus = []
+    if len(gpus) > 0:
+        print('TensorFlow found the following GPUs:')
+        for gpu in gpus:
+            print(gpu.name)
+    else:
+        print('No GPUs reported by TensorFlow')
+    return len(gpus)
+#%% Command-line driver
+if __name__ == '__main__':
+    print('*** Running Torch tests ***\n')
+    torch_test()
+    print('\n*** Running TF tests ***\n')
+    tf_test()