fix seg color bugs

fix seg color bugs

config.py

@@ -162,6 +162,9 @@ COLORS = {
     for i, cls in enumerate(CLASSES_DET)
 }
 
+# segment model classes
+CLASSES_SEG = CLASSES_DET
+
 # colors for segment masks
 MASK_COLORS = np.array([(255, 56, 56), (255, 157, 151), (255, 112, 31),
                         (255, 178, 29), (207, 210, 49), (72, 249, 10),

infer-seg-without-torch.py

@@ -4,7 +4,7 @@ from pathlib import Path
 import cv2
 import numpy as np
 
-from config import ALPHA, CLASSES, COLORS, MASK_COLORS
+from config import ALPHA, CLASSES_SEG, COLORS, MASK_COLORS
 from models.utils import blob, letterbox, path_to_list, seg_postprocess
 
 
@@ -60,7 +60,7 @@ def main(args: argparse.Namespace) -> None:
         for (bbox, score, label) in zip(bboxes, scores, labels):
             bbox = bbox.round().astype(np.int32).tolist()
             cls_id = int(label)
-            cls = CLASSES[cls_id]
+            cls = CLASSES_SEG[cls_id]
             color = COLORS[cls]
             cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2)
             cv2.putText(draw,

infer-seg.py

@@ -6,7 +6,7 @@ import cv2
 import numpy as np
 import torch
 
-from config import ALPHA, CLASSES, COLORS, MASK_COLORS
+from config import ALPHA, CLASSES_SEG, COLORS, MASK_COLORS
 from models.torch_utils import seg_postprocess
 from models.utils import blob, letterbox, path_to_list
 
@@ -63,7 +63,7 @@ def main(args: argparse.Namespace) -> None:
         for (bbox, score, label) in zip(bboxes, scores, labels):
             bbox = bbox.round().int().tolist()
             cls_id = int(label)
-            cls = CLASSES[cls_id]
+            cls = CLASSES_SEG[cls_id]
             color = COLORS[cls]
             cv2.rectangle(draw, bbox[:2], bbox[2:], color, 2)
             cv2.putText(draw,