|
|
|
|
|
|
|
|
|
import functools |
|
import json |
|
import multiprocessing as mp |
|
import numpy as np |
|
import os |
|
import time |
|
from fvcore.common.download import download |
|
from panopticapi.utils import rgb2id |
|
from PIL import Image |
|
|
|
from detectron2.data.datasets.builtin_meta import COCO_CATEGORIES |
|
|
|
|
|
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) + 255 |
|
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 = {} |
|
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) |
|
|
|
pool = mp.Pool(processes=max(mp.cpu_count() // 2, 4)) |
|
|
|
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 |
|
|
|
print("Start writing to {} ...".format(sem_seg_root)) |
|
start = time.time() |
|
pool.starmap( |
|
functools.partial(_process_panoptic_to_semantic, id_map=id_map), |
|
iter_annotations(), |
|
chunksize=100, |
|
) |
|
print("Finished. time: {:.2f}s".format(time.time() - start)) |
|
|
|
|
|
if __name__ == "__main__": |
|
dataset_dir = os.path.join(os.getenv("DETECTRON2_DATASETS", "datasets"), "coco") |
|
for s in ["val2017", "train2017"]: |
|
separate_coco_semantic_from_panoptic( |
|
os.path.join(dataset_dir, "annotations/panoptic_{}.json".format(s)), |
|
os.path.join(dataset_dir, "panoptic_{}".format(s)), |
|
os.path.join(dataset_dir, "panoptic_stuff_{}".format(s)), |
|
COCO_CATEGORIES, |
|
) |
|
|
|
|
|
|
|
dest_dir = os.path.join(dataset_dir, "annotations/") |
|
URL_PREFIX = "https://dl.fbaipublicfiles.com/detectron2/" |
|
download(URL_PREFIX + "annotations/coco/panoptic_val2017_100.json", dest_dir) |
|
with open(os.path.join(dest_dir, "panoptic_val2017_100.json")) as f: |
|
obj = json.load(f) |
|
|
|
def link_val100(dir_full, dir_100): |
|
print("Creating " + dir_100 + " ...") |
|
os.makedirs(dir_100, exist_ok=True) |
|
for img in obj["images"]: |
|
basename = os.path.splitext(img["file_name"])[0] |
|
src = os.path.join(dir_full, basename + ".png") |
|
dst = os.path.join(dir_100, basename + ".png") |
|
src = os.path.relpath(src, start=dir_100) |
|
os.symlink(src, dst) |
|
|
|
link_val100( |
|
os.path.join(dataset_dir, "panoptic_val2017"), |
|
os.path.join(dataset_dir, "panoptic_val2017_100"), |
|
) |
|
|
|
link_val100( |
|
os.path.join(dataset_dir, "panoptic_stuff_val2017"), |
|
os.path.join(dataset_dir, "panoptic_stuff_val2017_100"), |
|
) |
|
|
