megadetector 5.0.5__py3-none-any.whl → 5.0.7__py3-none-any.whl

detection/run_inference_with_yolov5_val.py

@@ -38,6 +38,7 @@
 #%% Imports
 
 import os
+import sys
 import uuid
 import glob
 import tempfile
@@ -48,7 +49,12 @@ from tqdm import tqdm
 
 from md_utils import path_utils
 from md_utils import process_utils
+from md_utils import string_utils
 from data_management import yolo_output_to_md_output
+from detection.run_detector import try_download_known_detector
+
+default_image_size_with_augmentation = int(1280 * 1.3)
+default_image_size_with_no_augmentation = 1280
 
 
 #%% Options class
@@ -59,17 +65,24 @@ class YoloInferenceOptions:
     
     input_folder = None
     model_filename = None
-    yolo_working_folder = None
     output_file = None
-
+    
     ## Optional ##
     
-    image_size = 1280 * 1.3
+    # Required for older YOLOv5 inference, not for newer ulytralytics inference
+    yolo_working_folder = None
+    
+    # Currently 'yolov5' and 'ultralytics' are supported, and really these are proxies for
+    # "the yolov5 repo" and "the ultralytics repo" (typically YOLOv8).
+    model_type = 'yolov5' 
+
+    image_size = default_image_size_with_augmentation
     conf_thres = '0.001'
     batch_size = 1
     device_string = '0'
     augment = True
-
+    half_precision_enabled = None
+    
     symlink_folder = None
     use_symlinks = True
     
@@ -80,24 +93,54 @@ class YoloInferenceOptions:
     
     # These are deliberately offset from the standard MD categories; YOLOv5
     # needs categories IDs to start at 0.
+    #
+    # This can also be a string that points to a YOLOv5 dataset.yaml file.
     yolo_category_id_to_name = {0:'animal',1:'person',2:'vehicle'}
     
     # 'error','skip','overwrite'
     overwrite_handling = 'skip'
+    
+    preview_yolo_command_only = False
+    
+    treat_copy_failures_as_warnings = False
+    
+    save_yolo_debug_output = False
             
     
 #%% Main function
 
 def run_inference_with_yolo_val(options):
 
-    ##%% Path handling
+    ##%% Input and path handling
+    
+    if options.model_type == 'yolov8':
+        
+        print('Warning: model type "yolov8" supplied, "ultralytics" is the preferred model type string for YOLOv8 models')
+        options.model_type = 'ultralytics'
+        
+    if (options.model_type == 'yolov5') and ('yolov8' in options.model_filename.lower()):
+        print('\n\n*** Warning: model type set as "yolov5", but your model filename contains "yolov8"... did you mean to use --model_type yolov8?" ***\n\n')        
     
+    if options.yolo_working_folder is None:
+        assert options.model_type == 'ultralytics', \
+            'A working folder is required to run YOLOv5 val.py'
+    else:
+        assert os.path.isdir(options.yolo_working_folder), \
+            'Could not find working folder {}'.format(options.yolo_working_folder)
+                
     assert os.path.isdir(options.input_folder) or os.path.isfile(options.input_folder), \
         'Could not find input {}'.format(options.input_folder)
-    assert os.path.isdir(options.yolo_working_folder), \
-        'Could not find working folder {}'.format(options.yolo_working_folder)
-    assert os.path.isfile(options.model_filename), \
-        'Could not find model file {}'.format(options.model_filename)
+    
+    if options.half_precision_enabled is not None:
+        assert options.half_precision_enabled in (0,1), \
+            'Invalid value {} for --half_precision_enabled (should be 0 or 1)'.format(
+                options.half_precision_enabled)
+            
+    # If the model filename is a known model string (e.g. "MDv5A", download the model if necessary)
+    model_filename = try_download_known_detector(options.model_filename)
+    
+    assert os.path.isfile(model_filename), \
+        'Could not find model file {}'.format(model_filename)
     
     if os.path.exists(options.output_file):
         if options.overwrite_handling == 'skip':
@@ -112,6 +155,17 @@ def run_inference_with_yolo_val(options):
             
     os.makedirs(os.path.dirname(options.output_file),exist_ok=True)
     
