Include clean PRIMA code and evaluator; document metrics, initialization, and MIT contributions
91c1e9e verified | """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 | |