megadetector 5.0.28__py3-none-any.whl → 10.0.0__py3-none-any.whl

megadetector/data_management/yolo_output_to_md_output.py

@@ -14,7 +14,7 @@ no detections are present.
 
 YOLOv5 output has one text file per image, like so:
 
-0 0.0141693 0.469758 0.0283385 0.131552 0.761428 
+0 0.0141693 0.469758 0.0283385 0.131552 0.761428
 
 That's [class, x_center, y_center, width_of_box, height_of_box, confidence]
 
@@ -40,6 +40,8 @@ import json
 import csv
 import os
 import re
+import sys
+import argparse
 
 from collections import defaultdict
 from tqdm import tqdm
@@ -57,39 +59,39 @@ def read_classes_from_yolo_dataset_file(fn):
     Reads a dictionary mapping integer class IDs to class names from a YOLOv5/YOLOv8
     dataset.yaml file or a .json file.  A .json file should contain a dictionary mapping
     integer category IDs to string category names.
-    
+
     Args:
-        fn (str): YOLOv5/YOLOv8 dataset file with a .yml or .yaml extension, a .json file 
+        fn (str): YOLOv5/YOLOv8 dataset file with a .yml or .yaml extension, a .json file
             mapping integer category IDs to category names, or a .txt file with a flat
             list of classes.
-            
+
     Returns:
         dict: a mapping from integer category IDs to category names
     """
-    
+
     category_id_to_name = {}
-    
+
     if fn.endswith('.yml') or fn.endswith('.yaml'):
-        
+
         with open(fn,'r') as f:
             lines = f.readlines()
-                
+
         pat = r'\d+:.+'
         for s in lines:
             if re.search(pat,s) is not None:
                 tokens = s.split(':')
                 assert len(tokens) == 2, 'Invalid token in category file {}'.format(fn)
                 category_id_to_name[int(tokens[0].strip())] = tokens[1].strip()
-                
+
     elif fn.endswith('.json'):
-        
+
         with open(fn,'r') as f:
             d_in = json.load(f)
             for k in d_in.keys():
                 category_id_to_name[int(k)] = d_in[k]
-        
+
     elif fn.endswith('.txt'):
-        
+
         with open(fn,'r') as f:
             lines = f.readlines()
         next_category_id = 0
@@ -99,20 +101,20 @@ def read_classes_from_yolo_dataset_file(fn):
                 continue
             category_id_to_name[next_category_id] = s
             next_category_id += 1
-            
+
     else:
-        
+
         raise ValueError('Unrecognized category file type: {}'.format(fn))
-        
+
     assert len(category_id_to_name) > 0, 'Failed to read class mappings from {}'.format(fn)
-    
+
     return category_id_to_name
-    
 
-def yolo_json_output_to_md_output(yolo_json_file, 
+
+def yolo_json_output_to_md_output(yolo_json_file,
                                   image_folder,
-                                  output_file, 
-                                  yolo_category_id_to_name,                              
+                                  output_file,
+                                  yolo_category_id_to_name,
                                   detector_name='unknown',
                                   image_id_to_relative_path=None,
                                   offset_yolo_class_ids=True,
@@ -121,56 +123,56 @@ def yolo_json_output_to_md_output(yolo_json_file,
                                   convert_slashes=True):
     """
     Converts a YOLOv5/YOLOv8 .json file to MD .json format.
-    
+
     Args:
-        
-        yolo_json_file (str): the .json file to convert from YOLOv5 format to MD output format
+        yolo_json_file (str): the YOLO-formatted .json file to convert to MD output format
         image_folder (str): the .json file contains relative path names, this is the path base
-        yolo_category_id_to_name (str or dict): the .json results file contains only numeric 
-            identifiers for categories, but we want names and numbers for the output format; 
-            yolo_category_id_to_name provides that mapping either as a dict or as a YOLOv5 
+        output_file (str): the MD-formatted .json file to write
+        yolo_category_id_to_name (str or dict): the .json results file contains only numeric
+            identifiers for categories, but we want names and numbers for the output format;
+            yolo_category_id_to_name provides that mapping either as a dict or as a YOLOv5
             dataset.yaml file.
-        detector_name (str, optional): a string that gets put in the output file, not otherwise 
+        detector_name (str, optional): a string that gets put in the output file, not otherwise
             used within this function
