diff --git a/skais/ais/ais_trajectory.py b/skais/ais/ais_trajectory.py
index 18ad39ce54ec3b8fe572c22926129827659a7256..6289cabc532fcc6e63384a5b76090e85bb2788b0 100644
--- a/skais/ais/ais_trajectory.py
+++ b/skais/ais/ais_trajectory.py
@@ -224,7 +224,7 @@ class AISTrajectory(AISPoints):
result.append((row['ts_sec'], current_label))
return result
- def generate_array_from_positions(self, height=256, width=256, link=True, bounding_box='fit', features=None):
+ def generate_array_from_positions(self, height=256, width=256, link=True, bounding_box='fit', features=None, node_size=0):
nb_channels = 1
if features is not None:
nb_channels = len(features)
@@ -234,10 +234,19 @@ class AISTrajectory(AISPoints):
if bounding_box != 'fit':
raise ValueError("feature not implemented")
positions = self.df[['longitude', 'latitude']].to_numpy()
- range_longitude = (min(positions[0, :]), max(positions[0, :]))
- range_latitude = (min(positions[1, :]), max(positions[1, :]))
+ range_longitude = (min(positions[:, 0]), max(positions[:, 0]))
+ range_latitude = (min(positions[:, 1]), max(positions[:, 1]))
for longitude, latitude in positions:
- x_coord = width * (longitude - range_longitude[0]) / (range_longitude[1] - range_longitude[0])
- y_coord = height * (longitude - range_latitude[0]) / (range_latitude[1] - range_latitude[0])
- data[x_coord, y_coord, :] = 1
+ x_coord = max(min(height - int(height * (latitude - range_latitude[0]) / (range_latitude[1] - range_latitude[0])) - 1, height - 1), 0)
+ y_coord = max(min(int((width - 1) * (longitude - range_longitude[0]) / (range_longitude[1] - range_longitude[0])), width - 1), 0)
+
+ x_lower_bound = max(0, x_coord - node_size)
+ x_upper_bound = min(height - 1, x_coord + node_size)
+
+ y_lower_bound = max(0, y_coord - node_size)
+ y_upper_bound = min(width - 1, y_coord + node_size)
+
+ for x in range(x_lower_bound, x_upper_bound + 1):
+ for y in range(y_lower_bound, y_upper_bound + 1):
+ data[x, y] = [1]
return data
diff --git a/skais/tests/ais/test_ais_trajectory.py b/skais/tests/ais/test_ais_trajectory.py
index bea4d8843011eaabc1ae0b81bf1dc023b3e62706..e4b6c6909161ff8852df6584ca0ed4a2455eb3b8 100644
--- a/skais/tests/ais/test_ais_trajectory.py
+++ b/skais/tests/ais/test_ais_trajectory.py
@@ -40,7 +40,7 @@ class TestAISTrajectory(unittest.TestCase):
for r, e in zip(result, expected):
np.testing.assert_array_equal(r, e)
-
+
def test_sliding_window_too_short(self):
trajectory = AISTrajectory(
pd.DataFrame(
@@ -319,7 +319,7 @@ class TestAISTrajectory(unittest.TestCase):
self.assertEqual(0, compute_trajectory.py_func(times, 800))
def test_apply_func_on_window(self):
- self.assertRaises(ValueError, apply_func_on_window,np.arange(10), 0, 0, 'not valid string')
+ self.assertRaises(ValueError, apply_func_on_window, np.arange(10), 0, 0, 'not valid string')
def test_apply_func_on_window_ignore(self):
result = apply_func_on_window(np.arange(10), np.mean, 2, 'ignore')
@@ -369,7 +369,7 @@ class TestAISTrajectory(unittest.TestCase):
trajectory = AISTrajectory(
pd.DataFrame(
{
- "label": [1 for _ in range(11)] + [2 for _ in range(10)]+ [1 for _ in range(10)],
+ "label": [1 for _ in range(11)] + [2 for _ in range(10)] + [1 for _ in range(10)],
"ts_sec": [i for i in range(0, 18001, 600)]
}
)
@@ -378,4 +378,54 @@ class TestAISTrajectory(unittest.TestCase):
result = trajectory.get_time_per_label_shift()
expected = [(0, 1), (6600, 2), (12600, 1)]
- self.assertListEqual(result, expected)
\ No newline at end of file
+ self.assertListEqual(result, expected)
+
+ def test_generate_array_from_positions(self):
+ trajectory = AISTrajectory(
+ pd.DataFrame(
+ {
+ "latitude": [0, 10, 0, -10],
+ "longitude": [0, 10, 10, -10],
+ "ts_sec": [i for i in range(4)]
+ }
+ )
+ )
+
+ result = trajectory.generate_array_from_positions(height=9, width=9, link=False, bounding_box='fit',
+ features=None, node_size=0).reshape((9, 9))
+ expected = np.array([[0, 0, 0, 0, 0, 0, 0, 0, 1],
+ [0, 0, 0, 0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 1, 0, 0, 0, 1],
+ [0, 0, 0, 0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 0, 0, 0, 0, 0, 0],
+ [1, 0, 0, 0, 0, 0, 0, 0, 0]])
+
+ np.testing.assert_array_equal(result, expected)
+
+ def test_generate_array_from_positions_node_size(self):
+ trajectory = AISTrajectory(
+ pd.DataFrame(
+ {
+ "latitude": [0, 10, 0, -10],
+ "longitude": [0, 10, 10, -10],
+ "ts_sec": [i for i in range(4)]
+ }
+ )
+ )
+
+ result = trajectory.generate_array_from_positions(height=9, width=9, link=False, bounding_box='fit',
+ features=None, node_size=1).reshape((9, 9))
+ expected = np.array([[0, 0, 0, 0, 0, 0, 0, 1, 1],
+ [0, 0, 0, 0, 0, 0, 0, 1, 1],
+ [0, 0, 0, 0, 0, 0, 0, 0, 0],
+ [0, 0, 0, 1, 1, 1, 0, 1, 1],
+ [0, 0, 0, 1, 1, 1, 0, 1, 1],
+ [0, 0, 0, 1, 1, 1, 0, 1, 1],
+ [0, 0, 0, 0, 0, 0, 0, 0, 0],
+ [1, 1, 0, 0, 0, 0, 0, 0, 0],
+ [1, 1, 0, 0, 0, 0, 0, 0, 0]])
+
+ np.testing.assert_array_equal(result, expected)
\ No newline at end of file