megadetector 5.0.6__py3-none-any.whl → 5.0.8__py3-none-any.whl

data_management/labelme_to_coco.py

@@ -2,7 +2,7 @@
 #
 # labelme_to_coco.py
 #
-# Converts a folder of labelme-formatted .json files to COCO.
+# Converts a folder of labelme-formatted .json files to COCO format.
 #
 ########
 
@@ -15,21 +15,201 @@ import uuid
 from md_utils import path_utils
 from md_visualization.visualization_utils import open_image
 
+from multiprocessing.pool import Pool, ThreadPool
+from functools import partial
+
 from tqdm import tqdm
 
 
-#%% Functions
+#%% Support functions
+
+def add_category(category_name,category_name_to_id,candidate_category_id=0):
+    """
+    Add the category [category_name] to the dict [category_name_to_id], by default
+    using the next available integer index.
+    """
+    
+    if category_name in category_name_to_id:
+        return category_name_to_id[category_name]
+    while candidate_category_id in category_name_to_id.values():
+        candidate_category_id += 1
+    category_name_to_id[category_name] = candidate_category_id
+    return candidate_category_id
+
+
+def _process_labelme_file(image_fn_relative,input_folder,use_folders_as_labels,
+                          no_json_handling,validate_image_sizes,
+                          category_name_to_id,allow_new_categories=True):
+    """
+    Internal function for processing each image; this support function facilitates parallelization.
+    """
+    
+    result = {}
+    result['im'] = None
+    result['annotations_this_image'] = None
+    result['status'] = None
+    
+    image_fn_abs = os.path.join(input_folder,image_fn_relative)
+    json_fn_abs = os.path.splitext(image_fn_abs)[0] + '.json'
+    
+    im = {}
+    im['id'] = image_fn_relative
+    im['file_name'] = image_fn_relative
+    
+    # If there's no .json file for this image...
+    if not os.path.isfile(json_fn_abs):
+        
+        # Either skip it...
+        if no_json_handling == 'skip':
+            print('Skipping image {} (no .json file)'.format(image_fn_relative))
+            result['status'] = 'skipped (no .json file)'
+            return result
+        
+        # ...or error
+        elif no_json_handling == 'error':
+            raise ValueError('Image file {} has no corresponding .json file'.format(
+                image_fn_relative))
+        
+        # ...or treat it as empty.
+        elif no_json_handling == 'empty':
+            try:
+                pil_im = open_image(image_fn_abs)
+            except Exception:
+                print('Warning: error opening image {}, skipping'.format(image_fn_abs))
+                result['status'] = 'image load error'
+                return result
+            im['width'] = pil_im.width
+            im['height'] = pil_im.height
+            
+            # Just in case we need to differentiate between "no .json file" and "a .json file with no annotations"
+            im['no_labelme_json'] = True
+            shapes = []
+        else:
+            raise ValueError('Unrecognized specifier {} for handling images with no .json files'.format(
+                no_json_handling))            
+    
+    # If we found a .json file for this image...
+    else:
+        
+        # Read the .json file
+        with open(json_fn_abs,'r') as f:
+            labelme_data = json.load(f)            
+        im['width'] = labelme_data['imageWidth']
+        im['height'] = labelme_data['imageHeight']
+        
+        if validate_image_sizes:
+            try:
+                pil_im = open_image(image_fn_abs)
+            except Exception:
+                print('Warning: error opening image {} for size validation, skipping'.format(image_fn_abs))
+                result['status'] = 'skipped (size validation error)'
+                return result
+            if not (im['width'] == pil_im.width and im['height'] == pil_im.height):
+                print('Warning: image size validation error for file {}'.format(image_fn_relative))
+                im['width'] = pil_im.width
+                im['height'] = pil_im.height
+                im['labelme_width'] = labelme_data['imageWidth']
+                im['labelme_height'] = labelme_data['imageHeight']
+
+        shapes = labelme_data['shapes']
+
+        if ('flags' in labelme_data) and (len(labelme_data['flags']) > 0):
+            im['flags'] = labelme_data['flags']
+
+    annotations_this_image = []
+    
+    if len(shapes) == 0:
+        
+        if allow_new_categories:
+            category_id = add_category('empty',category_name_to_id)
+        else:
+            assert 'empty' in category_name_to_id
+            category_id = category_name_to_id['empty']
+            
+        ann = {}
+        ann['id'] = str(uuid.uuid1())
+        ann['image_id'] = im['id']
+        ann['category_id'] = category_id
+        ann['sequence_level_annotation'] = False
+        annotations_this_image.append(ann)
+        
+    else:
+        
+        for shape in shapes:
+            
+            if shape['shape_type'] != 'rectangle':
+                print('Only rectangles are supported, skipping an annotation of type {} in {}'.format(
+                    shape['shape_type'],image_fn_relative))
+                continue
+            
+            if use_folders_as_labels:
+                category_name = os.path.basename(os.path.dirname(image_fn_abs))
+            else:
+                category_name = shape['label']                
+            
+            if allow_new_categories:
+                category_id = add_category(category_name,category_name_to_id)
+            else:
+                assert category_name in category_name_to_id
+                category_id = category_name_to_id[category_name]
+            
+            points = shape['points']
+            if len(points) != 2:
+                print('Warning: illegal rectangle with {} points for {}'.format(
+                    len(points),image_fn_relative))
+                continue
+            
+            p0 = points[0]
+            p1 = points[1]
+            x0 = min(p0[0],p1[0])
+            x1 = max(p0[0],p1[0])
+            y0 = min(p0[1],p1[1])
+            y1 = max(p0[1],p1[1])
+            
+            bbox = [x0,y0,abs(x1-x0),abs(y1-y0)]
+            ann = {}
+            ann['id'] = str(uuid.uuid1())
+            ann['image_id'] = im['id']
+            ann['category_id'] = category_id
+            ann['sequence_level_annotation'] = False
+            ann['bbox'] = bbox
+            annotations_this_image.append(ann)
+            
+        # ...for each shape
+        
+    result['im'] = im
+    result['annotations_this_image'] = annotations_this_image
+
+    return result
+    
+# ...def _process_labelme_file(...)
+
 
