update

seg2art/inference_util.py

@@ -33,15 +33,15 @@ colors = np.array(colors)
 
 
 def remap_label(arr):
-    # compare only first 1 channel to speed up
-    arr_r = arr[:, :, 0]
+    # compare only last color channel to speed up
+    arr_b = arr[:, :, 0]
 
     # remap color to label
     for i in range(len(colors)):
-        arr_r[arr_r == colors[i][0]] = values[i]
+        arr_b[arr_b == colors[i][2]] = values[i]
     # others to 15
-    arr_r[arr_r > 15] = 15
-    return arr_r
+    arr_b[arr_b > 15] = 15
+    return arr_b
 
 
 preprocess = transforms.Compose(