demo: Implement colormap of label

Author: synml

demo.py

@@ -1,5 +1,6 @@
 import os
 
+import matplotlib.colors
 import matplotlib.pyplot as plt
 import torch.nn.functional as F
 import torch.utils.data
@@ -26,6 +27,9 @@
         image_name = image_path.replace('\\', '/').split('/')[-1]
         image_names.append(image_name)
 
+    # label colormap 설정
+    cmap = matplotlib.colors.ListedColormap(utils.utils.get_cityscapes_label_colormap(short=True))
+
     # 예측 결과 저장
     step = 0
     result_dir = os.path.join('demo', model_name.lower())
@@ -48,6 +52,6 @@
         # 1 배치단위 처리
         assert masks.shape[0] == masks_pred.shape[0]
         for i in range(masks.shape[0]):
-            plt.imsave(os.path.join(result_dir, image_names[step]), masks_pred[i].cpu())
-            plt.imsave(os.path.join(groundtruth_dir, image_names[step]), masks[i])
+            plt.imsave(os.path.join(result_dir, image_names[step]), masks_pred[i].cpu(), cmap=cmap)
+            plt.imsave(os.path.join(groundtruth_dir, image_names[step]), masks[i], cmap=cmap)
             step += 1

utils/utils.py

@@ -1,7 +1,6 @@
 import configparser
 import os
 
-import matplotlib.colors
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
@@ -52,8 +51,8 @@ def get_model(model_name: str, num_classes: int, pretrained: str = None) -> torc
 
 # Cityscapes 데이터셋 라벨 색상 불러오기
 def get_cityscapes_label_colormap(short=False):
-    colormap = np.zeros((20, 3), dtype=np.uint8)
     if not short:
+        colormap = np.zeros((20, 3), dtype=np.uint8)
         colormap[0] = [0, 0, 0]
         colormap[1] = [128, 64, 128]
         colormap[2] = [244, 35, 232]
@@ -75,6 +74,7 @@ def get_cityscapes_label_colormap(short=False):
         colormap[18] = [0, 0, 230]
         colormap[19] = [119, 11, 32]
     else:
+        colormap = np.zeros((8, 3), dtype=np.uint8)
         colormap[0] = [0, 0, 0]
         colormap[1] = [128, 64, 128]
         colormap[2] = [70, 70, 70]
@@ -84,7 +84,7 @@ def get_cityscapes_label_colormap(short=False):
         colormap[6] = [220, 20, 60]
         colormap[7] = [0, 0, 142]
 
-    matplotlib.colors.ListedColormap(colormap)
+    return np.divide(colormap, 255).tolist()
 
 
 # Cityscapes 데이터셋 설정