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)