You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
99 lines
3.1 KiB
99 lines
3.1 KiB
# -*- coding: utf-8 -*- |
|
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved |
|
|
|
import time |
|
import functools |
|
import multiprocessing as mp |
|
import numpy as np |
|
import os |
|
from lvis import LVIS |
|
from pycocotools import mask as maskUtils |
|
|
|
|
|
def annToRLE(ann, img_size): |
|
h, w = img_size |
|
segm = ann['segmentation'] |
|
if type(segm) == list: |
|
# polygon -- a single object might consist of multiple parts |
|
# we merge all parts into one mask rle code |
|
rles = maskUtils.frPyObjects(segm, h, w) |
|
rle = maskUtils.merge(rles) |
|
elif type(segm['counts']) == list: |
|
# uncompressed RLE |
|
rle = maskUtils.frPyObjects(segm, h, w) |
|
else: |
|
# rle |
|
rle = ann['segmentation'] |
|
return rle |
|
|
|
|
|
def annToMask(ann, img_size): |
|
rle = annToRLE(ann, img_size) |
|
m = maskUtils.decode(rle) |
|
return m |
|
|
|
|
|
def _process_instance_to_semantic(anns, output_semantic, img): |
|
img_size = (img["height"], img["width"]) |
|
output = np.zeros(img_size, dtype=np.uint8) |
|
for ann in anns: |
|
mask = annToMask(ann, img_size) |
|
output[mask == 1] = ann["category_id"] // 5 |
|
# save as compressed npz |
|
np.savez_compressed(output_semantic, mask=output) |
|
# Image.fromarray(output).save(output_semantic) |
|
|
|
|
|
def create_lvis_semantic_from_instance(instance_json, sem_seg_root): |
|
""" |
|
Create semantic segmentation annotations from panoptic segmentation |
|
annotations, to be used by PanopticFPN. |
|
|
|
It maps all thing categories to contiguous ids starting from 1, and maps all unlabeled pixels to class 0 |
|
|
|
Args: |
|
instance_json (str): path to the instance json file, in COCO's format. |
|
sem_seg_root (str): a directory to output semantic annotation files |
|
""" |
|
os.makedirs(sem_seg_root, exist_ok=True) |
|
|
|
lvis_detection = LVIS(instance_json) |
|
|
|
def iter_annotations(): |
|
for img_id in lvis_detection.get_img_ids(): |
|
anns_ids = lvis_detection.get_ann_ids([img_id]) |
|
anns = lvis_detection.load_anns(anns_ids) |
|
img = lvis_detection.load_imgs([img_id])[0] |
|
file_name = os.path.splitext(img["file_name"])[0] |
|
output = os.path.join(sem_seg_root, file_name + '.npz') |
|
yield anns, output, img |
|
|
|
# # single process |
|
# print("Start writing to {} ...".format(sem_seg_root)) |
|
# start = time.time() |
|
# for anno, oup, img in iter_annotations(): |
|
# _process_instance_to_semantic( |
|
# anno, oup, img) |
|
# print("Finished. time: {:.2f}s".format(time.time() - start)) |
|
# return |
|
|
|
pool = mp.Pool(processes=max(mp.cpu_count() // 2, 4)) |
|
|
|
print("Start writing to {} ...".format(sem_seg_root)) |
|
start = time.time() |
|
pool.starmap( |
|
functools.partial( |
|
_process_instance_to_semantic), |
|
iter_annotations(), |
|
chunksize=100, |
|
) |
|
print("Finished. time: {:.2f}s".format(time.time() - start)) |
|
|
|
|
|
if __name__ == "__main__": |
|
dataset_dir = os.path.join(os.path.dirname(__file__), "lvis") |
|
for s in ["train"]: |
|
create_lvis_semantic_from_instance( |
|
os.path.join(dataset_dir, "lvis_v0.5_{}.json".format(s)), |
|
os.path.join(dataset_dir, "thing_{}".format(s)), |
|
)
|
|
|