format

Author: Matt Sokoloff

examples/integrations/detectron2/coco_utils.py

@@ -14,7 +14,9 @@
 
 def get_annotations(images, all_annotations):
     image_lookup = {image['id'] for image in images}
-    return [annot for annot in all_annotations if annot['image_id'] in image_lookup]
+    return [
+        annot for annot in all_annotations if annot['image_id'] in image_lookup
+    ]
 
 
 def partition_indices(total_n, splits):
@@ -27,34 +29,41 @@ def partition_indices(total_n, splits):
 
     splits = np.cumsum(splits)
     for idx in range(len((splits))):
-        start = 0 if idx == 0 else int(total_n * splits[idx-1])
+        start = 0 if idx == 0 else int(total_n * splits[idx - 1])
         end = int(splits[idx] * total_n)
         yield start, end
 
 
-
-def partition_coco(coco_instance_data, coco_panoptic_data = None, splits = None):
+def partition_coco(coco_instance_data, coco_panoptic_data=None, splits=None):
     images = coco_instance_data['images']
-    n_classes = len({category['id'] for category in coco_instance_data['categories']})
+    n_classes = len(
+        {category['id'] for category in coco_instance_data['categories']})
     random.shuffle(images)
     partitions = []
     for start, end in partition_indices(len(images), splits):
-        partition = {'instance' :  dict(
-            categories = coco_instance_data['categories'],
-            images = images[start:end],
-            annotations = get_annotations(images[start:end], coco_instance_data['annotations'])
-        )}
+        partition = {
+            'instance':
+                dict(categories=coco_instance_data['categories'],
+                     images=images[start:end],
+                     annotations=get_annotations(
+                         images[start:end], coco_instance_data['annotations']))
+        }
         if coco_panoptic_data is not None:
             partition['panoptic'] = dict(
-                categories = coco_panoptic_data['categories'],
-                images = images[start:end],
-                annotations = get_annotations(images[start:end], coco_panoptic_data['annotations'])
-            )
+                categories=coco_panoptic_data['categories'],
+                images=images[start:end],
+                annotations=get_annotations(images[start:end],
+                                            coco_panoptic_data['annotations']))
         partitions.append(partition)
     return partitions
 
 
-def visualize_object_inferences(metadata_catalog, coco_examples, predictor, scale = 1.0, max_images = 5, resize_dims = (768, 512)):
+def visualize_object_inferences(metadata_catalog,
+                                coco_examples,
+                                predictor,
+                                scale=1.0,
+                                max_images=5,
+                                resize_dims=(768, 512)):
     images = []
     for idx, example in enumerate(coco_examples):
         if idx > max_images:
@@ -67,9 +76,14 @@ def visualize_object_inferences(metadata_catalog, coco_examples, predictor, scal
     return Image.fromarray(np.vstack(images))
 
 
-def visualize_coco_examples(metadata_catalog, object_examples, panoptic_examples = None, scale = 1.0, max_images = 5, resize_dims = (768,512)):
+def visualize_coco_examples(metadata_catalog,
+                            object_examples,
+                            panoptic_examples=None,
+                            scale=1.0,
+                            max_images=5,
+                            resize_dims=(768, 512)):
     if panoptic_examples is not None:
-        lookup = {d['file_name'] : d for d in panoptic_examples}
+        lookup = {d['file_name']: d for d in panoptic_examples}
 
     images = []
     for idx, example in enumerate(object_examples):
@@ -86,7 +100,8 @@ def visualize_coco_examples(metadata_catalog, object_examples, panoptic_examples
     return Image.fromarray(np.vstack(images))
 
 
-def _process_panoptic_to_semantic(input_panoptic, output_semantic, segments, id_map):
+def _process_panoptic_to_semantic(input_panoptic, output_semantic, segments,
+                                  id_map):
     panoptic = np.asarray(Image.open(input_panoptic), dtype=np.uint32)
     panoptic = rgb2id(panoptic)
 
@@ -98,7 +113,8 @@ def _process_panoptic_to_semantic(input_panoptic, output_semantic, segments, id_
     Image.fromarray(output).save(output_semantic)
 
 
-def separate_coco_semantic_from_panoptic(panoptic_json, panoptic_root, sem_seg_root, categories):
+def separate_coco_semantic_from_panoptic(panoptic_json, panoptic_root,
+                                         sem_seg_root, categories):
     """
     Create semantic segmentation annotations from panoptic segmentation
     annotations, to be used by PanopticFPN.
@@ -144,4 +160,3 @@ def iter_annotations():
             futures.append(executor.submit(fn, *args))
         for _ in tqdm(as_completed(futures)):
             _.result()
-