fix patching keypoint_rcnn_inference
Summary: due to move of call file.
@@ -6,7 +6,7 @@ import torch
from detectron2.modeling import poolers
from detectron2.modeling.proposal_generator import rpn
-from detectron2.modeling.roi_heads import mask_head, roi_heads
+from detectron2.modeling.roi_heads import keypoint_head, mask_head
from detectron2.modeling.roi_heads.fast_rcnn import FastRCNNOutputs
from .c10 import (
@@ -125,15 +125,15 @@ class ROIHeadsPatcher:
format or not. Default to True.
"""
# NOTE: this requries the `keypoint_rcnn_inference` and `mask_rcnn_inference`
- # are called inside the same file as ROIHeads due to using mock.patch.
- roi_heads_mod = roi_heads.ROIHeads.__module__
+ # are called inside the same file as BaseXxxHead due to using mock.patch.
+ kpt_heads_mod = keypoint_head.BaseKeypointRCNNHead.__module__
mask_head_mod = mask_head.BaseMaskRCNNHead.__module__
mock_ctx_managers = [mock_fastrcnn_outputs_inference(tensor_mode)]
if getattr(self.heads, "keypoint_on", False):
mock_ctx_managers += [
mock_keypoint_rcnn_inference(
- tensor_mode, roi_heads_mod, self.use_heatmap_max_keypoint
+ tensor_mode, kpt_heads_mod, self.use_heatmap_max_keypoint
)
]
if getattr(self.heads, "mask_on", False):