-        image_id_to_relative_path (dict, optional): YOLOv5 .json uses only basenames (e.g. 
+        image_id_to_relative_path (dict, optional): YOLOv5 .json uses only basenames (e.g.
             abc1234.JPG); by default these will be appended to the input path to create pathnames.
             If you have a flat folder, this is fine.  If you want to map base names to relative paths in
-            a more complicated way, use this parameter.   
-        offset_yolo_class_ids (bool, optional): YOLOv5 class IDs always start at zero; if you want to 
-            make the output classes start at 1, set offset_yolo_class_ids to True.    
-        truncate_to_standard_md_precision (bool, optional): YOLOv5 .json includes lots of 
+            a more complicated way, use this parameter.
+        offset_yolo_class_ids (bool, optional): YOLOv5 class IDs always start at zero; if you want to
+            make the output classes start at 1, set offset_yolo_class_ids to True.
+        truncate_to_standard_md_precision (bool, optional): YOLOv5 .json includes lots of
             (not-super-meaningful) precision, set this to truncate to COORD_DIGITS and CONF_DIGITS.
-        image_id_to_error (dict, optional): if you want to include image IDs in the output file for which 
+        image_id_to_error (dict, optional): if you want to include image IDs in the output file for which
             you couldn't prepare the input file in the first place due to errors, include them here.
         convert_slashes (bool, optional): force all slashes to be forward slashes in the output file
-    """    
-        
+    """
+
     assert os.path.isfile(yolo_json_file), \
         'Could not find YOLO .json file {}'.format(yolo_json_file)
     assert os.path.isdir(image_folder), \
         'Could not find image folder {}'.format(image_folder)
-      
+
     if image_id_to_error is None:
         image_id_to_error = {}
-            
+
     print('Converting {} to MD format and writing results to {}'.format(
         yolo_json_file,output_file))
-    
+
     if isinstance(yolo_category_id_to_name,str):
         assert os.path.isfile(yolo_category_id_to_name), \
             'YOLO category mapping specified as a string, but file does not exist: {}'.format(
                 yolo_category_id_to_name)
         yolo_category_id_to_name = read_classes_from_yolo_dataset_file(yolo_category_id_to_name)
-    
+
     if image_id_to_relative_path is None:
-        
+
         image_files = path_utils.find_images(image_folder,recursive=True)
         image_files = [os.path.relpath(fn,image_folder) for fn in image_files]
-        
+
         # YOLOv5 identifies images in .json output by ID, which is the filename without
         # extension.  If a mapping is not provided, these need to be unique.
         image_id_to_relative_path = {}
-        
+
         for fn in image_files:
             image_id = os.path.splitext(os.path.basename(fn))[0]
             if image_id in image_id_to_relative_path:
@@ -180,49 +182,49 @@ def yolo_json_output_to_md_output(yolo_json_file,
             image_id_to_relative_path[image_id] = fn
 
     image_files_relative = sorted(list(image_id_to_relative_path.values()))
-    
+
     image_file_relative_to_image_id = {}
     for image_id in image_id_to_relative_path:
         relative_path = image_id_to_relative_path[image_id]
         assert relative_path not in image_file_relative_to_image_id, \
             'Duplication image IDs in YOLO output conversion for image {}'.format(relative_path)
         image_file_relative_to_image_id[relative_path] = image_id
-        
+
     with open(yolo_json_file,'r') as f:
         detections = json.load(f)
     assert isinstance(detections,list)
-    
+
     image_id_to_detections = defaultdict(list)
-    
+
     int_formatted_image_ids = False
-    
+
     # det = detections[0]
     for det in detections:
-        
+
         # This could be a string, but if the YOLOv5 inference script sees that the strings
         # are really ints, it converts to ints.
         image_id = det['image_id']
         image_id_to_detections[image_id].append(det)
         if isinstance(image_id,int):
             int_formatted_image_ids = True
-    
+
     # If there are any ints present, everything should be ints
     if int_formatted_image_ids:
         for det in detections:
             assert isinstance(det['image_id'],int), \
                 'Found mixed int and string image IDs'
-        
+
         # Convert the keys in image_id_to_error to ints
         #
-        # This should error if we're given non-int-friendly IDs        
-        int_formatted_image_id_to_error = {}        
+        # This should error if we're given non-int-friendly IDs
+        int_formatted_image_id_to_error = {}
         for image_id in image_id_to_error:
             int_formatted_image_id_to_error[int(image_id)] = \
                 image_id_to_error[image_id]
-        image_id_to_error = int_formatted_image_id_to_error        
-           
+        image_id_to_error = int_formatted_image_id_to_error
+
     # ...if image IDs are formatted as integers in YOLO output
-    
+
     # In a modified version of val.py, we use negative category IDs to indicate an error
     # that happened during inference (typically truncated images with valid headers,
     # so corruption was not detected during val.py's initial corruption check pass.
@@ -232,17 +234,17 @@ def yolo_json_output_to_md_output(yolo_json_file,
             error_string = det['error']
             print('Caught inference-time failure {} for image {}'.format(error_string,det['image_id']))
             image_id_to_error[det['image_id']] = error_string
-            
+
     output_images = []
-    
+
     # image_file_relative = image_files_relative[10]
     for image_file_relative in tqdm(image_files_relative):
-        
+
         im = {}
         im['file'] = image_file_relative
         if convert_slashes:
             im['file'] = im['file'].replace('\\','/')
-        
+
         image_id = image_file_relative_to_image_id[image_file_relative]
         if int_formatted_image_ids:
             image_id = int(image_id)
@@ -254,7 +256,7 @@ def yolo_json_output_to_md_output(yolo_json_file,
             detections = []
         else:
             detections = image_id_to_detections[image_id]
-        
+
         image_full_path = os.path.join(image_folder,image_file_relative)
         try:
             pil_im = vis_utils.open_image(image_full_path)
@@ -264,17 +266,17 @@ def yolo_json_output_to_md_output(yolo_json_file,
             im['failure'] = 'Conversion error: {}'.format(s)
             output_images.append(im)
             continue
-        
+
         im['detections'] = []
-        
+
         image_w = pil_im.size[0]
         image_h = pil_im.size[1]
-        
+
         # det = detections[0]
         for det in detections:
-            
+
             output_det = {}
-            
+
             yolo_cat_id = int(det['category_id'])
             if offset_yolo_class_ids:
                 yolo_cat_id += 1
@@ -284,63 +286,62 @@ def yolo_json_output_to_md_output(yolo_json_file,
                 conf = ct_utils.round_float(conf,CONF_DIGITS)
             output_det['conf'] = conf
             input_bbox = det['bbox']
-            
+
             # YOLO's COCO .json is not *that* COCO-like, but it is COCO-like in
             # that the boxes are already [xmin/ymin/w/h]
             box_xmin_absolute = input_bbox[0]
             box_ymin_absolute = input_bbox[1]
             box_width_absolute = input_bbox[2]
             box_height_absolute = input_bbox[3]
-            
+
             box_xmin_relative = box_xmin_absolute / image_w
             box_ymin_relative = box_ymin_absolute / image_h
             box_width_relative = box_width_absolute / image_w
             box_height_relative = box_height_absolute / image_h
-            
+
             output_bbox = [box_xmin_relative,box_ymin_relative,
                            box_width_relative,box_height_relative]
-            
+
             if truncate_to_standard_md_precision:
                 output_bbox = ct_utils.round_float_array(output_bbox,COORD_DIGITS)
-                
+
             output_det['bbox'] = output_bbox
             im['detections'].append(output_det)
-            
-        # ...for each detection            
-                
+
+        # ...for each detection
+
         output_images.append(im)
-        
+
     # ...for each image file
-    
+
     d = {}
     d['images'] = output_images
     d['info'] = {'format_version':'1.4','detector':detector_name}
     d['detection_categories'] = {}
-        
+
     for cat_id in yolo_category_id_to_name:
         yolo_cat_id = int(cat_id)
         if offset_yolo_class_ids:
             yolo_cat_id += 1
         d['detection_categories'][str(yolo_cat_id)] = yolo_category_id_to_name[cat_id]
-    
-    with open(output_file,'w') as f:
-        json.dump(d,f,indent=1)
-            
+
+    ct_utils.write_json(output_file, d)
+
 # ...def yolo_json_output_to_md_output(...)
-    
 
-def yolo_txt_output_to_md_output(input_results_folder, 
+
+def yolo_txt_output_to_md_output(input_results_folder,
                                  image_folder,
-                                 output_file, 
+                                 output_file,
                                  detector_tag=None,
                                  truncate_to_standard_md_precision=True):
     """
     Converts a folder of YOLO-output .txt files to MD .json format.
-    
-    Less finished than the .json conversion function; this .txt conversion assumes 
-    a hard-coded mapping representing the standard MD categories (in MD indexing, 
+
+    Less finished than the .json conversion function; this .txt conversion assumes
+    a hard-coded mapping representing the standard MD categories (in MD indexing,
     1/2/3=animal/person/vehicle; in YOLO indexing, 0/1/2=animal/person/vehicle).
-    
+
     Args:
         input_results_folder (str): the folder containing YOLO-output .txt files
         image_folder (str): the folder where images live, may be the same as
@@ -349,47 +350,47 @@ def yolo_txt_output_to_md_output(input_results_folder,
             results
         detector_tag (str, optional): string to put in the 'detector' field in the
             output file
-        truncate_to_standard_md_precision (bool, optional): set this to truncate to 
+        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)
     assert os.path.isdir(image_folder)
-    
+
     ## Enumerate results files and image files
-    
+
     yolo_results_files = os.listdir(input_results_folder)
     yolo_results_files = [f for f in yolo_results_files if f.lower().endswith('.txt')]
     # print('Found {} results files'.format(len(yolo_results_files)))
-    
+
     image_files = path_utils.find_images(image_folder,recursive=False)
     image_files_relative = [os.path.basename(f) for f in image_files]
     # print('Found {} images'.format(len(image_files)))
-            
+
     image_files_relative_no_extension = [os.path.splitext(f)[0] for f in image_files_relative]
-    
+
     ## Make sure that every results file corresponds to an image
-    
+
     for f in yolo_results_files:
         result_no_extension = os.path.splitext(f)[0]
         assert result_no_extension in image_files_relative_no_extension
-    
+
     ## Build MD output data
-    
+
     # Map 0-indexed YOLO categories to 1-indexed MD categories
     yolo_cat_map = { 0: 1, 1: 2, 2: 3 }
-    
+
     images_entries = []
 
     # image_fn = image_files_relative[0]
     for image_fn in image_files_relative:
-        
-        image_name, ext = os.path.splitext(image_fn)        
+
+        image_name, ext = os.path.splitext(image_fn)
         label_fn = image_name + '.txt'
         label_path = os.path.join(input_results_folder, label_fn)
-            
+
         detections = []
-        
+
         if not os.path.exists(label_path):
             # This is assumed to be an image with no detections
             pass
@@ -397,36 +398,36 @@ def yolo_txt_output_to_md_output(input_results_folder,
             with open(label_path, newline='') as f:
                 reader = csv.reader(f, delimiter=' ')
                 for row in reader:
-                    category = yolo_cat_map[int(row[0])]    
-                    api_box = ct_utils.convert_yolo_to_xywh([float(row[1]), float(row[2]), 
+                    category = yolo_cat_map[int(row[0])]
+                    api_box = ct_utils.convert_yolo_to_xywh([float(row[1]), float(row[2]),
                                                              float(row[3]), float(row[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': api_box
                     })
-                
+
         images_entries.append({
             'file': image_fn,
             'detections': detections
         })
-    
+
     # ...for each image
-    
+
     ## Save output file
-    
+
     detector_string = 'converted_from_yolo_format'
-    
+
     if detector_tag is not None:
         detector_string = detector_tag
-        
+
     output_content = {
         'info': {
             'detector': detector_string,
@@ -440,27 +441,152 @@ def yolo_txt_output_to_md_output(input_results_folder,
         },
         'images': images_entries
     }
-    
-    with open(output_file,'w') as f:
-        json.dump(output_content,f,indent=1)
-    
+
+    ct_utils.write_json(output_file, output_content)
+
 # ...def yolo_txt_output_to_md_output(...)
 
 
 #%% Interactive driver
 
 if False:
-    
+
     pass
 
-    #%%    
-    
+    #%%
+
     input_results_folder = os.path.expanduser('~/tmp/model-version-experiments/pt-test-kru/exp/labels')
     image_folder = os.path.expanduser('~/data/KRU-test')
-    output_file = os.path.expanduser('~/data/mdv5a-yolo-pt-kru.json')    
+    output_file = os.path.expanduser('~/data/mdv5a-yolo-pt-kru.json')
     yolo_txt_output_to_md_output(input_results_folder,image_folder,output_file)
 
 
 #%% Command-line driver
 
-# TODO
+def main():
+    """
+    Command-line interface to convert YOLOv5/YOLOv8 output (.json or .txt)
+    to MegaDetector output format.
+    """
+
+    parser = argparse.ArgumentParser(
+        description='Converts YOLOv5 output (.json or .txt) to MD output format.'
+    )
+
+    # The first argument determines which series of additional arguments are supported, for
+    # json/txt input
+    subparsers = parser.add_subparsers(dest='mode', required=True,
+        help="Mode of operation: 'json' for YOLO JSON output, 'txt' for YOLO TXT output.")
+
+
+    ## 'json' mode subparser
+
+    parser_json = subparsers.add_parser('json', help='Convert YOLO-formatted .json results.')
+
+    parser_json.add_argument(
+        'yolo_json_file', type=str,
+        help='Path to the input YOLO-formatted .json results file'
+    )
+    parser_json.add_argument(
+        'image_folder', type=str,
+        help='Path to the image folder'
+    )
+    parser_json.add_argument(
+        'output_file', type=str,
+        help='Path to the MD-formatted .json output file'
+    )
+    parser_json.add_argument(
+        'yolo_category_id_to_name_file', type=str,
+        help='Path to the .yml, .yaml, .json, or .txt file mapping YOLO category IDs to names'
+    )
+    parser_json.add_argument(
+        '--detector_name', type=str, default='unknown',
+        help="Detector name to store in the output file (default: 'unknown')"
+    )
+    parser_json.add_argument(
+        '--image_id_to_relative_path_file', type=str, default=None,
+        help='Path to a .json file mapping image IDs to relative paths'
+    )
+    parser_json.add_argument(
+        '--offset_yolo_class_ids', type=str, default='true', choices=['true', 'false'],
+        help="Offset YOLO class IDs in the output (default: 'true')"
+    )
+    parser_json.add_argument(
+        '--truncate_to_standard_md_precision', type=str, default='true', choices=['true', 'false'],
+        help="Truncate coordinates and confidences to standard MD precision (default: 'true')"
+    )
+    parser_json.add_argument(
+        '--convert_slashes', type=str, default='true', choices=['true', 'false'],
+        help="Convert backslashes to forward slashes in output file paths (default: 'true')"
+    )
+
+
+    ## 'txt' mode subparser
+
+    parser_txt = subparsers.add_parser('txt', help='Convert YOLO-formatted .txt results from a folder')
+    parser_txt.add_argument(
+        'input_results_folder', type=str,
+        help='Path to the folder containing YOLO .txt output files'
+    )
+    parser_txt.add_argument(
+        'image_folder', type=str,
+        help='Path to the image folder'
+    )
+    parser_txt.add_argument(
+        'output_file', type=str,
+        help='Path to the MD-formatted .json  file output'
+    )
+    parser_txt.add_argument(
+        '--detector_tag', type=str, default=None,
+        help='Detector tag to store in the output file'
+    )
+    parser_txt.add_argument(
+        '--truncate_to_standard_md_precision', type=str, default='true', choices=['true', 'false'],
+        help="Truncate coordinates and confidences to standard MD precision (default: 'true')."
+    )
+
+    args = parser.parse_args()
+
+    if args.mode == 'json':
+
+        image_id_to_relative_path = None
+        if args.image_id_to_relative_path_file:
+            try:
+                with open(args.image_id_to_relative_path_file, 'r') as f:
+                    image_id_to_relative_path = json.load(f)
+            except Exception as e:
+                print(f"Error loading image_id_to_relative_path_file: {e}")
+                sys.exit(1)
+
+        offset_yolo_class_ids = args.offset_yolo_class_ids.lower() == 'true'
+        truncate_json = args.truncate_to_standard_md_precision.lower() == 'true'
+        convert_slashes = args.convert_slashes.lower() == 'true'
+
+        yolo_json_output_to_md_output(
+            yolo_json_file=args.yolo_json_file,
+            image_folder=args.image_folder,
+            output_file=args.output_file,
+            yolo_category_id_to_name=args.yolo_category_id_to_name_file, # Function handles reading this file
+            detector_name=args.detector_name,
+            image_id_to_relative_path=image_id_to_relative_path,
+            offset_yolo_class_ids=offset_yolo_class_ids,
+            truncate_to_standard_md_precision=truncate_json,
+            convert_slashes=convert_slashes
+        )
+        print('Converted {} to {}'.format(args.yolo_json_file,args.output_file))
+
+    elif args.mode == 'txt':
+
+        truncate_txt = args.truncate_to_standard_md_precision.lower() == 'true'
+
+        yolo_txt_output_to_md_output(
+            input_results_folder=args.input_results_folder,
+            image_folder=args.image_folder,
+            output_file=args.output_file,
+            detector_tag=args.detector_tag,
+            truncate_to_standard_md_precision=truncate_txt
+        )
+        print('Converted results from {} to {}'.format(args.input_results_folder,args.output_file))
+
+if __name__ == '__main__':
+    main()