Spaces:
Running
Running
| # ref: https://github.com/MendelXu/SAN/blob/main/san/data/datasets/register_coco_stuff_164k.py | |
| import os | |
| from detectron2.data import DatasetCatalog, MetadataCatalog | |
| from detectron2.data.datasets import load_sem_seg | |
| from . import openseg_classes | |
| COCO_CATEGORIES = openseg_classes.get_coco_stuff_categories_with_prompt_eng() | |
| def _get_coco_stuff_meta(): | |
| # Id 0 is reserved for ignore_label, we change ignore_label for 0 | |
| # to 255 in our pre-processing. | |
| stuff_ids = [k["id"] for k in COCO_CATEGORIES] | |
| assert len(stuff_ids) == 171, len(stuff_ids) | |
| # For semantic segmentation, this mapping maps from contiguous stuff id | |
| # (in [0, 91], used in models) to ids in the dataset (used for processing results) | |
| stuff_dataset_id_to_contiguous_id = {k: i for i, k in enumerate(stuff_ids)} | |
| stuff_classes = [k["name"] for k in COCO_CATEGORIES] | |
| ret = { | |
| "stuff_dataset_id_to_contiguous_id": stuff_dataset_id_to_contiguous_id, | |
| "stuff_classes": stuff_classes, | |
| } | |
| return ret | |
| def register_all_coco_stuff_164k(root): | |
| root = os.path.join(root, "coco") | |
| meta = _get_coco_stuff_meta() | |
| for name, image_dirname, sem_seg_dirname in [ | |
| ("train", "train2017", "stuffthingmaps_detectron2/train2017"), | |
| ("test", "val2017", "stuffthingmaps_detectron2/val2017"), | |
| ]: | |
| image_dir = os.path.join(root, image_dirname) | |
| gt_dir = os.path.join(root, sem_seg_dirname) | |
| all_name = f"openvocab_coco_2017_{name}_stuff_sem_seg" | |
| DatasetCatalog.register( | |
| all_name, | |
| lambda x=image_dir, y=gt_dir: load_sem_seg( | |
| y, x, gt_ext="png", image_ext="jpg" | |
| ), | |
| ) | |
| MetadataCatalog.get(all_name).set( | |
| image_root=image_dir, | |
| sem_seg_root=gt_dir, | |
| evaluator_type="sem_seg", | |
| ignore_label=255, | |
| **meta, | |
| ) | |
| _root = os.getenv("DETECTRON2_DATASETS", "datasets") | |
| register_all_coco_stuff_164k(_root) |