File size: 4,223 Bytes
3d109ba | 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 | import torch
import numpy as np
import PIL
import torch.nn.functional as F
import torch.nn as nn
from typing import Dict, Iterator, List, Mapping, Optional, Sequence, Tuple, Union
# RGB colors used to visualize each semantic segmentation class.
SEGMENTATION_COLOR_MAP = dict(
TYPE_UNDEFINED=[0, 0, 0],
TYPE_EGO_VEHICLE=[102, 102, 102],
TYPE_CAR=[0, 0, 142],
TYPE_TRUCK=[0, 0, 70],
TYPE_BUS=[0, 60, 100],
TYPE_OTHER_LARGE_VEHICLE=[61, 133, 198],
TYPE_BICYCLE=[119, 11, 32],
TYPE_MOTORCYCLE=[0, 0, 230],
TYPE_TRAILER=[111, 168, 220],
TYPE_PEDESTRIAN=[220, 20, 60],
TYPE_CYCLIST=[255, 0, 0],
TYPE_MOTORCYCLIST=[180, 0, 0],
TYPE_BIRD=[127, 96, 0],
TYPE_GROUND_ANIMAL=[91, 15, 0],
TYPE_CONSTRUCTION_CONE_POLE=[230, 145, 56],
TYPE_POLE=[153, 153, 153],
TYPE_PEDESTRIAN_OBJECT=[234, 153, 153],
TYPE_SIGN=[246, 178, 107],
TYPE_TRAFFIC_LIGHT=[250, 170, 30],
TYPE_BUILDING=[70, 70, 70],
TYPE_ROAD=[128, 64, 128],
TYPE_LANE_MARKER=[234, 209, 220],
TYPE_ROAD_MARKER=[217, 210, 233],
TYPE_SIDEWALK=[244, 35, 232],
TYPE_VEGETATION=[107, 142, 35],
TYPE_SKY=[70, 130, 180],
TYPE_GROUND=[102, 102, 102],
TYPE_DYNAMIC=[102, 102, 102],
TYPE_STATIC=[102, 102, 102],
)
def _generate_color_map(
color_map_dict: Optional[
Mapping[int, Sequence[int]]] = None
) -> np.ndarray:
"""Generates a mapping from segmentation classes (rows) to colors (cols).
Args:
color_map_dict: An optional dict mapping from semantic classes to colors. If
None, the default colors in SEGMENTATION_COLOR_MAP will be used.
Returns:
A np array of shape [max_class_id + 1, 3], where each row encodes the color
for the corresponding class id.
"""
if color_map_dict is None:
color_map_dict = SEGMENTATION_COLOR_MAP
classes = list(color_map_dict.keys())
colors = list(color_map_dict.values())
color_map = np.zeros([#np.amax(classes) + 1
len(classes)
, 3], dtype=np.uint8)
for idx, color in enumerate(colors):
color_map[idx] = color
#color_map[classes] = colors
return color_map
DEFAULT_COLOR_MAP = _generate_color_map()
def get_panoptic_id(semantic_id, instance_id, semantic_interval=1000):
if isinstance(semantic_id, np.ndarray):
semantic_id = torch.from_numpy(semantic_id)
instance_id = torch.from_numpy(instance_id)
elif isinstance(semantic_id, PIL.Image.Image):
semantic_id = torch.from_numpy(np.array(semantic_id))
instance_id = torch.from_numpy(np.array(instance_id))
elif isinstance(semantic_id, torch.Tensor):
pass
else:
raise ValueError("semantic_id type is not supported!")
return semantic_id * semantic_interval + instance_id
def get_panoptic_encoding(semantic_id, instance_id, ):
# 将 semantic-id 和 instance-id 编码成 panoptic one-hot编码
panoptic_id = get_panoptic_id(semantic_id, instance_id)
unique_panoptic_classes = panoptic_id.unique()
num_panoptic_classes = unique_panoptic_classes.shape[0]
# construct id map dict: panoptic_id -> num_class_idx
id_to_idx_dict = {}
for i in range(num_panoptic_classes):
id_to_idx_dict[unique_panoptic_classes[i]] = i
# convert to one-hot encoding
panoptic_encoding = torch.zeros((num_panoptic_classes, ), dtype=torch.float32)
def feat_encode(obj_id, id_to_idx, gt_label_embedding: nn.Embedding = None, output_both=False, only_idx=False):
""" 根据 obj_id 和 id_to_idx_dict 编码成 one-hot """
map_ids = torch.zeros_like(obj_id) #obj_id.clone()
# 将 gt-obj-id 替换成 global-obj-idx ，然后转成 one-hot
for key, value in id_to_idx.items():
map_ids[obj_id == key] = value
# query embedding
if gt_label_embedding is not None:
gt_label = gt_label_embedding(map_ids.flatten().long())
else:
gt_label = None
if output_both:
return map_ids, gt_label
else:
if only_idx:
return map_ids.long().flatten()
else:
return gt_label
|