如何在Detectron2中为同一类别的Bounding Box分配相同颜色？

方法一：通过Metadata指定类别固定颜色 #

Detectron2的Visualizer会优先使用Metadata中定义的类别颜色，你可以直接给数据集Metadata配置每个类别的RGB颜色：

from detectron2.data import MetadataCatalog

# 替换成你的数据集名称和类别列表
train_metadata = MetadataCatalog.get("your_dataset_name")
# 按类别顺序设置对应RGB颜色，比如猫(蓝)、狗(红)
train_metadata.set(thing_classes=["cat", "dog"], thing_colors=[(0, 0, 255), (255, 0, 0)])

之后保持你原有的可视化代码不变，同一类别的Bounding Box就会自动使用对应颜色绘制。

方法二：自定义颜色生成器 #

如果不想修改Metadata，可以重写颜色生成逻辑，让同一类别实例使用相同颜色：

from detectron2.utils.visualizer import Visualizer, _ColorGenerator
import glob
import random
import cv2

# 自定义颜色生成器，按类别索引返回固定颜色
class FixedColorGenerator(_ColorGenerator):
    def __init__(self, class_colors):
        self.class_colors = class_colors
    def __call__(self, instance):
        # 根据实例的类别ID取对应颜色
        return self.class_colors[instance.class_id]

# 按类别顺序定义RGB颜色，索引对应类别ID
class_colors = [(0, 0, 255), (255, 0, 0)]  # 猫、狗对应颜色

# 原可视化逻辑修改部分
for imageName in random.sample(glob.glob(os.path.join(test_path, '*.png')), 3):
    im = cv2.imread(imageName)
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1],
                   metadata=train_metadata,
                   scale=0.8
                   )
    # 替换Visualizer的颜色生成器
    v.color_generator = FixedColorGenerator(class_colors)
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    cv2_imshow(out.get_image()[:, :, ::-1])

关键说明 #

• 颜色采用RGB格式，因为代码中已将OpenCV的BGR图像转为RGB供Visualizer处理。
• 若有更多类别，只需按类别顺序在颜色列表中添加对应RGB值即可。

内容的提问来源于stack exchange，提问作者Seunghyeon