Spaces:
Runtime error
Runtime error
File size: 2,194 Bytes
ba4c371 |
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 |
"""
This file may have been modified by Bytedance Ltd. and/or its affiliates (“Bytedance's Modifications”).
All Bytedance's Modifications are Copyright (year) Bytedance Ltd. and/or its affiliates.
Reference: 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) |