megadetector 5.0.11__py3-none-any.whl → 5.0.13__py3-none-any.whl

megadetector/data_management/coco_to_yolo.py

@@ -0,0 +1,662 @@
+"""
+
+coco_to_yolo.py
+
+Converts a COCO-formatted dataset to a YOLO-formatted dataset, flattening
+the dataset (to a single folder) in the process.
+
+If the input and output folders are the same, writes .txt files to the input folder,
+and neither moves nor modifies images.
+
+Currently ignores segmentation masks, and errors if an annotation has a 
+segmentation polygon but no bbox.
+
+Has only been tested on a handful of COCO Camera Traps data sets; if you
+use it for more general COCO conversion, YMMV.
+
+"""
+
+#%% Imports and constants
+
+import json
+import os
+import shutil
+
+from collections import defaultdict
+from tqdm import tqdm
+
+from megadetector.utils.path_utils import safe_create_link,find_images
+
+
+#%% Support functions
+
+def write_yolo_dataset_file(yolo_dataset_file,
+                            dataset_base_dir,
+                            class_list,
+                            train_folder_relative=None,
+                            val_folder_relative=None,
+                            test_folder_relative=None):
+    """
+    Write a YOLOv5 dataset.yaml file to the absolute path [yolo_dataset_file] (should
+    have a .yaml extension, though it's only a warning if it doesn't).  
+    
+    Args:
+        yolo_dataset_file (str): the file, typically ending in .yaml or .yml, to write.  
+            Does not have to be within dataset_base_dir.
+        dataset_base_dir (str): the absolute base path of the YOLO dataset
+        class_list (list or str): an ordered list of class names (the first item will be class 0, 
+            etc.), or the name of a text file containing an ordered list of class names (one per 
+            line, starting from class zero).
+    """
+    
+    # Read class names
+    if isinstance(class_list,str):
+        with open(class_list,'r') as f:
+            class_lines = f.readlines()
+        class_lines = [s.strip() for s in class_lines]    
+        class_list = [s for s in class_lines if len(s) > 0]
+
+    if not (yolo_dataset_file.endswith('.yml') or yolo_dataset_file.endswith('.yaml')):
+        print('Warning: writing dataset file to a non-yml/yaml extension:\n{}'.format(
+            yolo_dataset_file))
+        
+    # Write dataset.yaml
+    with open(yolo_dataset_file,'w') as f:
+        
+        f.write('# Train/val sets\n')
+        f.write('path: {}\n'.format(dataset_base_dir))
+        if train_folder_relative is not None:
+            f.write('train: {}\n'.format(train_folder_relative))
+        if val_folder_relative is not None:
+            f.write('val: {}\n'.format(val_folder_relative))
+        if test_folder_relative is not None:
+            f.write('val: {}\n'.format(test_folder_relative))
+            
+        f.write('\n')
+        
+        f.write('# Classes\n')
+        f.write('names:\n')
+        for i_class,class_name in enumerate(class_list):
+            f.write('  {}: {}\n'.format(i_class,class_name))
+
+# ...def write_yolo_dataset_file(...)
+
+            
+def coco_to_yolo(input_image_folder,
+                 output_folder,
+                 input_file,
+                 source_format='coco',
+                 overwrite_images=False,
+                 create_image_and_label_folders=False,
+                 class_file_name='classes.txt',
+                 allow_empty_annotations=False,
+                 clip_boxes=False,
+                 image_id_to_output_image_json_file=None,
+                 images_to_exclude=None,
+                 path_replacement_char='#',
+                 category_names_to_exclude=None,
+                 category_names_to_include=None,
+                 write_output=True,
+                 flatten_paths=True):
+    """
+    Converts a COCO-formatted dataset to a YOLO-formatted dataset, optionally flattening the 
+    dataset to a single folder in the process.
+    
+    If the input and output folders are the same, writes .txt files to the input folder,
+    and neither moves nor modifies images.
+    
+    Currently ignores segmentation masks, and errors if an annotation has a 
+    segmentation polygon but no bbox.
+    
+    Args:
+        input_image_folder (str): the folder where images live; filenames in the COCO .json
+            file [input_file] should be relative to this folder
+        output_folder (str): the base folder for the YOLO dataset
+        input_file (str): a .json file in COCO format; can be the same as [input_image_folder], in which case
+            images are left alone.
+        source_format (str, optional): can be 'coco' (default) or 'coco_camera_traps'.  The only difference
+            is that when source_format is 'coco_camera_traps', we treat an image with a non-bbox
+            annotation with a category id of 0 as a special case, i.e. that's how an empty image
+            is indicated.  The original COCO standard is a little ambiguous on this issue.  If
+            source_format is 'coco', we either treat images as empty or error, depending on the value
+            of [allow_empty_annotations].  [allow_empty_annotations] has no effect if source_format is
+            'coco_camera_traps'.
+        create_image_and_label_folder (bool, optional): whether to create separate folders called 'images' and
+            'labels' in the YOLO output folder.  If create_image_and_label_folders is False, 
+            a/b/c/image001.jpg will become a#b#c#image001.jpg, and the corresponding text file will 
+            be a#b#c#image001.txt.  If create_image_and_label_folders is True, a/b/c/image001.jpg will become 
+            images/a#b#c#image001.jpg, and the corresponding text file will be 
+            labels/a#b#c#image001.txt.    
+        clip_boxes (bool, optional): whether to clip bounding box coordinates to the range [0,1] before
+            converting to YOLO xywh format
+        image_id_to_output_image_json_file (str, optional): an optional *output* file, to which we will write
+            a mapping from image IDs to output file names
+        images_to_exclude (list, optional): a list of image files (relative paths in the input folder) that we 
+            should ignore
+        path_replacement_char (str, optional): only relevant if [flatten_paths] is True; this is used to replace
+            path separators, e.g. if [path_replacement_char] is '#' and [flatten_paths] is True, a/b/c/d.jpg
+            becomes a#b#c#d.jpg
+        category_names_to_exclude (str, optional): category names that should not be represented in the
+            YOLO output; only impacts annotations, does not prevent copying images.  There's almost no reason
+            you would want to specify this and [category_names_to_include]. 
+        category_names_to_include (str, optional): allow-list of category names that should be represented in the
+            YOLO output; only impacts annotations, does not prevent copying images.  There's almost no reason
+            you would want to specify this and [category_names_to_exclude]. 
+        write_output (bool, optional): determines whether we actually copy images and write annotations;
+            setting this to False mostly puts this function in "dry run" "mode.  The class list
+            file is written regardless of the value of write_output.
+    
+    Returns:
+        dict: information about the coco --> yolo mapping, containing at least the fields:
+        
+        - class_list_filename: the filename to which we wrote the flat list of class names required 
+          by the YOLO format.
+        - source_image_to_dest_image: a dict mapping source images to destination images
+        - coco_id_to_yolo_id: a dict mapping COCO category IDs to YOLO category IDs        
+    """
+        
+    ## Validate input
+    
+    if category_names_to_include is not None and category_names_to_exclude is not None:
+        raise ValueError('category_names_to_include and category_names_to_exclude are mutually exclusive')
+        
+    if output_folder is None:
+        output_folder = input_image_folder
+    
+    if images_to_exclude is not None:
+        images_to_exclude = set(images_to_exclude)
+            
+    if category_names_to_exclude is None:
+        category_names_to_exclude = {}
+            
+    assert os.path.isdir(input_image_folder)
+    assert os.path.isfile(input_file)
+    os.makedirs(output_folder,exist_ok=True)
+    
+    if (output_folder == input_image_folder) and (overwrite_images) and \
+        (not create_image_and_label_folders) and (not flatten_paths):
+            print('Warning: output folder and input folder are the same, disabling overwrite_images')
+            overwrite_images = False
+            
+    ## Read input data
+    
+    with open(input_file,'r') as f:
+        data = json.load(f)
+        
+        
+    ## Parse annotations
+  
+    image_id_to_annotations = defaultdict(list)
+    
+    # i_ann = 0; ann = data['annotations'][0]
+    for i_ann,ann in enumerate(data['annotations']):
+        
+        # Make sure no annotations have *only* segmentation data 
+        if ( \
+            ('segmentation' in ann.keys()) and \
+            (ann['segmentation'] is not None) and \
+            (len(ann['segmentation']) > 0) ) \
+            and \
+            (('bbox' not in ann.keys()) or (ann['bbox'] is None) or (len(ann['bbox'])==0)):
+                raise ValueError('Oops: segmentation data present without bbox information, ' + \
+                                 'this script isn\'t ready for this dataset')
+        
+        image_id_to_annotations[ann['image_id']].append(ann)
+        
+    print('Parsed annotations for {} images'.format(len(image_id_to_annotations)))
+        
+    # Re-map class IDs to make sure they run from 0...n-classes-1
+    #
+    # Note: this allows unused categories in the output data set.  This is OK for
+    # some training pipelines, not for others.
+    next_category_id = 0
+    coco_id_to_yolo_id = {}    
+    coco_id_to_name = {}
+    yolo_id_to_name = {}
+    coco_category_ids_to_exclude = set()
+    
+    for category in data['categories']:
+        coco_id_to_name[category['id']] = category['name']
+        if (category_names_to_include is not None) and \
+            (category['name'] not in category_names_to_include):
+            coco_category_ids_to_exclude.add(category['id'])
+            continue
+        elif (category['name'] in category_names_to_exclude):
+            coco_category_ids_to_exclude.add(category['id'])
+            continue               
+        assert category['id'] not in coco_id_to_yolo_id
+        coco_id_to_yolo_id[category['id']] = next_category_id
+        yolo_id_to_name[next_category_id] = category['name']
+        next_category_id += 1
+        
+    
+    ## Process images (everything but I/O)
+    
+    # List of dictionaries with keys 'source_image','dest_image','bboxes','dest_txt'
+    images_to_copy = []
+    
+    missing_images = []
+    excluded_images = []
+    
+    image_names = set()
+    
+    typical_image_extensions = set(['.jpg','.jpeg','.png','.gif','.tif','.bmp'])
+    
+    printed_empty_annotation_warning = False
+    
+    image_id_to_output_image_name = {}
+    
+    print('Processing annotations')
+    
+    n_clipped_boxes = 0
+    n_total_boxes = 0
+    
+    # i_image = 0; im = data['images'][i_image]
+    for i_image,im in tqdm(enumerate(data['images']),total=len(data['images'])):
+                
+        output_info = {}
+        source_image = os.path.join(input_image_folder,im['file_name'])        
+        output_info['source_image'] = source_image
+        
+        if images_to_exclude is not None and im['file_name'] in images_to_exclude:
+            excluded_images.append(im['file_name'])
+            continue
+        
+        tokens = os.path.splitext(im['file_name'])
+        if tokens[1].lower() not in typical_image_extensions:
+            print('Warning: unusual image file name {}'.format(im['file_name']))
+                
+        if flatten_paths:
+            image_name = tokens[0].replace('\\','/').replace('/',path_replacement_char) + \
+                '_' + str(i_image).zfill(6)            
+        else:
+            image_name = tokens[0]
+
+        assert image_name not in image_names, 'Image name collision for {}'.format(image_name)
+        image_names.add(image_name)
+        
+        assert im['id'] not in image_id_to_output_image_name
+        image_id_to_output_image_name[im['id']] = image_name
+        
+        dest_image_relative = image_name + tokens[1]
+        output_info['dest_image_relative'] = dest_image_relative
+        dest_txt_relative = image_name + '.txt'
+        output_info['dest_txt_relative'] = dest_txt_relative
+        output_info['bboxes'] = []
+        
+        # assert os.path.isfile(source_image), 'Could not find image {}'.format(source_image)
+        if not os.path.isfile(source_image):
+            print('Warning: could not find image {}'.format(source_image))
+            missing_images.append(im['file_name'])
+            continue
+        
+        image_id = im['id']
+        
+        image_bboxes = []
+            
+        if image_id in image_id_to_annotations:
+                        
+            for ann in image_id_to_annotations[image_id]:
+                
+                # If this annotation has no bounding boxes...
+                if 'bbox' not in ann or ann['bbox'] is None or len(ann['bbox']) == 0:
+                    
+                    if source_format == 'coco':
+                    
+                        if not allow_empty_annotations:
+                            # This is not entirely clear from the COCO spec, but it seems to be consensus
+                            # that if you want to specify an image with no objects, you don't include any
+                            # annotations for that image.
+                            raise ValueError('If an annotation exists, it should have content')
+                        else:
+                            continue
+                    
+                    elif source_format == 'coco_camera_traps':
+                        
+                        # We allow empty bbox lists in COCO camera traps; this is typically a negative
+                        # example in a dataset that has bounding boxes, and 0 is typically the empty 
+                        # category.
+                        if ann['category_id'] != 0:
+                            if not printed_empty_annotation_warning:
+                                printed_empty_annotation_warning = True
+                                print('Warning: non-bbox annotation found with category {}'.format(
+                                    ann['category_id']))
+                        continue
+                    
+                    else:
+                        
+                        raise ValueError('Unrecognized COCO variant: {}'.format(source_format))
+                    
+                # ...if this is an empty annotation
+                
+                coco_bbox = ann['bbox']
+                
+                # This category isn't in our category list.  This typically corresponds to whole sets
+                # of images that were excluded from the YOLO set.
+                if ann['category_id'] in coco_category_ids_to_exclude:
+                    continue
+                
+                yolo_category_id = coco_id_to_yolo_id[ann['category_id']]
+                
+                # COCO: [x_min, y_min, width, height] in absolute coordinates
+                # YOLO: [class, x_center, y_center, width, height] in normalized coordinates
+                
+                # Convert from COCO coordinates to YOLO coordinates
+                img_w = im['width']
+                img_h = im['height']
+                                
+                if source_format in ('coco','coco_camera_traps'):
+                    
+                    x_min_absolute = coco_bbox[0]
+                    y_min_absolute = coco_bbox[1]
+                    box_w_absolute = coco_bbox[2]
+                    box_h_absolute = coco_bbox[3]
+                    
+                    x_center_absolute = (x_min_absolute + (x_min_absolute + box_w_absolute)) / 2
+                    y_center_absolute = (y_min_absolute + (y_min_absolute + box_h_absolute)) / 2
+                    
+                    x_center_relative = x_center_absolute / img_w
+                    y_center_relative = y_center_absolute / img_h
+                    
+                    box_w_relative = box_w_absolute / img_w
+                    box_h_relative = box_h_absolute / img_h
+                    
+                else:
+                    
+                    raise ValueError('Unrecognized source format {}'.format(source_format))
+                
+                if clip_boxes:
+                    
+                    clipped_box = False
+                    
+                    box_right = x_center_relative + (box_w_relative / 2.0)                    
+                    if box_right > 1.0:
+                        clipped_box = True
+                        overhang = box_right - 1.0
+                        box_w_relative -= overhang
+                        x_center_relative -= (overhang / 2.0)
+
+                    box_bottom = y_center_relative + (box_h_relative / 2.0)                                        
+                    if box_bottom > 1.0:
+                        clipped_box = True
+                        overhang = box_bottom - 1.0
+                        box_h_relative -= overhang
+                        y_center_relative -= (overhang / 2.0)
+                    
+                    box_left = x_center_relative - (box_w_relative / 2.0)
+                    if box_left < 0.0:
+                        clipped_box = True
+                        overhang = abs(box_left)
+                        box_w_relative -= overhang
+                        x_center_relative += (overhang / 2.0)
+                        
+                    box_top = y_center_relative - (box_h_relative / 2.0)
+                    if box_top < 0.0:
+                        clipped_box = True
+                        overhang = abs(box_top)
+                        box_h_relative -= overhang
+                        y_center_relative += (overhang / 2.0)
+                        
+                    if clipped_box:
+                        n_clipped_boxes += 1
+                
+                yolo_box = [yolo_category_id,
+                            x_center_relative, y_center_relative, 
+                            box_w_relative, box_h_relative]
+                
+                image_bboxes.append(yolo_box)
+                n_total_boxes += 1
+                
+            # ...for each annotation 
+            
+        # ...if this image has annotations
+        
+        output_info['bboxes'] = image_bboxes
+        
+        images_to_copy.append(output_info)        
+    
+    # ...for each image
+        
+    print('\nWriting {} boxes ({} clipped) for {} images'.format(n_total_boxes,
+                                                               n_clipped_boxes,len(images_to_copy)))
+    print('{} missing images (of {})'.format(len(missing_images),len(data['images'])))
+    
+    if images_to_exclude is not None:
+        print('{} excluded images (of {})'.format(len(excluded_images),len(data['images'])))
+            
+    
+    ## Write output
+    
+    print('Generating class list')
+    
+    class_list_filename = os.path.join(output_folder,class_file_name)
+    with open(class_list_filename, 'w') as f:
+        print('Writing class list to {}'.format(class_list_filename))
+        for i_class in range(0,len(yolo_id_to_name)):
+            # Category IDs should range from 0..N-1
+            assert i_class in yolo_id_to_name
+            f.write(yolo_id_to_name[i_class] + '\n')
+    
+    if image_id_to_output_image_json_file is not None:
+        print('Writing image ID mapping to {}'.format(image_id_to_output_image_json_file))
+        with open(image_id_to_output_image_json_file,'w') as f:
+            json.dump(image_id_to_output_image_name,f,indent=1)
+    
+
+    if (output_folder == input_image_folder) and (not create_image_and_label_folders):
+        print('Creating annotation files (not copying images, input and output folder are the same)')
+    else:
+        print('Copying images and creating annotation files')
+
+    if create_image_and_label_folders:
+        dest_image_folder = os.path.join(output_folder,'images')
+        dest_txt_folder = os.path.join(output_folder,'labels')
+    else:
+        dest_image_folder = output_folder
+        dest_txt_folder = output_folder
+        
+    source_image_to_dest_image = {}
+    
+    # TODO: parallelize this loop
+    #
+    # output_info = images_to_copy[0]
+    for output_info in tqdm(images_to_copy):
+
+        source_image = output_info['source_image']
+        dest_image_relative = output_info['dest_image_relative']
+        dest_txt_relative = output_info['dest_txt_relative']
+        
+        dest_image = os.path.join(dest_image_folder,dest_image_relative)
+        dest_txt = os.path.join(dest_txt_folder,dest_txt_relative)
+        
+        source_image_to_dest_image[source_image] = dest_image
+        
+        if write_output:
+            
+            os.makedirs(os.path.dirname(dest_image),exist_ok=True)        
+            os.makedirs(os.path.dirname(dest_txt),exist_ok=True)
+            
+            if not create_image_and_label_folders:
+                assert os.path.dirname(dest_image) == os.path.dirname(dest_txt)
+            
+            if (not os.path.isfile(dest_image)) or (overwrite_images):
+                shutil.copyfile(source_image,dest_image)
+        
+            bboxes = output_info['bboxes']        
+            
+            # Only write an annotation file if there are bounding boxes.  Images with 
+            # no .txt files are treated as hard negatives, at least by YOLOv5:
+            #
+            # https://github.com/ultralytics/yolov5/issues/3218
+            #
+            # I think this is also true for images with empty .txt files, but 
+            # I'm using the convention suggested on that issue, i.e. hard 
+            # negatives are expressed as images without .txt files.
+            if len(bboxes) > 0:
+                
+                with open(dest_txt,'w') as f:
+                    
+                    # bbox = bboxes[0]
+                    for bbox in bboxes:
+                        assert len(bbox) == 5
+                        s = '{} {} {} {} {}'.format(bbox[0],bbox[1],bbox[2],bbox[3],bbox[4])
+                        f.write(s + '\n')
+                        
+        # ...if we're actually writing output
+        
+    # ...for each image
+    
+    coco_to_yolo_info = {}
+    coco_to_yolo_info['class_list_filename'] = class_list_filename
+    coco_to_yolo_info['source_image_to_dest_image'] = source_image_to_dest_image
+    coco_to_yolo_info['coco_id_to_yolo_id'] = coco_id_to_yolo_id
+    
+    return coco_to_yolo_info
+
+# ...def coco_to_yolo(...)
+
+
+def create_yolo_symlinks(source_folder,images_folder,labels_folder,
+                         class_list_file=None,
+                         class_list_output_name='object.data',
+                         force_lowercase_image_extension=False):
+    """
+    Given a YOLO-formatted folder of images and .txt files, creates a folder
+    of symlinks to all the images, and a folder of symlinks to all the labels. 
+    Used to support preview/editing tools that assume images and labels are in separate 
+    folders.
+    
+    :meta private:
+    """    
+    
+    assert source_folder != images_folder and source_folder != labels_folder
+    
+    os.makedirs(images_folder,exist_ok=True)
+    os.makedirs(labels_folder,exist_ok=True)
+    
+    image_files_relative = find_images(source_folder,recursive=True,return_relative_paths=True)
+    
+    # image_fn_relative = image_files_relative[0]=    
+    for image_fn_relative in tqdm(image_files_relative):
+        
+        source_file_abs = os.path.join(source_folder,image_fn_relative)
+        target_file_abs = os.path.join(images_folder,image_fn_relative)
+        
+        if force_lowercase_image_extension:
+            tokens = os.path.splitext(target_file_abs)
+            target_file_abs = tokens[0] + tokens[1].lower()
+            
+        os.makedirs(os.path.dirname(target_file_abs),exist_ok=True)
+        safe_create_link(source_file_abs,target_file_abs)
+        source_annotation_file_abs = os.path.splitext(source_file_abs)[0] + '.txt'
+        if os.path.isfile(source_annotation_file_abs):
+            target_annotation_file_abs = \
+                os.path.splitext(os.path.join(labels_folder,image_fn_relative))[0] + '.txt'
+            os.makedirs(os.path.dirname(target_annotation_file_abs),exist_ok=True)
+            safe_create_link(source_annotation_file_abs,target_annotation_file_abs)
+
+    # ...for each image  
+
+    if class_list_file is not None:
+        target_class_list_file = os.path.join(labels_folder,class_list_output_name)
+        safe_create_link(class_list_file,target_class_list_file)
+
+# ...def create_yolo_symlinks(...)
+
+
+#%% Interactive driver
+
+if False:
+    
+    pass
+
+    #%% Options
+    
+    input_file = os.path.expanduser('~/data/md-test-coco.json')
+    image_folder = os.path.expanduser('~/data/md-test')    
+    output_folder = os.path.expanduser('~/data/md-test-yolo')    
+    create_image_and_label_folders=False
+    class_file_name='classes.txt'
+    allow_empty_annotations=False
+    clip_boxes=False
+    image_id_to_output_image_json_file=None
+    images_to_exclude=None
+    path_replacement_char='#'
+    category_names_to_exclude=None
+                                  
+                                               
+    #%% Programmatic execution
+    
+    coco_to_yolo_results = coco_to_yolo(image_folder,output_folder,input_file,
+                     source_format='coco',
+                     overwrite_images=False,
+                     create_image_and_label_folders=create_image_and_label_folders,
+                     class_file_name=class_file_name,
+                     allow_empty_annotations=allow_empty_annotations,
+                     clip_boxes=clip_boxes)                     
+    
+    create_yolo_symlinks(source_folder=output_folder,
+                         images_folder=output_folder + '/images',
+                         labels_folder=output_folder + '/labels',
+                         class_list_file=coco_to_yolo_results['class_list_filename'],
+                         class_list_output_name='object.data',
+                         force_lowercase_image_extension=True)
+
+
+    #%% Prepare command-line example
+
+    s = 'python coco_to_yolo.py {} {} {} --create_bounding_box_editor_symlinks'.format(
+        image_folder,output_folder,input_file)
+    print(s)
+    import clipboard; clipboard.copy(s)    
+
+
+#%% Command-line driver
+
+import sys,argparse
+
+def main():
+
+    parser = argparse.ArgumentParser(
+        description='Convert COCO-formatted data to YOLO format, flattening the image structure')
+    
+    # input_image_folder,output_folder,input_file
+    
+    parser.add_argument(
+        'input_folder',
+        type=str,
+        help='Path to input images')
+    
+    parser.add_argument(
+        'output_folder',
+        type=str,
+        help='Path to flat, YOLO-formatted dataset')
+    
+    parser.add_argument(
+        'input_file',
+        type=str,
+        help='Path to COCO dataset file (.json)')
+    
+    parser.add_argument(
+        '--create_bounding_box_editor_symlinks',
+        action='store_true',
+        help='Prepare symlinks so the whole folder appears to contain "images" and "labels" folderss')        
+    
+    if len(sys.argv[1:]) == 0:
+        parser.print_help()
+        parser.exit()
+
+    args = parser.parse_args()
+
+    coco_to_yolo_results = coco_to_yolo(args.input_folder,args.output_folder,args.input_file)
+    
+    if args.create_bounding_box_editor_symlinks:
+        create_yolo_symlinks(source_folder=args.output_folder,
+                             images_folder=args.output_folder + '/images',
+                             labels_folder=args.output_folder + '/labels',
+                             class_list_file=coco_to_yolo_results['class_list_filename'],
+                             class_list_output_name='object.data',
+                             force_lowercase_image_extension=True)
+    
+if __name__ == '__main__':
+    main()

megadetector/data_management/databases/add_width_and_height_to_db.py

@@ -0,0 +1,33 @@
+"""
+
+add_width_and_height_to_db.py
+
+Grabs width and height from actual image files for a .json database that is missing w/h.
+
+TODO: this is a one-off script waiting to be cleaned up for more general use.
+
+"""
+
+#%% Imports and constants
+
+import json
+from PIL import Image
+
+datafile = '/datadrive/snapshotserengeti/databases/snapshotserengeti.json'
+image_base = '/datadrive/snapshotserengeti/images/'
+
+def main():
+
+    with open(datafile,'r') as f:
+        data = json.load(f)
+
+    for im in data['images']:
+        if 'height' not in im:
+            im_w, im_h = Image.open(image_base+im['file_name']).size
+            im['height'] = im_h
+            im['width'] = im_w
+
+    json.dump(data, open(datafile,'w'))
+
+if __name__ == '__main__':
+    main()