You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
80 lines
2.4 KiB
Python
80 lines
2.4 KiB
Python
import unittest
|
|
|
|
import numpy as np
|
|
|
|
from ballistic_trajectory import predict_ballistic
|
|
from helpers import box_wh
|
|
|
|
|
|
def trajectory_observations(with_outlier=False):
|
|
observations = []
|
|
for index, timestamp in enumerate(np.linspace(-0.4, 0.0, 9)):
|
|
center = np.array([
|
|
100.0 + 50.0 * timestamp + 10.0 * timestamp * timestamp,
|
|
80.0 - 10.0 * timestamp + 3.0 * timestamp * timestamp,
|
|
])
|
|
if with_outlier and index == 3:
|
|
center += np.array([90.0, -70.0])
|
|
size = np.array([
|
|
20.0 * np.exp(0.4 * timestamp),
|
|
10.0 * np.exp(0.4 * timestamp),
|
|
])
|
|
observations.append({
|
|
"ts": float(timestamp),
|
|
"center": center,
|
|
"box": [
|
|
center[0] - 0.5 * size[0],
|
|
center[1] - 0.5 * size[1],
|
|
center[0] + 0.5 * size[0],
|
|
center[1] + 0.5 * size[1],
|
|
],
|
|
})
|
|
return observations
|
|
|
|
|
|
class BallisticTrajectoryTests(unittest.TestCase):
|
|
def test_robust_fit_ignores_single_bad_observation(self):
|
|
prediction = predict_ballistic(
|
|
trajectory_observations(with_outlier=True),
|
|
0.2,
|
|
320,
|
|
240,
|
|
)
|
|
|
|
np.testing.assert_allclose(prediction["center"], [110.4, 78.12], atol=0.2)
|
|
np.testing.assert_allclose(prediction["velocity"], [50.0, -10.0], atol=0.5)
|
|
np.testing.assert_allclose(prediction["acceleration"], [20.0, 6.0], atol=1.0)
|
|
|
|
def test_approach_growth_predicts_larger_box(self):
|
|
observations = trajectory_observations()
|
|
prediction = predict_ballistic(observations, 0.2, 320, 240)
|
|
last_size = box_wh(observations[-1]["box"])
|
|
|
|
self.assertGreater(box_wh(prediction["box"])[0], last_size[0])
|
|
self.assertGreater(box_wh(prediction["box"])[1], last_size[1])
|
|
|
|
def test_prediction_horizon_is_limited(self):
|
|
prediction = predict_ballistic(
|
|
trajectory_observations(),
|
|
2.0,
|
|
320,
|
|
240,
|
|
max_horizon_sec=0.55,
|
|
)
|
|
|
|
self.assertAlmostEqual(prediction["horizon"], 0.55)
|
|
|
|
def test_too_short_history_returns_no_prediction(self):
|
|
prediction = predict_ballistic(
|
|
trajectory_observations()[:3],
|
|
0.2,
|
|
320,
|
|
240,
|
|
)
|
|
|
|
self.assertIsNone(prediction)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|