+    
+    ##%% Other input handling
+    
+    if isinstance(options.yolo_category_id_to_name,str):
+        assert os.path.isfile(options.yolo_category_id_to_name)
+        yolo_dataset_file = options.yolo_category_id_to_name
+        options.yolo_category_id_to_name = \
+            yolo_output_to_md_output.read_classes_from_yolo_dataset_file(yolo_dataset_file)
+        print('Loaded {} category mappings from {}'.format(
+            len(options.yolo_category_id_to_name),yolo_dataset_file))
+
     temporary_folder = None
     symlink_folder_is_temp_folder = False
     yolo_folder_is_temp_folder = False
@@ -185,24 +239,38 @@ def run_inference_with_yolo_val(options):
             else:
                 shutil.copyfile(image_fn,symlink_full_path)
         except Exception as e:
-            image_id_to_error[image_id] = str(e)
-            print('Warning: error copying/creating link for input file {}: {}'.format(
-                image_fn,str(e)))
-            continue
+            error_string = str(e)
+            image_id_to_error[image_id] = error_string
+            # Always break if the user is trying to create symlinks on Windows without
+            # permission, 100% of images will always fail in this case.
+            if ('a required privilege is not held by the client' in error_string.lower()) or \
+               (not options.treat_copy_failures_as_warnings):
+                   print('\nError copying/creating link for input file {}: {}'.format(
+                       image_fn,error_string))
+                   
+                   raise
+            else:
+                print('Warning: error copying/creating link for input file {}: {}'.format(
+                    image_fn,error_string))
+                continue
         
     # ...for each image
 
 
-    ##%% Create the dataset file
+    ##%% Create the dataset file if necessary
+    
+    # This may have been passed in as a string, but at this point, we should have
+    # loaded the dataset file.
+    assert isinstance(options.yolo_category_id_to_name,dict)
     
     # Category IDs need to be continuous integers starting at 0
     category_ids = sorted(list(options.yolo_category_id_to_name.keys()))
     assert category_ids[0] == 0
     assert len(category_ids) == 1 + category_ids[-1]
     
-    dataset_file = os.path.join(yolo_results_folder,'dataset.yaml')
+    yolo_dataset_file = os.path.join(yolo_results_folder,'dataset.yaml')
     
-    with open(dataset_file,'w') as f:
+    with open(yolo_dataset_file,'w') as f:
         f.write('path: {}\n'.format(symlink_folder_inner))
         f.write('train: .\n')
         f.write('val: .\n')
@@ -217,50 +285,154 @@ def run_inference_with_yolo_val(options):
                                         options.yolo_category_id_to_name[category_id]))
 
 
-    ##%% Prepare YOLOv5 command
+    ##%% Prepare Python command or YOLO CLI command
     
     image_size_string = str(round(options.image_size))
-    cmd = 'python val.py --data "{}"'.format(dataset_file)
-    cmd += ' --weights "{}"'.format(options.model_filename)
-    cmd += ' --batch-size {} --imgsz {} --conf-thres {} --task test'.format(
-        options.batch_size,image_size_string,options.conf_thres)
-    cmd += ' --device "{}" --save-json'.format(options.device_string)
-    cmd += ' --project "{}" --name "{}" --exist-ok'.format(yolo_results_folder,'yolo_results')
     
