| 4 | import torch |
| 5 | |
| 6 | class SegHelper: |
| 7 | def __init__(self,opt=None,idx2color_path='../../backup/idx2color.txt',num_class=32): |
| 8 | self.opt = opt |
| 9 | self.num_classes = num_class |
| 10 | self.idx2color_path = idx2color_path |
| 11 | f = open(self.idx2color_path, 'r') |
| 12 | self.idx2color = {k:[] for k in range(self.num_classes)} |
| 13 | for j in range(256): |
| 14 | line = f.readline() |
| 15 | line = line.strip(' \n').strip('[').strip(']').strip(' ').split() |
| 16 | line = [int(l) for l in line if l.isdigit()] |
| 17 | self.idx2color[j] = line # color in rgb order |
| 18 | |
| 19 | self.color2idx = {tuple(v):k for k,v in self.idx2color.items()} |
| 20 | name2idx = Dict_indexing() |
| 21 | self.name2idx = {k: name2idx[k]['index'] for k in name2idx.keys()} |
| 22 | self.idx2name = {v:k for k,v in self.name2idx.items()} |
| 23 | self.idx2name_padding = {v:'BG' for v in range(self.num_classes,256)} |
| 24 | self.idx2name.update(self.idx2name_padding) |
| 25 | |
| 26 | def unique(self,array): |
| 27 | uniq, index = np.unique(array, return_index=True, axis=0) |
| 28 | return uniq[index.argsort()] |
| 29 | |
| 30 | def extract_color_from_seg(self,img_seg): |
| 31 | colors = img_seg.reshape(-1, img_seg.shape[-1]) # (H*W,3) # color channel in rgb order |
| 32 | unique_colors = self.unique(colors) # (num_class_in_img,3) |
| 33 | return unique_colors |
| 34 | |
| 35 | def extract_class_from_seg(self,img_seg): |
| 36 | unique_colors = self.extract_color_from_seg(img_seg) # (num_class_in_img,3) # color channel in rgb order |
| 37 | classes_idx = [self.color2idx[tuple(color.tolist())]for color in unique_colors] |
| 38 | classes_str = [self.idx2name[idx] for idx in classes_idx] |
| 39 | return classes_idx, classes_str |
| 40 | |
| 41 | def colormap2classmap(self,seg_array): |
| 42 | seg_array_flattened = torch.LongTensor(seg_array.reshape(-1,3)).cuda() |
| 43 | seg_map_class_flattened = torch.zeros((seg_array.shape[0],seg_array.shape[1],1)).view(-1,1).cuda() |
| 44 | for color, cls in self.color2idx.items(): |
| 45 | matching_indices = (seg_array_flattened == torch.LongTensor(color).cuda()) |
| 46 | matching_indices = (matching_indices.sum(dim=1)==3) |
| 47 | seg_map_class_flattened[matching_indices] = cls |
| 48 | seg_map_class = seg_map_class_flattened.view(seg_array.shape[0],seg_array.shape[1],1) |
| 49 | return seg_map_class |
| 50 | |
| 51 | def classmap2colormap(self,seg_map_class): |
| 52 | seg_map_class_flattened = seg_map_class.view(-1,1) |
| 53 | seg_map_color_flattened = torch.zeros(seg_map_class.shape[0]*seg_map_class.shape[1],3).cuda().long() |
| 54 | for cls, color in self.idx2color.items(): |
| 55 | matching_indices = (seg_map_class_flattened == torch.LongTensor([cls]).cuda()) |
| 56 | seg_map_color_flattened[matching_indices.view(-1)] = torch.LongTensor(color).cuda() |
| 57 | seg_map_color_flattened = seg_map_color_flattened.view(seg_map_class.shape[0],seg_map_class.shape[1],3) |
| 58 | return seg_map_color_flattened |
| 59 | |
| 60 | def split_SemAndChange(self,seg_map_class): |
| 61 | seg_map_change_class = seg_map_class//50 |
| 62 | seg_map_semantic_class = torch.fmod(seg_map_class,50) |
| 63 | return seg_map_semantic_class, seg_map_change_class |
no outgoing calls
no test coverage detected