fixed bug in CoarseDropout 'apply_to_keypoints' (#1330)

* fixed bug in CoarseDropout 'apply_to_keypoints'

* added pytests, improved code efficiency

albumentations/augmentations/dropout/coarse_dropout.py

@@ -167,14 +167,12 @@ def _keypoint_in_hole(self, keypoint: KeypointType, hole: Tuple[int, int, int, i
     def apply_to_keypoints(
         self, keypoints: Sequence[KeypointType], holes: Iterable[Tuple[int, int, int, int]] = (), **params
     ) -> List[KeypointType]:
-        result = []
+        result = set(keypoints)
         for hole in holes:
-            remaining_keypoints = []
             for kp in keypoints:
-                if not self._keypoint_in_hole(kp, hole):
-                    remaining_keypoints.append(kp)
-            result = remaining_keypoints
-        return result
+                if self._keypoint_in_hole(kp, hole):
+                    result.discard(kp)
+        return list(result)
 
     def get_transform_init_args_names(self):
         return (

tests/test_keypoint.py

@@ -304,3 +304,21 @@ def test_coarse_dropout():
 
     result = aug(image=np.zeros((128, 128)), keypoints=((10, 10), (20, 30)))
     assert len(result["keypoints"]) == 0
+
+
+@pytest.mark.parametrize(
+    ["keypoints", "expected_keypoints", "holes"],
+    [
+        [[(50, 50, 0, 0), (75, 75, 0, 0)], [], [(40, 40, 60, 60), (70, 70, 80, 80), (10, 10, 20, 20)]],
+        [[(50, 50, 0, 0), (75, 75, 0, 0)], [], [(10, 10, 20, 20), (40, 40, 60, 60), (70, 70, 80, 80)]],
+        [[(50, 50, 0, 0), (75, 75, 0, 0)], [], [(40, 40, 60, 60), (10, 10, 20, 20), (70, 70, 80, 80)]],
+        [[(50, 50, 0, 0), (75, 75, 0, 0)], [(75, 75, 0, 0)], [(40, 40, 60, 60), (10, 10, 20, 20)]],
+        [[(50, 50, 0, 0), (75, 75, 0, 0)], [(50, 50, 0, 0)], [(70, 70, 80, 80), (10, 10, 20, 20)]],
+        [[(50, 50, 0, 0), (75, 75, 0, 0)], [(50, 50, 0, 0), (75, 75, 0, 0)], [(10, 10, 20, 20)]],
+    ],
+)
+def test_coarse_dropout_remove_keypoints(keypoints, expected_keypoints, holes):
+    t = A.CoarseDropout()
+    result_keypoints = t.apply_to_keypoints(keypoints, holes)
+
+    assert set(result_keypoints) == set(expected_keypoints)