-def labelme_to_coco(input_folder,output_file=None,category_id_to_category_name=None,
-                    empty_category_name='empty',empty_category_id=None,info_struct=None,
-                    relative_paths_to_include=None,relative_paths_to_exclude=None,
-                    use_folders_as_labels=False,recursive=True,no_json_handling='skip',
-                    validate_image_sizes=True):
+#%% Main function
+
+def labelme_to_coco(input_folder,
+                    output_file=None,
+                    category_id_to_category_name=None,
+                    empty_category_name='empty',
+                    empty_category_id=None,
+                    info_struct=None,
+                    relative_paths_to_include=None,
+                    relative_paths_to_exclude=None,
+                    use_folders_as_labels=False,
+                    recursive=True,
+                    no_json_handling='skip',
+                    validate_image_sizes=True,
+                    max_workers=1, 
+                    use_threads=True):
     """
     Find all images in [input_folder] that have corresponding .json files, and convert
     to a COCO .json file.
     
-    Currently only supports bounding box annotations.
+    Currently only supports bounding box annotations and image-level flags (i.e., does not
+    support point or general polygon annotations).
+    
+    Labelme's image-level flags don't quite fit the COCO annotations format, so they are attached
+    to image objects, rather than annotation objects.
     
     If output_file is None, just returns the resulting dict, does not write to file.    
     
@@ -48,32 +228,59 @@ def labelme_to_coco(input_folder,output_file=None,category_id_to_category_name=N
         
     * 'skip': ignore image files with no corresponding .json files
     * 'empty': treat image files with no corresponding .json files as empty
-    * 'error': throw an error when an image file has no corresponding .json file 
+    * 'error': throw an error when an image file has no corresponding .json file    
     """
     
+    if max_workers > 1:
+        assert category_id_to_category_name is not None, \
+            'When parallelizing labelme --> COCO conversion, you must supply a category mapping'
+            
     if category_id_to_category_name is None:
         category_name_to_id = {}
     else:
         category_name_to_id = {v: k for k, v in category_id_to_category_name.items()}