-    if options.augment:
-        cmd += ' --augment'
+    if options.model_type == 'yolov5':
+        
+        cmd = 'python val.py --task test --data "{}"'.format(yolo_dataset_file)
+        cmd += ' --weights "{}"'.format(model_filename)
+        cmd += ' --batch-size {} --imgsz {} --conf-thres {}'.format(
+            options.batch_size,image_size_string,options.conf_thres)
+        cmd += ' --device "{}" --save-json'.format(options.device_string)
+        cmd += ' --project "{}" --name "{}" --exist-ok'.format(yolo_results_folder,'yolo_results')
+        
+        if options.augment:
+            cmd += ' --augment'
+                
+        # --half is a store_true argument for YOLOv5's val.py
+        if (options.half_precision_enabled is not None) and (options.half_precision_enabled == 1):
+            cmd += ' --half'
+        
+        # Sometimes useful for debugging
+        # cmd += ' --save_conf --save_txt'
+        
+    elif options.model_type == 'ultralytics':
+                
+        if options.augment:
+            augment_string = 'augment'
+        else:
+            augment_string = ''
+            
+        cmd = 'yolo val {} model="{}" imgsz={} batch={} data="{}" project="{}" name="{}" device="{}"'.\
+            format(augment_string,model_filename,image_size_string,options.batch_size,
+                   yolo_dataset_file,yolo_results_folder,'yolo_results',options.device_string)
+        cmd += ' save_json exist_ok'
+        
+        if (options.half_precision_enabled is not None):
+            if options.half_precision_enabled == 1:
+                cmd += ' --half=True'
+            else:
+                assert options.half_precision_enabled == 0
+                cmd += ' --half=False'
+        
+        # Sometimes useful for debugging
+        # cmd += ' save_conf save_txt'
+            
+    else:
+        
+        raise ValueError('Unrecognized model type {}'.format(options.model_type))
+        
+    # print(cmd); import clipboard; clipboard.copy(cmd)
+
     
+    ##%% Run YOLO command
+    
+    if options.yolo_working_folder is not None:
+        current_dir = os.getcwd()
+        os.chdir(options.yolo_working_folder)
 
-    ##%% Run YOLOv5 command
+    print('Running YOLO inference command:\n{}\n'.format(cmd))
     
-    current_dir = os.getcwd()
-    os.chdir(options.yolo_working_folder)    
-    execution_result = process_utils.execute_and_print(cmd)
-    assert execution_result['status'] == 0, 'Error running YOLOv5'
+    if options.preview_yolo_command_only:
+        
+        if options.remove_symlink_folder:
+            try:
+                print('Removing YOLO symlink folder {}'.format(symlink_folder))
+                shutil.rmtree(symlink_folder)
+            except Exception:
+                print('Warning: error removing symlink folder {}'.format(symlink_folder))
+                pass
+        if options.remove_yolo_results_folder:
+            try:
+                print('Removing YOLO results folder {}'.format(yolo_results_folder))
+                shutil.rmtree(yolo_results_folder)
+            except Exception:
+                print('Warning: error removing YOLO results folder {}'.format(yolo_results_folder))
+                pass
+        
+        sys.exit()
+    
+    execution_result = process_utils.execute_and_print(cmd,encoding='utf-8',verbose=True)
+    assert execution_result['status'] == 0, 'Error running {}'.format(options.model_type)
     yolo_console_output = execution_result['output']
+      
+    if options.save_yolo_debug_output:
+        with open(os.path.join(yolo_results_folder,'yolo_console_output.txt'),'w') as f:
+            for s in yolo_console_output:
+                f.write(s + '\n')
+        with open(os.path.join(yolo_results_folder,'image_id_to_file.json'),'w') as f:
+            json.dump(image_id_to_file,f,indent=1)
+        with open(os.path.join(yolo_results_folder,'image_id_to_error.json'),'w') as f:
+            json.dump(image_id_to_error,f,indent=1)
+                
+        
+    # YOLO console output contains lots of ANSI escape codes, remove them for easier parsing
+    yolo_console_output = [string_utils.remove_ansi_codes(s) for s in yolo_console_output]
     
+    # Find errors that occrred during the initial corruption check; these will not be included in the
+    # output.  Errors that occur during inference will be handled separately.
     yolo_read_failures = []
+    
     for line in yolo_console_output:
+        # Lines look like:
+        #
+        # For ultralytics val:
+        #
+        # val: WARNING ⚠️ /a/b/c/d.jpg: ignoring corrupt image/label: [Errno 13] Permission denied: '/a/b/c/d.jpg'
+        # line = "val: WARNING ⚠️ /a/b/c/d.jpg: ignoring corrupt image/label: [Errno 13] Permission denied: '/a/b/c/d.jpg'"
+        #
+        # For yolov5 val.py:
+        #
+        # test: WARNING: a/b/c/d.jpg: ignoring corrupt image/label: cannot identify image file '/a/b/c/d.jpg'
+        # line = "test: WARNING: a/b/c/d.jpg: ignoring corrupt image/label: cannot identify image file '/a/b/c/d.jpg'"
         if 'cannot identify image file' in line:
             tokens = line.split('cannot identify image file')
             image_name = tokens[-1].strip()
             assert image_name[0] == "'" and image_name [-1] == "'"
             image_name = image_name[1:-1]
             yolo_read_failures.append(image_name)            
