File size: 6,188 Bytes
4258b21 fb39cb7 4258b21 43ab69e 4258b21 43ab69e 4258b21 fb39cb7 4258b21 fb39cb7 4258b21 fb39cb7 c0aa824 fb39cb7 c0aa824 fb39cb7 4258b21 |
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 |
import os
import cv2
import numpy as np
import glob
from detectron2.utils.visualizer import ColorMode, Visualizer
from detectron2.data import MetadataCatalog
# Define the ADE20K_150_CATEGORIES with the color mapping
# ADE20K_150_CATEGORIES = [
# {"color": [128, 64, 128], "id": 0, "isthing": 0, "name": "road"},
# {"color": [244, 35, 232], "id": 1, "isthing": 0, "name": "sidewalk"},
# {"color": [70, 70, 70], "id": 2, "isthing": 0, "name": "building"},
# {"color": [102, 102, 156], "id": 3, "isthing": 0, "name": "wall"},
# {"color": [190, 153, 153], "id": 4, "isthing": 0, "name": "fence"},
# {"color": [153, 153, 153], "id": 5, "isthing": 0, "name": "pole"},
# {"color": [250, 170, 30], "id": 6, "isthing": 0, "name": "traffic light"},
# {"color": [220, 220, 0], "id": 7, "isthing": 1, "name": "traffic sign"},
# {"color": [107, 142, 35], "id": 8, "isthing": 1, "name": "vegetation"},
# {"color": [152, 251, 152], "id": 9, "isthing": 0, "name": "terrain"},
# {"color": [70, 130, 180], "id": 10, "isthing": 1, "name": "sky"},
# {"color": [220, 20, 60], "id": 11, "isthing": 0, "name": "person"},
# {"color": [255, 0, 0], "id": 12, "isthing": 1, "name": "rider"},
# {"color": [0, 0, 142], "id": 13, "isthing": 0, "name": "car"},
# {"color": [0, 0, 70], "id": 14, "isthing": 1, "name": "truck"},
# {"color": [0, 60, 100], "id": 15, "isthing": 1, "name": "bus"},
# {"color": [0, 80, 100], "id": 16, "isthing": 0, "name": "train"},
# {"color": [0, 0, 230], "id": 17, "isthing": 0, "name": "motorcycle"},
# {"color": [119, 11, 32], "id": 18, "isthing": 1, "name": "bicycle"}
# ]
ADE20K_150_CATEGORIES = [
{"color": [250, 250, 250], "id": 0, "isthing": 0, "name": "background"},
{"color": [194, 255, 0], "id": 1, "isthing": 0, "name": "aeroplane"},
{"color": [6, 230, 230], "id": 2, "isthing": 0, "name": "bicycle"},
{"color": [80, 150, 20], "id": 3, "isthing": 0, "name": "bird"},
{"color": [4, 200, 3], "id": 4, "isthing": 0, "name": "boat"},
{"color": [255, 5, 153], "id": 5, "isthing": 0, "name": "bottle"},
{"color": [0, 255, 245], "id": 6, "isthing": 0, "name": "bus"},
{"color": [204, 5, 255], "id": 7, "isthing": 1, "name": "car"},
{"color": [255, 0, 102], "id": 8, "isthing": 1, "name": "cat "},
{"color": [4, 250, 7], "id": 9, "isthing": 0, "name": "chair"},
{"color": [224, 5, 255], "id": 10, "isthing": 1, "name": "cow"},
{"color": [235, 255, 7], "id": 11, "isthing": 0, "name": "diningtable"},
{"color": [150, 5, 61], "id": 12, "isthing": 1, "name": "dog"},
{"color": [0, 255, 112], "id": 13, "isthing": 0, "name": "horse"},
{"color": [8, 255, 51], "id": 14, "isthing": 1, "name": "motorbike"},
{"color": [255, 6, 82], "id": 15, "isthing": 1, "name": "person"},
{"color": [143, 255, 140], "id": 16, "isthing": 0, "name": "pottedplant"},
{"color": [204, 255, 4], "id": 17, "isthing": 0, "name": "sheep"},
{"color": [255, 51, 7], "id": 18, "isthing": 1, "name": "sofa"},
{"color": [204, 70, 3], "id": 19, "isthing": 1, "name": "train"},
{"color": [0, 102, 200], "id": 20, "isthing": 1, "name": "tvmonitor"},
]
# Create a dictionary mapping from class ID to color
id_to_color = {category["id"]: category["color"] for category in ADE20K_150_CATEGORIES}
def _get_ade20k_full_meta():
# Id 0 is reserved for ignore_label, we change ignore_label for 0
# to 255 in our pre-processing, so all ids are shifted by 1.
stuff_ids = [k["id"] for k in ADE20K_150_CATEGORIES]
# assert len(stuff_ids) == 847, 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"].split(",")[0] for k in ADE20K_150_CATEGORIES]
color = [k["color"] for k in ADE20K_150_CATEGORIES]
ret = {
"stuff_dataset_id_to_contiguous_id": stuff_dataset_id_to_contiguous_id,
"stuff_classes": stuff_classes,
"color":color
}
return ret
def register_all_ade20k_full(root):
root = os.path.join(root, "ADE20K_2021_17_01")
meta = _get_ade20k_full_meta()
for name, dirname in [("val", "validation")]:
image_dir = os.path.join(root, "images_detectron2", dirname)
gt_dir = os.path.join(root, "annotations_detectron2", dirname)
name = f"sem_seg"
MetadataCatalog.get(name).set(
stuff_classes=meta["stuff_classes"][:],
stuff_colors = meta["color"][:],
image_root=image_dir,
sem_seg_root=gt_dir,
evaluator_type="sem_seg",
ignore_label=65535, # NOTE: gt is saved in 16-bit TIFF images
)
def visualize_segmentation(image_path, mask_path, vis_file):
_root = os.getenv("DETECTRON2_DATASETS", "datasets")
register_all_ade20k_full(_root)
mask = cv2.imread(mask_path)[:,:,0]
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
h,w = image.shape[:2]
metadata = MetadataCatalog.get("sem_seg")
visualizer = Visualizer(image, metadata, instance_mode=ColorMode.IMAGE)
vis_output = visualizer.draw_sem_seg(
mask, alpha=0.5
)
vis_image = vis_output.get_image()[:, :, ::-1]
vis_output.save(vis_file)
return vis_image
# # Set paths
# image_path = "/content/drive/MyDrive/DatasetDM/DataDiffusion/CityScapes/Image/"
# mask_path = "/content/drive/MyDrive/DatasetDM/DataDiffusion/CityScapes/Mask/"
# save_path = 'Results/Vis'
# os.makedirs(save_path, exist_ok=True)
# # List image and mask files
# image_files = glob.glob(os.path.join(image_path, '*.jpg'))
# mask_files = glob.glob(os.path.join(mask_path, '*.png'))
# # Visualize and save segmentation maps for each pair of image and mask files
# for image_file, mask_file in zip(image_files, mask_files):
# if not os.path.isfile(mask_file):
# continue
# save_file = os.path.join(save_path, os.path.basename(image_file))
# visualize_segmentation(image_file, mask_file, save_file)
|