fruits_nuts_metadata = MetadataCatalog.get("fruits_nuts")
if __name__ == "__main__":
cfg = get_cfg()
cfg.merge_from_file(
"/media/dell/新加卷2/downcode/detectron2_train(ubuntu)/configs/COCO-InstanceSegmentation//mask_rcnn_R_50_FPN_3x.yaml"
)
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
print('loading from: {}'.format(cfg.MODEL.WEIGHTS))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set the testing threshold for this model
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3
cfg.DATASETS.TEST = ("fruits_nuts", )
predictor = DefaultPredictor(cfg)
data_f = './data/images/2.jpg'
im = cv2.imread(data_f)
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1],
metadata=fruits_nuts_metadata,
scale=0.8,
instance_mode=ColorMode.IMAGE_BW # remove the colors of unsegmented pixels
)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
img = v.get_image()[:, :, ::-1]
cv2.imshow('rr', img)
cv2.waitKey(0)
v.draw_instance_predictions看这个函数,类别没显示而已,mask已经画好在你这图上了,是有的
那怎么修改代码可以把类别显示出来?