Fix missing labels after albumentations (ultralytics#4455)

* fix missing labels after augmentation * Update datasets.py Cleanup Co-authored-by: Huu Quan <[email protected]> Co-authored-by: Glenn Jocher <[email protected]>
fhkiel-mlaip · Aug 18, 2021 · 22be136 · 22be136
1 parent 996a0a5
commit 22be136
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/utils/datasets.py b/utils/datasets.py
@@ -568,6 +568,7 @@ def __getitem__(self, index):
         if self.augment:
             # Albumentations
             img, labels = self.albumentations(img, labels)
+            nl = len(labels) # update after albumentations
 
             # HSV color-space
             augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])