Update to new albumentations (#1209)

* Fixed augmentation imports and test_train.py to use the new import paths.

* Add reruns in pytest ini to fix flaky web socket tests.

Author: DocGarbanzo

donkeycar/__init__.py

@@ -3,7 +3,7 @@
 from pyfiglet import Figlet
 import logging
 
-__version__ = '5.2.dev2'
+__version__ = '5.2.dev3'
 
 logging.basicConfig(level=os.environ.get('LOGLEVEL', 'INFO').upper())
 

donkeycar/pipeline/augmentations.py

@@ -2,7 +2,8 @@
 import logging
 import albumentations as A
 from albumentations import GaussianBlur
-from albumentations.augmentations.transforms import RandomBrightnessContrast
+from albumentations.augmentations import RandomBrightnessContrast
+
 
 from donkeycar.config import Config
 
@@ -11,14 +12,14 @@
 
 
 class ImageAugmentation:
-    def __init__(self, cfg, key, prob=0.5, always_apply=False):
+    def __init__(self, cfg, key, prob=0.5):
         aug_list = getattr(cfg, key, [])
-        augmentations = [ImageAugmentation.create(a, cfg, prob, always_apply)
+        augmentations = [ImageAugmentation.create(a, cfg, prob)
                          for a in aug_list]
         self.augmentations = A.Compose(augmentations)
 
     @classmethod
-    def create(cls, aug_type: str, config: Config, prob, always) -> \
+    def create(cls, aug_type: str, config: Config, prob) -> \
             albumentations.core.transforms_interface.BasicTransform:
         """ Augmentation factory. Cropping and trapezoidal mask are
             transformations which should be applied in training, validation
@@ -30,13 +31,13 @@ def create(cls, aug_type: str, config: Config, prob, always) -> \
             logger.info(f'Creating augmentation {aug_type} {b_limit}')
             return RandomBrightnessContrast(brightness_limit=b_limit,
                                             contrast_limit=b_limit,
-                                            p=prob, always_apply=always)
+                                            p=prob)
 
         elif aug_type == 'BLUR':
             b_range = getattr(config, 'AUG_BLUR_RANGE', 3)
             logger.info(f'Creating augmentation {aug_type} {b_range}')
             return GaussianBlur(sigma_limit=b_range, blur_limit=(13, 13),
-                                p=prob, always_apply=always)
+                                p=prob)
 
     # Parts interface
     def run(self, img_arr):

donkeycar/tests/pytest.ini

@@ -5,3 +5,4 @@ filterwarnings =
 
 log_cli = True
 log_cli_level = INFO
+reruns = 3

donkeycar/tests/test_train.py

@@ -106,6 +106,8 @@ def car_dir(tmpdir_factory, base_config, imu_fields) -> str:
         record['localizer/location'] = 3 * count // len(tub)
         tub_full.write_record(record)
         count += 1
+    tub_full.close()
+    tub.close()
     return car_dir