"""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