-            
+        elif 'ignoring corrupt image/label' in line:
+            assert 'WARNING' in line
+            if '⚠️' in line:
+                assert line.startswith('val'), \
+                    'Unrecognized line in YOLO output: {}'.format(line)
+                tokens = line.split('ignoring corrupt image/label')
+                image_name = tokens[0].split('⚠️')[-1].strip()
+            else:
+                assert line.startswith('test'), \
+                    'Unrecognized line in YOLO output: {}'.format(line)
+                tokens = line.split('ignoring corrupt image/label')
+                image_name = tokens[0].split('WARNING:')[-1].strip()
+            assert image_name.endswith(':')
+            image_name = image_name[0:-1]
+            yolo_read_failures.append(image_name)
+                    
     # image_file = yolo_read_failures[0]
     for image_file in yolo_read_failures:
         image_id = os.path.splitext(os.path.basename(image_file))[0]
         assert image_id in image_id_to_file
         if image_id not in image_id_to_error:
-            image_id_to_error[image_id] = 'YOLOv5 read failure'
+            image_id_to_error[image_id] = 'YOLO read failure'
     
-    os.chdir(current_dir)
+    if options.yolo_working_folder is not None:
+        os.chdir(current_dir)
         
     
     ##%% Convert results to MD format
     
-    json_files = glob.glob(yolo_results_folder+ '/yolo_results/*.json')
+    json_files = glob.glob(yolo_results_folder + '/yolo_results/*.json')
     assert len(json_files) == 1    
     yolo_json_file = json_files[0]
 
@@ -288,7 +460,7 @@ def run_inference_with_yolo_val(options):
         image_folder=image_base,
         output_file=options.output_file,
         yolo_category_id_to_name=options.yolo_category_id_to_name,
-        detector_name=os.path.basename(options.model_filename),
+        detector_name=os.path.basename(model_filename),
         image_id_to_relative_path=image_id_to_relative_path,
         image_id_to_error=image_id_to_error)
 
@@ -312,7 +484,7 @@ def run_inference_with_yolo_val(options):
 
 #%% Command-line driver
 
-import argparse,sys
+import argparse
 from md_utils.ct_utils import args_to_object
 
 def main():
@@ -329,14 +501,14 @@ def main():
     parser.add_argument(
         'output_file',type=str,
         help='.json file where output will be written')