-        
     for category_name in category_name_to_id:
         try:
             category_name_to_id[category_name] = int(category_name_to_id[category_name])
         except ValueError:
             raise ValueError('Category IDs must be ints or string-formatted ints')
+    
+    # If the user supplied an explicit empty category ID, and the empty category
+    # name is already in category_name_to_id, make sure they match.
+    if empty_category_id is not None:
+        if empty_category_name in category_name_to_id:
+            assert category_name_to_id[empty_category_name] == empty_category_id, \
+                'Ambiguous empty category specification'
+        if empty_category_id in category_id_to_category_name:
+            assert category_id_to_category_name[empty_category_id] == empty_category_name, \
+                'Ambiguous empty category specification'
+    else:
+        if empty_category_name in category_name_to_id:
+            empty_category_id = category_name_to_id[empty_category_name]
 
+    del category_id_to_category_name
+            
     # Enumerate images
+    print('Enumerating images in {}'.format(input_folder))    
     image_filenames_relative = path_utils.find_images(input_folder,recursive=recursive,
-                                                      return_relative_paths=True)    
-    
-    def add_category(category_name,candidate_category_id=0):
-        if category_name in category_name_to_id:
-            return category_name_to_id[category_name]
-        while candidate_category_id in category_name_to_id.values():
-            candidate_category_id += 1
-        category_name_to_id[category_name] = candidate_category_id
-        return candidate_category_id
+                                                      return_relative_paths=True,
+                                                      convert_slashes=True)    
+    
+    # Remove any images we're supposed to skip
+    if (relative_paths_to_include is not None) or (relative_paths_to_exclude is not None):
+        image_filenames_relative_to_process = []
+        for image_fn_relative in image_filenames_relative:
+            if relative_paths_to_include is not None and image_fn_relative not in relative_paths_to_include:
+                continue
+            if relative_paths_to_exclude is not None and image_fn_relative in relative_paths_to_exclude:
+                continue
+            image_filenames_relative_to_process.append(image_fn_relative)
+        print('Processing {} of {} images'.format(
+            len(image_filenames_relative_to_process),
+            len(image_filenames_relative)))
+        image_filenames_relative = image_filenames_relative_to_process
     
+    # If the user supplied a category ID to use for empty images...
     if empty_category_id is not None:
         try:
             empty_category_id = int(empty_category_id)
