changed to mmcv-lite

egoforce_runtime_patches.py

@@ -12,6 +12,74 @@ def _torchvision_nms():
     return nms
 
 
+def _torchvision_roi_align():
+    from torchvision.ops import roi_align
+
+    return roi_align
+
+
+def _torchvision_roi_align_module():
+    from torchvision.ops import RoIAlign
+
+    return RoIAlign
+
+
+def _torchvision_roi_pool_module():
+    from torchvision.ops import RoIPool
+
+    return RoIPool
+
+
+def _bbox_overlaps(
+    bboxes1: Any,
+    bboxes2: Any,
+    mode: str = "iou",
+    aligned: bool = False,
+    offset: int = 0,
+    eps: float = 1e-6,
+) -> Any:
+    import torch
+
+    if bboxes1.numel() == 0 or bboxes2.numel() == 0:
+        if aligned:
+            return bboxes1.new_zeros((bboxes1.shape[0],))
+        return bboxes1.new_zeros((bboxes1.shape[0], bboxes2.shape[0]))
+
+    if aligned:
+        lt = torch.maximum(bboxes1[:, :2], bboxes2[:, :2])
+        rb = torch.minimum(bboxes1[:, 2:], bboxes2[:, 2:])
+        wh = (rb - lt + offset).clamp(min=0)
+        overlap = wh[:, 0] * wh[:, 1]
+        area1 = (bboxes1[:, 2] - bboxes1[:, 0] + offset) * (bboxes1[:, 3] - bboxes1[:, 1] + offset)
+        area2 = (bboxes2[:, 2] - bboxes2[:, 0] + offset) * (bboxes2[:, 3] - bboxes2[:, 1] + offset)
+    else:
+        lt = torch.maximum(bboxes1[:, None, :2], bboxes2[None, :, :2])
+        rb = torch.minimum(bboxes1[:, None, 2:], bboxes2[None, :, 2:])
+        wh = (rb - lt + offset).clamp(min=0)
+        overlap = wh[..., 0] * wh[..., 1]
+        area1 = ((bboxes1[:, 2] - bboxes1[:, 0] + offset) * (bboxes1[:, 3] - bboxes1[:, 1] + offset))[:, None]
+        area2 = ((bboxes2[:, 2] - bboxes2[:, 0] + offset) * (bboxes2[:, 3] - bboxes2[:, 1] + offset))[None, :]
+
+    if mode == "iof":
+        union = area1
+    elif mode == "giou":
+        union = area1 + area2 - overlap
+        if aligned:
+            enclosed_lt = torch.minimum(bboxes1[:, :2], bboxes2[:, :2])
+            enclosed_rb = torch.maximum(bboxes1[:, 2:], bboxes2[:, 2:])
+        else:
+            enclosed_lt = torch.minimum(bboxes1[:, None, :2], bboxes2[None, :, :2])
+            enclosed_rb = torch.maximum(bboxes1[:, None, 2:], bboxes2[None, :, 2:])
+        enclosed_wh = (enclosed_rb - enclosed_lt + offset).clamp(min=0)
+        enclosed_area = enclosed_wh[..., 0] * enclosed_wh[..., 1]
+        iou = overlap / union.clamp(min=eps)
+        return iou - (enclosed_area - union) / enclosed_area.clamp(min=eps)
+    else:
+        union = area1 + area2 - overlap
+
+    return overlap / union.clamp(min=eps)
+
+
 def _nms(
     boxes: Any,
     scores: Any,
@@ -104,16 +172,28 @@ def apply_runtime_patches() -> None:
     if ops_module is None:
         ops_module = types.ModuleType("mmcv.ops")
         sys.modules["mmcv.ops"] = ops_module
+    ops_module.__path__ = []
 
     nms_module = sys.modules.get("mmcv.ops.nms")
     if nms_module is None:
         nms_module = types.ModuleType("mmcv.ops.nms")
         sys.modules["mmcv.ops.nms"] = nms_module
 
+    roi_align_module = sys.modules.get("mmcv.ops.roi_align")
+    if roi_align_module is None:
+        roi_align_module = types.ModuleType("mmcv.ops.roi_align")
+        sys.modules["mmcv.ops.roi_align"] = roi_align_module
+
     ops_module.nms = _nms
     ops_module.batched_nms = _batched_nms
+    ops_module.bbox_overlaps = _bbox_overlaps
+    ops_module.roi_align = _torchvision_roi_align()
+    ops_module.RoIAlign = _torchvision_roi_align_module()
+    ops_module.RoIPool = _torchvision_roi_pool_module()
     nms_module.nms = _nms
     nms_module.batched_nms = _batched_nms
+    roi_align_module.roi_align = ops_module.roi_align
+    roi_align_module.RoIAlign = ops_module.RoIAlign
 
     if mmcv is not None:
         mmcv.ops = ops_module