-    parser.add_argument(
-        'yolo_working_folder',type=str,
-        help='folder in which to execute val.py')
     
     parser.add_argument(
-        '--image_size', default=options.image_size, type=int,
-        help='image size for model execution (default {})'.format(
-            options.image_size))
+        '--yolo_working_folder',type=str,default=None,
+        help='folder in which to execute val.py (not necessary for YOLOv8 inference)')
+    parser.add_argument(
+        '--image_size', default=None, type=int,
+        help='image size for model execution (default {} when augmentation is enabled, else {})'.format(
+            default_image_size_with_augmentation,default_image_size_with_no_augmentation))
     parser.add_argument(
         '--conf_thres', default=options.conf_thres, type=float,
         help='confidence threshold for including detections in the output file (default {})'.format(
@@ -344,13 +516,23 @@ def main():
     parser.add_argument(
         '--batch_size', default=options.batch_size, type=int,
         help='inference batch size (default {})'.format(options.batch_size))
+    parser.add_argument(
+        '--half_precision_enabled', default=None, type=int,
+        help='use half-precision-inference (1 or 0) (default is the underlying model\'s default, probably half for YOLOv8 and full for YOLOv8')
     parser.add_argument(
         '--device_string', default=options.device_string, type=str,
-        help='CUDA device specifier, e.g. "0" or "cpu" (default {})'.format(options.device_string))
+        help='CUDA device specifier, typically "0" or "1" for CUDA devices, "mps" for M1/M2 devices, or "cpu" (default {})'.format(options.device_string))
     parser.add_argument(
         '--overwrite_handling', default=options.overwrite_handling, type=str,
         help='action to take if the output file exists (skip, error, overwrite) (default {})'.format(
-            options.overwrite_handling)    )
+            options.overwrite_handling))
+    parser.add_argument(
+        '--yolo_dataset_file', default=None, type=str,
+        help='YOLOv5 dataset.yml file from which we should load category information ' + \
+            '(otherwise defaults to MD categories)')
+    parser.add_argument(
+        '--model_type', default=options.model_type, type=str,
+        help='Model type ("yolov5" or "ultralytics" ("yolov8" behaves the same as "ultralytics")) (default {})'.format(options.model_type))
 
     parser.add_argument(
         '--symlink_folder', type=str,
@@ -367,11 +549,19 @@ def main():
     parser.add_argument(
         '--no_remove_yolo_results_folder', action='store_true',
         help='don\'t remove the temporary folder full of YOLO intermediate files')
+    parser.add_argument(
+        '--save_yolo_debug_output', action='store_true',
+        help='write yolo console output to a text file in the results folder, along with additional debug files')
+    
+    parser.add_argument(
+        '--preview_yolo_command_only', action='store_true',
+        help='don\'t run inference, just preview the YOLO inference command (still creates symlinks)')
     
     if options.augment:
         default_augment_enabled = 1
     else:
         default_augment_enabled = 0
+            
     parser.add_argument(
         '--augment_enabled', default=default_augment_enabled, type=int,
         help='enable/disable augmentation (default {})'.format(default_augment_enabled))
@@ -381,8 +571,27 @@ def main():
         parser.exit()
 
     args = parser.parse_args()
-    
+        
+    # If the caller hasn't specified an image size, choose one based on whether augmentation
+    # is enabled.
+    if args.image_size is None:        
+        assert args.augment_enabled in (0,1), \
+            'Illegal augment_enabled value {}'.format(args.augment_enabled)
+        if args.augment_enabled == 1:
+            args.image_size = default_image_size_with_augmentation
+        else:
+            args.image_size = default_image_size_with_no_augmentation
+        augment_enabled_string = 'enabled'
+        if not args.augment_enabled:
+            augment_enabled_string = 'disabled'
+        print('Augmentation is {}, using default image size {}'.format(
+            augment_enabled_string,args.image_size))
+        
     args_to_object(args, options)
+    
+    if args.yolo_dataset_file is not None:
+        options.yolo_category_id_to_name = args.yolo_dataset_file
+        
     options.remove_symlink_folder = (not options.no_remove_symlink_folder)
     options.remove_yolo_results_folder = (not options.no_remove_yolo_results_folder)
     options.use_symlinks = (not options.no_use_symlinks)
@@ -400,16 +609,6 @@ if __name__ == '__main__':
 #%% Scrap
 
 if False:
-    
-    #%% Run from a set of options
-    
-    options = YoloInferenceOptions()
-
-    args = {'augment': 1, 'batch_size': 1, 'conf_thres': 0.005, 'device_string': '1', 'image_size': 1664.0, 'input_folder': '/home/user/postprocessing/usgs-kissel/usgs-kissel-2023-09-11-aug-v5a.0.0/chunk031.json', 'model_filename': '/home/user/models/camera_traps/megadetector/md_v5.0.0/md_v5a.0.0.pt', 'output_file': '/home/user/postprocessing/usgs-kissel/usgs-kissel-2023-09-11-aug-v5a.0.0/chunk031_results.json', 'overwrite_handling': 'skip', 'symlink_folder': '/home/user/postprocessing/usgs-kissel/usgs-kissel-2023-09-11-aug-v5a.0.0/symlinks/symlinks_031', 'yolo_results_folder': '/home/user/postprocessing/usgs-kissel/usgs-kissel-2023-09-11-aug-v5a.0.0/yolo_results/yolo_results_031', 'yolo_working_folder': '/home/user/git/yolov5', 'remove_symlink_folder': True, 'remove_yolo_results_folder': True, 'use_symlinks': False, 'augment': True}
-    
-    for k in args:
-        setattr(options, k, args[k])
-
 
     #%% Test driver (folder)
     
@@ -452,8 +651,10 @@ if False:
     options.remove_temporary_symlink_folder = False
     options.remove_yolo_results_file = False
     
-    cmd = f'python run_inference_with_yolov5_val.py {model_filename} {input_folder} {output_file} {yolo_working_folder} ' + \
-          f' --image_size {options.image_size} --conf_thres {options.conf_thres} --batch_size {options.batch_size} ' + \
+    cmd = f'python run_inference_with_yolov5_val.py {model_filename} {input_folder} ' + \
+          f'{output_file} --yolo_working_folder {yolo_working_folder} ' + \
+          f' --image_size {options.image_size} --conf_thres {options.conf_thres} ' + \
+          f' --batch_size {options.batch_size} ' + \
           f' --symlink_folder {options.symlink_folder} --yolo_results_folder {options.yolo_results_folder} ' + \
           ' --no_remove_symlink_folder --no_remove_yolo_results_folder'
       
@@ -470,6 +671,74 @@ if False:
         import clipboard; clipboard.copy(cmd)
     
     
+    #%% Test driver (folder) (YOLOv8 model)
+    
+    project_name = 'yolov8-inference-test'
+    input_folder = os.path.expanduser('~/data/usgs-kissel-training-resized/val')
+    dataset_file = os.path.expanduser('~/data/usgs-kissel-training-yolo/dataset.yaml')
+    output_folder = os.path.expanduser(f'~/tmp/{project_name}')
+    model_filename = os.path.expanduser(
+        '~/models/usgs-tegus/usgs-tegus-yolov8x-2023.10.25-b-1-img640-e200-best.pt')
+    model_name = os.path.splitext(os.path.basename(model_filename))[0]
+    
+    assert os.path.isdir(input_folder)
+    assert os.path.isfile(dataset_file)
+    assert os.path.isfile(model_filename)    
+    
+    symlink_folder = os.path.join(output_folder,'symlinks')
+    yolo_results_folder = os.path.join(output_folder,'yolo_results')
+    
+    output_file = os.path.join(output_folder,'{}_{}-md_format.json'.format(
+        project_name,model_name))
+    
+    options = YoloInferenceOptions()
+    
+    options.model_type = 'yolov8'    
+    options.yolo_category_id_to_name = dataset_file    
+    options.yolo_working_folder = None    
+    options.output_file = output_file
+    
+    options.augment = False
+    options.conf_thres = '0.001'
+    options.batch_size = 1
+    options.device_string = '0'
+
+    if options.augment:
+        options.image_size = round(640 * 1.3)
+    else:
+        options.image_size = 640
+    
+    options.input_folder = input_folder
+    options.model_filename = model_filename
+    
+    options.yolo_results_folder = yolo_results_folder 
+    options.symlink_folder = symlink_folder
+    options.use_symlinks = False
+    
+    options.remove_temporary_symlink_folder = False
+    options.remove_yolo_results_file = False
+    
+    cmd = f'python run_inference_with_yolov5_val.py {model_filename} ' + \
+          f'{input_folder} {output_file}' + \
+          f' --image_size {options.image_size} --conf_thres {options.conf_thres} ' + \
+          f' --batch_size {options.batch_size} --symlink_folder {options.symlink_folder} ' + \
+          f'--yolo_results_folder {options.yolo_results_folder} --model_type {options.model_type}' + \
+          f' --yolo_dataset_file {options.yolo_category_id_to_name}' + \
+          ' --no_remove_symlink_folder --no_remove_yolo_results_folder'
+      
+    if not options.use_symlinks:
+        cmd += ' --no_use_symlinks'
+    if not options.augment:
+        cmd += ' --augment_enabled 0'
+        
+    print(cmd)
+    execute_in_python = False
+    if execute_in_python:
+        run_inference_with_yolo_val(options)
+    else:
+        import clipboard; clipboard.copy(cmd)
+    
+    
     #%% Preview results
     
     postprocessing_output_folder = os.path.join(output_folder,'yolo-val-preview')