File size: 7,122 Bytes
91c1e9e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | """Image-source and mask utilities required by M4SEG."""
import os
import json
import gzip
from pathlib import Path
import numpy as np
from dataset.utils.coco import COCO, maskUtils
def load_json_file(file_path):
opener = gzip.open if str(file_path).endswith('.gz') else open
with opener(file_path, "rt", encoding='utf-8') as file:
return json.load(file)
def _write_list(items):
if len(items) == 0:
return ""
if len(items) == 1:
return items[0]
if len(items) == 2:
return " and ".join(items)
return ", ".join(items[:-1]) + f", and {items[-1]}"
def source_metadata_path(dataset_dir, source, mode, original_path):
"""Use original annotation tables when present, otherwise released lookup catalogs."""
original_path = Path(original_path)
if original_path.is_file():
return original_path
root = Path(dataset_dir)
if source in {'ade20k_part234', 'paco_lvis', 'pascal_part'}:
root = root.parent
split = 'test' if mode == 'val' and source != 'paco_lvis' else mode
for metadata in (root / 'm4seg' / 'metadata', root / 'metadata'):
path = metadata / 'source_catalogs' / source / (split + '.json.gz')
if path.is_file():
return path
raise FileNotFoundError('Source metadata is missing. Download the M4Seg metadata directory '
'under dataset_dir/m4seg, or supply the original annotation table.')
def init_ade20k_part234(dataset_dir, mode, area_threshold=None):
assert mode in ["train", "val"], f"Invalid mode for ADE20K Part 234: {mode}"
ade20k_part234_api = COCO(
source_metadata_path(dataset_dir, 'ade20k_part234', mode, os.path.join(
dataset_dir, "ADE20KPart234", f"ade20k_instance_{mode}_mixed_parts.json"
)),
area_threshold,
)
all_obj_classes = ade20k_part234_api.loadCats(ade20k_part234_api.getCatIds())
class_map_ade20k_part234 = {"object": {}, "part": {}}
for cat in all_obj_classes:
name = cat["name"]
class_map_ade20k_part234["object"][cat["id"]] = name
part_class_path = Path(__file__).with_name('ade20k_part234_classes_mixed_parts.json')
root = Path(dataset_dir).parent
for metadata in (root / 'm4seg' / 'metadata', root / 'metadata'):
candidate = metadata / 'ade20k_part234_classes.json'
if candidate.is_file():
part_class_path = candidate
break
all_part_classes = load_json_file(part_class_path)
for idx, cat in enumerate(all_part_classes):
obj, part = cat.split("'s ")
class_map_ade20k_part234["part"][idx] = (obj, part)
img_ids = ade20k_part234_api.getImgIds()
return class_map_ade20k_part234, img_ids, ade20k_part234_api
def init_paco_lvis(dataset_dir, mode, area_threshold=None):
assert mode in ["train", "val", "test"], f"Invalid mode for PACO LVIS: {mode}"
paco_lvis_api = COCO(
source_metadata_path(dataset_dir, 'paco_lvis', mode, os.path.join(
dataset_dir, "paco_lvis", "annotations", f"paco_lvis_v1_{mode}.json"
)),
area_threshold,
)
all_classes = paco_lvis_api.loadCats(paco_lvis_api.getCatIds())
class_map_paco_lvis = {"object": {}, "part": {}}
for cat in all_classes:
cat_split = cat["name"].strip().split(":")
if len(cat_split) == 1:
name = cat_split[0].split("_(")[0]
name = " ".join(name.split("_"))
class_map_paco_lvis["object"][cat["id"]] = name
else:
assert len(cat_split) == 2
obj, part = cat_split
obj = obj.split("_(")[0]
part = part.split("_(")[0]
obj = " ".join(obj.split("_"))
part = " ".join(part.split("_"))
name = (obj, part)
class_map_paco_lvis["part"][cat["id"]] = name
# class_map_paco_lvis[cat["id"]] = name
img_ids = paco_lvis_api.getImgIds()
return class_map_paco_lvis, img_ids, paco_lvis_api
def init_pascal_part(dataset_dir, mode, area_threshold=None):
assert mode in ["train", "val"], f"Invalid mode for PASCAL-Part: {mode}"
pascal_part_api = COCO(
source_metadata_path(dataset_dir, 'pascal_part', mode,
os.path.join(dataset_dir, "PASCAL-Part", f"{mode}.json")),
area_threshold,
)
all_classes = pascal_part_api.loadCats(pascal_part_api.getCatIds())
class_map_pascal_part = {"object": {}, "part": {}}
for cat in all_classes:
cat_split = cat["name"].strip().split(":")
if len(cat_split) == 1:
name = cat_split[0].split("_(")[0]
name = " ".join(name.split("_"))
class_map_pascal_part["object"][cat["id"]] = name
else:
assert len(cat_split) == 2
obj, part = cat_split
obj = obj.split("_(")[0]
part = part.split("_(")[0]
obj = " ".join(obj.split("_"))
part = " ".join(part.split("_"))
name = (obj, part)
class_map_pascal_part["part"][cat["id"]] = name
img_ids = pascal_part_api.getImgIds()
return class_map_pascal_part, img_ids, pascal_part_api
def init_pascal(dataset_dir, mode, area_threshold=None):
class_map_pascal_part, img_ids, pascal_part_api = init_pascal_part(
os.path.join(dataset_dir, "Part_Segm"), mode, area_threshold
)
return class_map_pascal_part["object"], img_ids, pascal_part_api
def init_coco(dataset_dir, mode, area_threshold=None):
assert mode in ["train", "val"], f"Invalid mode for COCO: {mode}"
coco_api = COCO(
source_metadata_path(dataset_dir, 'coco_objects', mode, os.path.join(
dataset_dir, "coco_2017", "annotations", f"instances_{mode}2017.json"
)),
area_threshold,
)
all_classes = coco_api.loadCats(coco_api.getCatIds())
class_map = {}
for cat in all_classes:
name = cat["name"]
class_map[cat["id"]] = name
img_ids = coco_api.getImgIds()
return class_map, img_ids, coco_api
def init_ade20k_coco(dataset_dir, mode, area_threshold=None):
assert mode in ["train", "val"], f"Invalid mode for ADE20K: {mode}"
ade20k_api = COCO(
source_metadata_path(dataset_dir, 'ade20k_objects', mode,
os.path.join(dataset_dir, "Semantic_Segm", f"ade20k_instances_{mode}.json")),
area_threshold,
)
all_classes = ade20k_api.loadCats(ade20k_api.getCatIds())
class_map = {}
for cat in all_classes:
name = cat["name"]
class_map[cat["id"]] = name
img_ids = ade20k_api.getImgIds()
return class_map, img_ids, ade20k_api
def segm_to_mask(segm, h, w):
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 = segm
m = maskUtils.decode(rle)
return m
|