add mask preprocessing to coco utils

Author: Matt Sokoloff

examples/integrations/detectron2/coco_utils.py

@@ -59,20 +59,76 @@ def visualize_coco_examples(coco_examples, metadata_catalog, scale = 1.0, max_im
     return Image.fromarray(np.vstack(images))
 
 
-def visualize_object_inferences(coco_examples, metadata_catalog, predictor, scale = 1.0, max_images = 5, resize_dims = (512, 768)):
-    images = []
-    for idx, example in enumerate(coco_examples):
-        if idx > max_images:
-            break
-        im = cv2.imread(example['file_name'])
-        outputs = predictor(im)
-        v = Visualizer(im[:, :, ::-1], metadata_catalog, scale=scale)
-        out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
-        images.append(cv2.resize(out.get_image()[:, :, ::-1], resize_dims))
-    Image.fromarray(np.vstack(images))
-
-
-
-def visualize_panoptic_inferences():
-    ...
+ def visualize_object_inferences(coco_examples, metadata_catalog, predictor, scale = 1.0, max_images = 5, resize_dims = (768, 512)):
+     images = []
+     for idx, example in enumerate(coco_examples):
+         if idx > max_images:
+             break
+         im = cv2.imread(example['file_name'])
+         outputs = predictor(im)
+         v = Visualizer(im[:, :, ::-1], metadata_catalog, scale=scale)
+         out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+         images.append(cv2.resize(out.get_image()[:, :, ::-1], resize_dims))
+     return Image.fromarray(np.vstack(images))
+
+
+
+ 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)
+
+    output = np.zeros_like(panoptic, dtype=np.uint8)
+    for seg in segments:
+        cat_id = seg["category_id"]
+        new_cat_id = id_map[cat_id]
+        output[panoptic == seg["id"]] = new_cat_id
+    Image.fromarray(output).save(output_semantic)
+
+
+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.
+
+    It maps all thing categories to class 0, and maps all unlabeled pixels to class 255.
+    It maps all stuff categories to contiguous ids starting from 1.
+
+    Args:
+        panoptic_json (str): path to the panoptic json file, in COCO's format.
+        panoptic_root (str): a directory with panoptic annotation files, in COCO's format.
+        sem_seg_root (str): a directory to output semantic annotation files
+        categories (list[dict]): category metadata. Each dict needs to have:
+            "id": corresponds to the "category_id" in the json annotations
+            "isthing": 0 or 1
+    """
+    os.makedirs(sem_seg_root, exist_ok=True)
+
+    stuff_ids = [k["id"] for k in categories if k["isthing"] == 0]
+    thing_ids = [k["id"] for k in categories if k["isthing"] == 1]
+    id_map = {}  # map from category id to id in the output semantic annotation
+    assert len(stuff_ids) <= 254
+    for i, stuff_id in enumerate(stuff_ids):
+        id_map[stuff_id] = i + 1
+    for thing_id in thing_ids:
+        id_map[thing_id] = 0
+    id_map[0] = 255
+
+    with open(panoptic_json) as f:
+        obj = json.load(f)
+
+    def iter_annotations():
+        for anno in obj["annotations"]:
+            file_name = anno["file_name"]
+            segments = anno["segments_info"]
+            input = os.path.join(panoptic_root, file_name)
+            output = os.path.join(sem_seg_root, file_name)
+            yield input, output, segments
+
+    fn = functools.partial(_process_panoptic_to_semantic, id_map=id_map)
+    futures = []
+    with ThreadPoolExecutor(max_workers=12) as executor:
+        for args in iter_annotations():
+            futures.append(executor.submit(fn, *args))
+        for _ in tqdm(as_completed(futures)):
+            _.result()