megadetector 5.0.22__py3-none-any.whl → 5.0.24__py3-none-any.whl

megadetector/data_management/yolo_output_to_md_output.py

@@ -2,7 +2,7 @@
 
 yolo_output_to_md_output.py
 
-Converts the output of YOLOv5's detect.py or val.py to the MD API output format.
+Converts the output of YOLOv5's detect.py or val.py to the MD output format.
 
 **Converting .txt files**
 
@@ -74,7 +74,7 @@ def read_classes_from_yolo_dataset_file(fn):
         with open(fn,'r') as f:
             lines = f.readlines()
                 
-        pat = '\d+:.+'
+        pat = r'\d+:.+'
         for s in lines:
             if re.search(pat,s) is not None:
                 tokens = s.split(':')
@@ -281,7 +281,7 @@ def yolo_json_output_to_md_output(yolo_json_file,
             output_det['category'] = str(int(yolo_cat_id))
             conf = det['score']
             if truncate_to_standard_md_precision:
-                conf = ct_utils.truncate_float(conf,CONF_DIGITS)
+                conf = ct_utils.round_float(conf,CONF_DIGITS)
             output_det['conf'] = conf
             input_bbox = det['bbox']
             
@@ -301,7 +301,7 @@ def yolo_json_output_to_md_output(yolo_json_file,
                            box_width_relative,box_height_relative]
             
             if truncate_to_standard_md_precision:
-                output_bbox = ct_utils.truncate_float_array(output_bbox,COORD_DIGITS)
+                output_bbox = ct_utils.round_float_array(output_bbox,COORD_DIGITS)
                 
             output_det['bbox'] = output_bbox
             im['detections'].append(output_det)
@@ -332,7 +332,8 @@ def yolo_json_output_to_md_output(yolo_json_file,
 def yolo_txt_output_to_md_output(input_results_folder, 
                                  image_folder,
                                  output_file, 
-                                 detector_tag=None):
+                                 detector_tag=None,
+                                 truncate_to_standard_md_precision=True):
     """
     Converts a folder of YOLO-output .txt files to MD .json format.
     
@@ -347,7 +348,9 @@ def yolo_txt_output_to_md_output(input_results_folder,
         output_file (str): the MD-formatted .json file to which we should write
             results
         detector_tag (str, optional): string to put in the 'detector' field in the
-            output file            
+            output file
+        truncate_to_standard_md_precision (bool, optional): set this to truncate to 
+            COORD_DIGITS and CONF_DIGITS, like the standard MD pipeline does.
     """
     
     assert os.path.isdir(input_results_folder)
@@ -398,12 +401,16 @@ def yolo_txt_output_to_md_output(input_results_folder,
                     api_box = ct_utils.convert_yolo_to_xywh([float(row[1]), float(row[2]), 
                                                              float(row[3]), float(row[4])])
     
-                    conf = ct_utils.truncate_float(float(row[5]), precision=4)
+                    conf = float(row[5])
+                    
+                    if truncate_to_standard_md_precision:
+                        conf = ct_utils.round_float(conf, precision=CONF_DIGITS)
+                        api_box = ct_utils.round_float_array(api_box, precision=COORD_DIGITS)
                     
                     detections.append({
                         'category': str(category),
                         'conf': conf,
-                        'bbox': ct_utils.truncate_float_array(api_box, precision=4)
+                        'bbox': api_box
                     })
                 
         images_entries.append({

megadetector/detection/process_video.py

@@ -28,6 +28,7 @@ from uuid import uuid1
 from megadetector.detection import run_detector_batch
 from megadetector.visualization import visualize_detector_output
 from megadetector.utils.ct_utils import args_to_object
+from megadetector.utils.ct_utils import dict_to_kvp_list, parse_kvp_list
 from megadetector.utils.path_utils import insert_before_extension, clean_path
 from megadetector.detection.video_utils import video_to_frames
 from megadetector.detection.video_utils import run_callback_on_frames
@@ -163,7 +164,7 @@ class ProcessVideoOptions:
         self.max_width = None
         
         #: Run the model at this image size (don't mess with this unless you know what you're
-        #: getting into)
+        #: getting into)... if you just want to pass smaller frames to MD, use max_width        
         self.image_size = None
         
         #: Enable image augmentation
@@ -178,6 +179,9 @@ class ProcessVideoOptions:
         #: frame result for each video (default), or every frame that was processed?
         self.include_all_processed_frames = False
         
+        #: Detector-specific options
+        self.detector_options = None
+        
 # ...class ProcessVideoOptions
 
 
@@ -402,7 +406,7 @@ def process_video(options):
             print('Warning: frame_folder specified, but keep_extracted_frames is ' + \
                   'not; no raw frames will be written')
         
-        detector = load_detector(options.model_file)
+        detector = load_detector(options.model_file,detector_options=options.detector_options)
         
         def frame_callback(image_np,image_id):
             return detector.generate_detections_one_image(image_np,
@@ -475,7 +479,8 @@ def process_video(options):
                 class_mapping_filename=options.class_mapping_filename,
                 quiet=True,
                 augment=options.augment,
-                image_size=options.image_size)
+                image_size=options.image_size,
+                detector_options=options.detector_options)
         
             results = _add_frame_numbers_to_results(results)
         
@@ -612,7 +617,7 @@ def process_video_folder(options):
             print('Warning: frame_folder specified, but keep_extracted_frames is ' + \
                   'not; no raw frames will be written')
         
-        detector = load_detector(options.model_file)
+        detector = load_detector(options.model_file,detector_options=options.detector_options)
         
         def frame_callback(image_np,image_id):
             return detector.generate_detections_one_image(image_np,
@@ -719,7 +724,8 @@ def process_video_folder(options):
                 class_mapping_filename=options.class_mapping_filename,
                 quiet=True,
                 augment=options.augment,
-                image_size=options.image_size)
+                image_size=options.image_size,
+                detector_options=options.detector_options)
         
             _add_frame_numbers_to_results(results)
             
@@ -910,6 +916,8 @@ def options_to_command(options):
         cmd += ' --force_extracted_frame_folder_deletion'
     if options.force_rendered_frame_folder_deletion:
         cmd += ' --force_rendered_frame_folder_deletion'
+    if options.detector_options is not None and len(options.detector_options) > 0:
+        cmd += '--detector_options {}'.format(dict_to_kvp_list(options.detector_options))        
 
     return cmd
 
@@ -1209,14 +1217,23 @@ def main():
     parser.add_argument('--allow_empty_videos',
                         action='store_true',
                         help='By default, videos with no retrievable frames cause an error, this makes it a warning')
-            
+    
+    parser.add_argument(
+        '--detector_options',
+        nargs='*',
+        metavar='KEY=VALUE',
+        default='',
+        help='Detector-specific options, as a space-separated list of key-value pairs')
+        
     if len(sys.argv[1:]) == 0:
         parser.print_help()
         parser.exit()
         
     args = parser.parse_args()
-    options = ProcessVideoOptions()
+    options = ProcessVideoOptions()    
     args_to_object(args,options)
+    
+    options.detector_options = parse_kvp_list(args.detector_options)
 
     if os.path.isdir(options.input_video_file):
         process_video_folder(options)