Merge pull request MhLiao#259 from jiangxiluning/master

augment keypoints in one run
priyank9320 · May 31, 2021 · d0775a6 · d0775a6
2 parents 4ac194d + 8383c0f
commit d0775a6
Showing 1 changed file with 17 additions and 17 deletions.
diff --git a/data/processes/augment_data.py b/data/processes/augment_data.py
@@ -55,28 +55,28 @@ def process(self, data):
 
 
 class AugmentDetectionData(AugmentData):
-    def may_augment_annotation(self, aug, data, shape):
+
+    def may_augment_annotation(self, aug: imgaug.augmenters.Augmenter, data, shape):
         if aug is None:
             return data
-
+        
         line_polys = []
+        keypoints = []
+        texts = []
         for line in data['lines']:
-            if self.only_resize:
-                new_poly = [(p[0], p[1]) for p in line['poly']]
-            else:
-                new_poly = self.may_augment_poly(aug, shape, line['poly'])
+            texts.append(line['text'])
+            for p in line['poly']:
+                keypoints.append(imgaug.Keypoint(p[0], p[1]))
+
+        keypoints = aug.augment_keypoints([imgaug.KeypointsOnImage(keypoints=keypoints, shape=shape)])[0].keypoints
+        new_polys = np.array([p.x, p.y] for p in keypoints).reshape([-1, 4, 2])
+        for i in range(len(texts)):
+            poly = new_polys[i]
             line_polys.append({
-                'points': new_poly,
-                'ignore': line['text'] == '###',
-                'text': line['text'],
+                'points': poly,
+                'ignore': texts[i] == '###',
+                'text': texts[i]
             })
+
         data['polys'] = line_polys
-        return data
-
-    def may_augment_poly(self, aug, img_shape, poly):
-        keypoints = [imgaug.Keypoint(p[0], p[1]) for p in poly]
-        keypoints = aug.augment_keypoints(
-            [imgaug.KeypointsOnImage(keypoints, shape=img_shape)])[0].keypoints
-        poly = [(p.x, p.y) for p in keypoints]
-        return poly