From 9a3b0349c3fd0fddfdc30370cba88f5633e7db6a Mon Sep 17 00:00:00 2001 From: triplemu Date: Tue, 4 Jun 2024 10:25:25 +0800 Subject: [PATCH] fix seg color bugs --- config.py | 3 +++ infer-seg-without-torch.py | 4 ++-- infer-seg.py | 4 ++-- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/config.py b/config.py index 17ca86c..d14f2aa 100644 --- a/config.py +++ b/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), diff --git a/infer-seg-without-torch.py b/infer-seg-without-torch.py index fa08b62..d568099 100644 --- a/infer-seg-without-torch.py +++ b/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, diff --git a/infer-seg.py b/infer-seg.py index 0284244..9608e05 100644 --- a/infer-seg.py +++ b/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,