在detectron2中有类ID而不是类名

Question

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 
cfg.DATASETS.TEST = ("fruit_test", )
predictor = DefaultPredictor(cfg)

image_path = "/content/detectron2_custom_dataset/testimages/test2.jpg"

def on_image(image_path,predictor):
    im = cv2.imread(image_path)
    outputs = predictor(im)
    v = Visualizer(im[:,:,::-1], metadata = {}, scale=0.5, instance_mode = ColorMode.SEGMENTATION)
    v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    plt.figure(figsize=(14,10))
    plt.imshow(v.get_image())
    plt.show()

on_image(image_path, predictor)

总之，我想用我现在上传的模型测试我的模型，我不想在图像上有类 ID。我想要班级名称，例如橙色，香蕉，苹果

Answer 1

在寻找一个简单的解决方案之后，我想到了这个，这是最简单的 IMO。

创建一个新的元数据类

Class Metadata:
    def get(self, _):
        return ['apple','banana','orange','etc'] #your class labels

然后，在 Visualizer 行中提供元数据

v = Visualizer(im[:, :, ::-1], Metadata, scale=0.5, instance_mode = ColorMode.SEGMENTATION)

否则，您需要使用（虚拟）数据或什至 1 张带有注释的图像注册模型，然后加载其元数据。我觉得在这个阶段有点无关紧要，如果你只需要它用于推理代码。

Answer 2

您可以通过填充包含映射的

metadata

kwarg 来获取标签。

v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TEST[0]), scale=0.5, instance_mode = ColorMode.SEGMENTATION)

Answer 3

我会将 Visualizer 行中的“元数据 = {}”更改为

{"thing_classes":['orange','banana','apple']}

整条线是

v = Visualizer(im[:,:,::-1], {"thing_classes":['orange','banana','apple']}, scale=0.5, instance_mode = ColorMode.SEGMENTATION)

基于 draw_instance_predictions 函数中的以下代码 (detectron2/utils/visualizer.py)

labels = _create_text_labels(classes, scores, self.metadata.get("thing_classes", None))

Answer 4

这条简单的线对我有用。

MetadataCatalog.get("YOUR_DATASET").thing_classes = [LIST OF YOUR CLASS LABELS]