diff --git a/skais/process/data_augmentation/augmentation_engine.py b/skais/process/data_augmentation/augmentation_engine.py
index 27004aa58f514f39de26a75242f8d92a4d53f490..64ecb499ab660c94d27e98ce21058f70e3b75ad1 100644
--- a/skais/process/data_augmentation/augmentation_engine.py
+++ b/skais/process/data_augmentation/augmentation_engine.py
@@ -28,7 +28,7 @@ class AugmentationEngine:
self.pipelines.append(Pipeline([translator, flip]))
def transform(self, x, verbose=0):
- results = x.copy()
+ results = []
iterator = self.pipelines
if verbose > 0:
diff --git a/skais/process/data_augmentation/flip.py b/skais/process/data_augmentation/flip.py
index 578ca56b93ed09dda3f219524411b5c20ae86b71..0917becd8cef0e1ab55a405cba84ec7b2bcd0418 100644
--- a/skais/process/data_augmentation/flip.py
+++ b/skais/process/data_augmentation/flip.py
@@ -14,6 +14,4 @@ class Flip(DataTransformer):
df = trajectory.df.copy()
df['latitude'] = -trajectory.df['latitude']
result.append(AISTrajectory(df))
- else:
- result += x.copy()
return result
diff --git a/skais/tests/process/data_augmentation/test_engine.py b/skais/tests/process/data_augmentation/test_engine.py
index c77c55db68997644971bf9b2ce3dbf2a0b1e4e4d..9ea8d4cd92fe338e037e12d21083adb1672c1cfe 100644
--- a/skais/tests/process/data_augmentation/test_engine.py
+++ b/skais/tests/process/data_augmentation/test_engine.py
@@ -5,7 +5,8 @@ from skais.process.data_augmentation.augmentation_engine import AugmentationEngi
import pandas as pd
-class Test_Engine(unittest.TestCase):
+
+class TestEngine(unittest.TestCase):
def setUp(self):
t1 = AISTrajectory(
pd.DataFrame(
@@ -28,7 +29,6 @@ class Test_Engine(unittest.TestCase):
self.trajectories = [t1, t2]
-
def test_transform_simple_translation(self):
engine = AugmentationEngine(translation_values=[(10, 0), (20, 0)], keep_original=False)
@@ -74,20 +74,116 @@ class Test_Engine(unittest.TestCase):
)
expected = [t1, t2, t3, t4]
- self.assertEqual(len(expected), len(result))
+ self.assertEqual(len(result), len(expected))
for t1, t2 in zip(result, expected):
pd.testing.assert_frame_equal(t1.df, t2.df)
def test_transform_simple_flip(self):
- engine = AugmentationEngine(flip_values=[None, 0], keep_original=False)
+ engine = AugmentationEngine(flip_values=[(None, 0)], keep_original=False)
+
+ result = engine.transform(self.trajectories)
+
+ t1 = AISTrajectory(
+ pd.DataFrame(
+ {
+ 'ts_sec': [i for i in range(10)],
+ 'latitude': [0 for _ in range(10)],
+ 'longitude': [12 + i for i in range(10)]
+ }
+ )
+ )
+ t2 = AISTrajectory(
+ pd.DataFrame(
+ {
+ 'ts_sec': [i for i in range(10)],
+ 'latitude': [12 - i for i in range(10)],
+ 'longitude': [12 + i for i in range(10)]
+ }
+ )
+ )
+ expected = [t1, t2]
+
+ self.assertEqual(len(expected), len(result))
+ for t1, t2 in zip(result, expected):
+ pd.testing.assert_frame_equal(t1.df, t2.df)
+
+ def test_transform_flip_and_translate(self):
+ engine = AugmentationEngine(translation_values=[(10, 0)], flip_values=[(None, 0)], keep_original=False)
result = engine.transform(self.trajectories)
+ t1 = AISTrajectory(
+ pd.DataFrame(
+ {
+ 'ts_sec': [i for i in range(10)],
+ 'latitude': [0 for _ in range(10)],
+ 'longitude': [22 + i for i in range(10)]
+ }
+ )
+ )
+ t2 = AISTrajectory(
+ pd.DataFrame(
+ {
+ 'ts_sec': [i for i in range(10)],
+ 'latitude': [-12 + i for i in range(10)],
+ 'longitude': [22 + i for i in range(10)]
+ }
+ )
+ )
+
+ t3 = AISTrajectory(
+ pd.DataFrame(
+ {
+ 'ts_sec': [i for i in range(10)],
+ 'latitude': [0 for _ in range(10)],
+ 'longitude': [12 + i for i in range(10)]
+ }
+ )
+ )
+ t4 = AISTrajectory(
+ pd.DataFrame(
+ {
+ 'ts_sec': [i for i in range(10)],
+ 'latitude': [12 - i for i in range(10)],
+ 'longitude': [12 + i for i in range(10)]
+ }
+ )
+ )
+
+ t5 = AISTrajectory(
+ pd.DataFrame(
+ {
+ 'ts_sec': [i for i in range(10)],
+ 'latitude': [0 for _ in range(10)],
+ 'longitude': [22 + i for i in range(10)]
+ }
+ )
+ )
+ t6 = AISTrajectory(
+ pd.DataFrame(
+ {
+ 'ts_sec': [i for i in range(10)],
+ 'latitude': [12 - i for i in range(10)],
+ 'longitude': [22 + i for i in range(10)]
+ }
+ )
+ )
+
+ expected = [t1, t2, t3, t4, t5, t6]
+
+ self.assertEqual(len(expected), len(result))
+ for t1, t2 in zip(result, expected):
+ pd.testing.assert_frame_equal(t1.df, t2.df)
+
+ def test_engine_verbose(self):
+ engine = AugmentationEngine(flip_values=[(None, 0)], keep_original=False)
+
+ result = engine.transform(self.trajectories, verbose=1)
t1 = AISTrajectory(
pd.DataFrame(
{
'ts_sec': [i for i in range(10)],
- 'latitude': [0 - i for i in range(10)],
+ 'latitude': [0 for _ in range(10)],
'longitude': [12 + i for i in range(10)]
}
)
@@ -107,5 +203,35 @@ class Test_Engine(unittest.TestCase):
for t1, t2 in zip(result, expected):
pd.testing.assert_frame_equal(t1.df, t2.df)
+ def test_engine_keep_original(self):
+ engine = AugmentationEngine(flip_values=[(None, 0)], keep_original=True)
+
+ result = engine.transform(self.trajectories, verbose=1)
+
+ t1 = AISTrajectory(
+ pd.DataFrame(
+ {
+ 'ts_sec': [i for i in range(10)],
+ 'latitude': [0 for _ in range(10)],
+ 'longitude': [12 + i for i in range(10)]
+ }
+ )
+ )
+ t2 = AISTrajectory(
+ pd.DataFrame(
+ {
+ 'ts_sec': [i for i in range(10)],
+ 'latitude': [12 - i for i in range(10)],
+ 'longitude': [12 + i for i in range(10)]
+ }
+ )
+ )
+ expected = self.trajectories + [t1, t2]
+
+ self.assertEqual(len(expected), len(result))
+ for t1, t2 in zip(result, expected):
+ pd.testing.assert_frame_equal(t1.df, t2.df)
+
+
if __name__ == '__main__':
unittest.main()