如何使用 detectorron2 和 DefaultPredictor 对多个图像进行推理

Question

我已经训练了模型，现在我想用它来检测许多图像中的对象。我看到默认预测器允许您仅在图像上进行检测，我该怎么办？

我对这个世界真的很陌生。我尝试的方法是使用 for 循环，但它不起作用。还有其他方法吗？

%cd /kaggle/working/detectron2
import glob
cfg.MODEL.WEIGHTS = os.path.join("/kaggle/working/detectron2/output", "model_final.pth") # path to the model we trained
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.0001 # set a testing threshold
pred = DefaultPredictor(cfg)
os.chdir("/kaggle/working/detectron2/images")
for img in glob.glob('.jpg'):
    inputs = cv2.imread(img)
    outputs = pred(inputs)
    print(outputs)

Answer 1

好的，我是这样解决的：

%cd /kaggle/working/detectron2
import glob
cfg.MODEL.WEIGHTS = os.path.join("/kaggle/working/detectron2/output", "model_final.pth")   # path to the model we trained
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.0001   # set a testing threshold
pred = DefaultPredictor(cfg)
for img in glob.glob('/kaggle/working/detectron2/images/*.jpg'):
    inputs = cv2.imread(img)
    outputs = pred(inputs)
    print(outputs)

我删除了 os.chdir()

Answer 2

引用DefaultPredictor的源码文档：

如果您想做更复杂的事情，请参考其源代码作为示例来手动构建和使用模型。

即，要对多个图像进行推理，请复制并重命名 DefaultPredictor 类（github 链接），然后将 _call 方法替换为：

def __call__(self, original_images):
    with torch.no_grad():  # https://github.com/sphinx-doc/sphinx/issues/4258
        # Apply pre-processing to images.
        inputs = []
        for original_image in original_images:
            if self.input_format == "RGB":
                # whether the model expects BGR inputs or RGB
                original_image = original_image[:, :, ::-1]

            height, width = original_image.shape[:2]
            image = self.aug.get_transform(original_image).apply_image(original_image)
            image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
            image.to(self.cfg.MODEL.DEVICE)
            inputs.append({"image": image, "height": height, "width": width})

        predictions = self.model(inputs)
        return predictions

使用与 DefaultPredictor 相同的方式，但传递图像数组：

predictor = MyPredictor(cfg)
predictions = predictor(your_images)