visualize.py

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)