如何为detectron2内置模型做输入?

How to do input for detectron2 builtinmodel?

我训练了一个模型,现在我想用它来检测图像中的物体。使用 DefaultDetector 仅返回边界框,我需要掩码。我看到你也可以用这个方法进行推理:

model.eval()
with torch.no_grad():
    outputs = model(inputs)

我认为这才是他应该使用的。问题是我不知道如何设置输入,从图像开始。

import torch
import glob
cfg = get_cfg()

cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/"
                                              "mask_rcnn_R_101_FPN_3x.yaml"))
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.SOLVER.IMS_PER_BATCH = 1
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only has one class
cfg.INPUT.FORMAT = "BGR"
#Just run these lines if you have the trained model im memory
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set the testing threshold for this model
#build model
model = build_model(cfg)

DetectionCheckpointer(model).load("output/model_final.pth")

model.eval()#make sure its in eval mode

image = cv2.imread("/kaggle/working/detectron2/images/73-ab1.jpg")
height, width = image.shape[:2]
image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
image = ImageList.from_tensors([image])

with torch.no_grad():
    inputs = image
    outputs = model(inputs)

不过很遗憾,我想我错了,谁能赐教一下?

有关内置模型,请参阅 Model Input Format

基本上,代码中的模型不需要 ImageList 对象,而是 listdict 对象,其中每个 dict 需要提供有关的特定信息一张图片,如上面链接的文档中所述。

因此,您的推理代码需要更正为以下内容。

image = cv2.imread("/kaggle/working/detectron2/images/73-ab1.jpg")
height, width = image.shape[:2]
image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
inputs = [{"image": image, "height": height, "width": width}]

with torch.no_grad():
    outputs = model(inputs)

您也可以在代码中看到这一点 - the forward method of the GeneralizedRCNN class