Dororo99/Ours_S3GS_Waymo

1import torch
2import numpy as np
3import PIL
4import torch.nn.functional as F
5import torch.nn as nn
6from typing import Dict, Iterator, List, Mapping, Optional, Sequence, Tuple, Union
7
8# RGB colors used to visualize each semantic segmentation class.
9SEGMENTATION_COLOR_MAP = dict(
10    TYPE_UNDEFINED=[0, 0, 0],
11    TYPE_EGO_VEHICLE=[102, 102, 102],
12    TYPE_CAR=[0, 0, 142],
13    TYPE_TRUCK=[0, 0, 70],
14    TYPE_BUS=[0, 60, 100],
15    TYPE_OTHER_LARGE_VEHICLE=[61, 133, 198],
16    TYPE_BICYCLE=[119, 11, 32],
17    TYPE_MOTORCYCLE=[0, 0, 230],
18    TYPE_TRAILER=[111, 168, 220],
19    TYPE_PEDESTRIAN=[220, 20, 60],
20    TYPE_CYCLIST=[255, 0, 0],
21    TYPE_MOTORCYCLIST=[180, 0, 0],
22    TYPE_BIRD=[127, 96, 0],
23    TYPE_GROUND_ANIMAL=[91, 15, 0],
24    TYPE_CONSTRUCTION_CONE_POLE=[230, 145, 56],
25    TYPE_POLE=[153, 153, 153],
26    TYPE_PEDESTRIAN_OBJECT=[234, 153, 153],
27    TYPE_SIGN=[246, 178, 107],
28    TYPE_TRAFFIC_LIGHT=[250, 170, 30],
29    TYPE_BUILDING=[70, 70, 70],
30    TYPE_ROAD=[128, 64, 128],
31    TYPE_LANE_MARKER=[234, 209, 220],
32    TYPE_ROAD_MARKER=[217, 210, 233],
33    TYPE_SIDEWALK=[244, 35, 232],
34    TYPE_VEGETATION=[107, 142, 35],
35    TYPE_SKY=[70, 130, 180],
36    TYPE_GROUND=[102, 102, 102],
37    TYPE_DYNAMIC=[102, 102, 102],
38    TYPE_STATIC=[102, 102, 102],
39)
40
41def _generate_color_map(
42    color_map_dict: Optional[
43        Mapping[int, Sequence[int]]] = None
44) -> np.ndarray:
45  """Generates a mapping from segmentation classes (rows) to colors (cols).
46
47  Args:
48    color_map_dict: An optional dict mapping from semantic classes to colors. If
49      None, the default colors in SEGMENTATION_COLOR_MAP will be used.
50  Returns:
51    A np array of shape [max_class_id + 1, 3], where each row encodes the color
52      for the corresponding class id.
53  """
54  if color_map_dict is None:
55    color_map_dict = SEGMENTATION_COLOR_MAP
56  classes = list(color_map_dict.keys())
57  colors = list(color_map_dict.values())
58  color_map = np.zeros([#np.amax(classes) + 1
59                        len(classes)
60                        , 3], dtype=np.uint8)
61  for idx, color in enumerate(colors):
62      color_map[idx] = color
63  #color_map[classes] = colors
64  return color_map
65
66DEFAULT_COLOR_MAP = _generate_color_map()
67
68def get_panoptic_id(semantic_id, instance_id, semantic_interval=1000):
69    if isinstance(semantic_id, np.ndarray):
70        semantic_id = torch.from_numpy(semantic_id)
71        instance_id = torch.from_numpy(instance_id)
72    elif isinstance(semantic_id, PIL.Image.Image):
73        semantic_id = torch.from_numpy(np.array(semantic_id))
74        instance_id = torch.from_numpy(np.array(instance_id))
75    elif isinstance(semantic_id, torch.Tensor):
76        pass
77    else:
78        raise ValueError("semantic_id type is not supported!")
79
80    return semantic_id * semantic_interval + instance_id
81
82def get_panoptic_encoding(semantic_id, instance_id, ):
83    # 将 semantic-id 和 instance-id 编码成 panoptic one-hot编码
84    panoptic_id = get_panoptic_id(semantic_id, instance_id)
85    unique_panoptic_classes = panoptic_id.unique()
86    num_panoptic_classes = unique_panoptic_classes.shape[0]
87    # construct id map dict: panoptic_id -> num_class_idx
88    id_to_idx_dict = {}
89    for i in range(num_panoptic_classes):
90        id_to_idx_dict[unique_panoptic_classes[i]] = i
91
92    # convert to one-hot encoding
93    panoptic_encoding = torch.zeros((num_panoptic_classes, ), dtype=torch.float32)
94
95
96def feat_encode(obj_id, id_to_idx, gt_label_embedding: nn.Embedding = None, output_both=False, only_idx=False):
97    """ 根据 obj_id 和 id_to_idx_dict 编码成 one-hot """
98
99    map_ids = torch.zeros_like(obj_id) #obj_id.clone()
100    # 将 gt-obj-id 替换成 global-obj-idx ，然后转成 one-hot
101    for key, value in id_to_idx.items():
102        map_ids[obj_id == key] = value
103
104    # query embedding
105    if gt_label_embedding is not None:
106        gt_label = gt_label_embedding(map_ids.flatten().long())
107    else:
108        gt_label = None
109
110    if output_both:
111        return map_ids, gt_label
112    else:
113        if only_idx:
114            return map_ids.long().flatten()
115        else:
116            return gt_label
117
118        
119
120