First lazy BBaug

dereyly · Jun 21, 2020 · cc4dd43 · cc4dd43
1 parent d20acc7
commit cc4dd43
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/mmdet/datasets/pipelines/transforms.py b/mmdet/datasets/pipelines/transforms.py
@@ -7,6 +7,11 @@
 from mmdet.core import PolygonMasks
 from mmdet.core.evaluation.bbox_overlaps import bbox_overlaps
 from ..builder import PIPELINES
+try:
+    import bbaug
+    from bbaug import policies
+except ImportError:
+    use_bbaug = None
 
 try:
     from imagecorruptions import corrupt
@@ -950,6 +955,9 @@ def __init__(self):
         pass
 
     def __call__(self, results):
+        if use_bbaug is None:
+            raise RuntimeError('BBaug is not installed')
+
         policy = policy_container.select_random_policy()
         img_aug, bbs_aug = policy_container.apply_augmentation(policy, results['img'], results['gt_bboxes'], results['gt_labels'])