@@ -81,122 +288,52 @@ def labelme_to_coco(input_folder,output_file=None,category_id_to_category_name=N
             raise ValueError('Category IDs must be ints or string-formatted ints')
         
     if empty_category_id is None:
-        empty_category_id = add_category(empty_category_name)
-        
-    images = []
-    annotations = []
-    
-    # image_fn_relative = image_filenames_relative[0]
-    for image_fn_relative in tqdm(image_filenames_relative):
-        
-        if relative_paths_to_include is not None and image_fn_relative not in relative_paths_to_include:
-            continue
-        if relative_paths_to_exclude is not None and image_fn_relative in relative_paths_to_exclude:
-            continue
-        
-        image_fn_abs = os.path.join(input_folder,image_fn_relative)
-        json_fn_abs = os.path.splitext(image_fn_abs)[0] + '.json'
-        
-        im = {}
-        im['id'] = image_fn_relative
-        im['file_name'] = image_fn_relative
-        
-        # If there's no .json file for this image...
-        if not os.path.isfile(json_fn_abs):
+        empty_category_id = add_category(empty_category_name,category_name_to_id)
             
-            # Either skip it...
-            if no_json_handling == 'skip':
-                continue
+    if max_workers <= 1:
+        
+        image_results = []
+        for image_fn_relative in tqdm(image_filenames_relative):
             
-            # ...or error
-            elif no_json_handling == 'error':
-                raise ValueError('Image file {} has no corresponding .json file'.format(
-                    image_fn_relative))
+            result = _process_labelme_file(image_fn_relative,input_folder,use_folders_as_labels,
+                                      no_json_handling,validate_image_sizes,
+                                      category_name_to_id,allow_new_categories=True)        
+            image_results.append(result)
             
-            # ...or treat it as empty.
-            elif no_json_handling == 'empty':
-                try:
-                    pil_im = open_image(image_fn_abs)
-                except Exception:
-                    print('Warning: error opening image {}, skipping'.format(image_fn_abs))
-                    continue
-                im['width'] = pil_im.width
-                im['height'] = pil_im.height
-                shapes = []
-            else:
-                raise ValueError('Unrecognized specifier {} for handling images with no .json files'.format(
-                    no_json_handling))            
+    else:                      
         
-        # If we found a .json file for this image...
+        n_workers = min(max_workers,len(image_filenames_relative))
+        assert category_name_to_id is not None
+        
+        if use_threads:
+            pool = ThreadPool(n_workers)
         else:
-            
-            # Read the .json file
-            with open(json_fn_abs,'r') as f:
-                labelme_data = json.load(f)            
-            im['width'] = labelme_data['imageWidth']
-            im['height'] = labelme_data['imageHeight']
-            
-            if validate_image_sizes:
-                try:
-                    pil_im = open_image(image_fn_abs)
-                except Exception:
-                    print('Warning: error opening image {}, skipping'.format(image_fn_abs))
-                    continue                
-                assert im['width'] == pil_im.width and im['height'] == pil_im.height, \
-                    'Image size validation error for file {}'.format(image_fn_relative)                
-                
-            shapes = labelme_data['shapes']
+            pool = Pool(n_workers)
         
-        if len(shapes) == 0:
-            
-            category_id = add_category('empty')
-            ann = {}
-            ann['id'] = str(uuid.uuid1())
-            ann['image_id'] = im['id']
-            ann['category_id'] = category_id
-            ann['sequence_level_annotation'] = False
-            annotations.append(ann)
-            
+        image_results = list(tqdm(pool.imap(
+            partial(_process_labelme_file,
+                input_folder=input_folder,
+                use_folders_as_labels=use_folders_as_labels,
+                no_json_handling=no_json_handling,
+                validate_image_sizes=validate_image_sizes,
+                category_name_to_id=category_name_to_id,
+                allow_new_categories=False
+                ),image_filenames_relative), total=len(image_filenames_relative)))
+        
+    images = []
+    annotations = []
+    
+    # Flatten the lists of images and annotations
+    for result in image_results:
+        im = result['im']
+        annotations_this_image = result['annotations_this_image']
+        
+        if im is None:
+            assert annotations_this_image is None
         else:
+            images.append(im)
+            annotations.extend(annotations_this_image)
             
-            for shape in shapes:
-                if shape['shape_type'] != 'rectangle':
-                    print('Only rectangles are supported, skipping an annotation of type {} in {}'.format(
-                        shape['shape_type'],image_fn_relative))
-                    continue
-                
-                if use_folders_as_labels:
-                    category_name = os.path.basename(os.path.dirname(image_fn_abs))
-                else:
-                    category_name = shape['label']                
-                
-                category_id = add_category(category_name)
-            
-                points = shape['points']
-                assert len(points) == 2, 'Illegal rectangle with {} points'.format(
-                    len(points))
-                
-                p0 = points[0]
-                p1 = points[1]
-                x0 = min(p0[0],p1[0])
-                x1 = max(p0[0],p1[0])
-                y0 = min(p0[1],p1[1])
-                y1 = max(p0[1],p1[1])
-                bbox = [x0,y0,abs(x1-x0),abs(y1-y0)]
-                ann = {}
-                ann['id'] = str(uuid.uuid1())
-                ann['image_id'] = im['id']
-                ann['category_id'] = category_id
-                ann['sequence_level_annotation'] = False
-                ann['bbox'] = bbox
-                annotations.append(ann)
-                
-            # ...for each shape
-            
-        images.append(im)
-                  
-    # ..for each image                
-    
     output_dict = {}
     output_dict['images'] = images
     output_dict['annotations'] = annotations