Simsim.py 3.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  1. import matplotlib.pyplot as plt
  2. import numpy as np
  3. from target import Target
  4. from tracker import Tracker
  5. import logging
  6. class Simsim:
  7. def __init__(self) -> None:
  8. self.trackers = []
  9. self.edges = []
  10. self.targets = []
  11. self.map_size = [1000, 1000]
  12. self.rate = 30
  13. plt.figure(figsize=(8, 8))
  14. plt.ion()
  15. plt.axis('equal')
  16. def add_tracker(self, name, position, sensor_rad):
  17. tracker = Tracker(name, len(self.trackers),
  18. position, sensor_rad, self.rate)
  19. self.trackers.append(tracker)
  20. def add_edges(self, edges):
  21. self.edges.extend(edges)
  22. for e in edges:
  23. self.trackers[e[0]].neighbor.add(e[1])
  24. self.trackers[e[1]].neighbor.add(e[0])
  25. def add_target(self, name, position):
  26. target = Target(name, len(self.targets), position, self.rate)
  27. self.targets.append(target)
  28. def _update_all(self):
  29. """Simulate once, update all trackers and targets
  30. """
  31. for tracker in self.trackers:
  32. tracker.job()
  33. for target in self.targets:
  34. target.job()
  35. pass
  36. def run(self, log_lvl=logging.WARN, ground_truth=False):
  37. logging.basicConfig(format='%(asctime)s.%(msecs)03d %(levelname)s: %(message)s',
  38. datefmt='%m/%d/%Y %H:%M:%S', level=log_lvl)
  39. logging.warning("The lvl is WARN")
  40. while 1:
  41. self._update_all()
  42. plt.cla()
  43. plt.xlim(0, self.map_size[0])
  44. plt.ylim(0, self.map_size[1])
  45. if ground_truth:
  46. for target in self.targets:
  47. plt.scatter(
  48. target.position[0], target.position[1], marker='d', s=10, color='r')
  49. for e in self.edges:
  50. t0 = self.trackers[e[0]]
  51. t1 = self.trackers[e[1]]
  52. plt.plot([t0.position[0], t1.position[0]],
  53. [t0.position[1], t1.position[1]],
  54. linewidth=1, color='g', alpha=0.5)
  55. for t in self.trackers:
  56. # draw the trackers
  57. plt.scatter(t.position[0], t.position[1],
  58. marker='s', s=20, c='b')
  59. plt.annotate(t.id, (t.position[0], t.position[1]+10))
  60. # draw the camera coverage
  61. circle = plt.Circle(
  62. (t.position), t.sensor.coverage_radius, fill=False, color='grey', alpha=0.3)
  63. plt.gcf().gca().add_artist(circle)
  64. # draw the detections
  65. dets_relative = t.sensor.get_detection(self.targets)
  66. for det in dets_relative:
  67. det_abs = det+t.position
  68. plt.scatter(det_abs[0], det_abs[1],
  69. marker='^', s=8) # , alpha=det[2]
  70. plt.draw()
  71. plt.pause(1/self.rate)
  72. s = Simsim()
  73. s.add_tracker('tracker', np.array([120.0, 180.0]), sensor_rad=100)
  74. # s.trackers[0].speed = np.array([10, 10])
  75. s.add_tracker('tracker', np.array([220, 380]), sensor_rad=150)
  76. # s.trackers[1].speed = np.array([5, 10])
  77. s.add_tracker('tracker', np.array([500, 380]), sensor_rad=200)
  78. s.add_tracker('tracker', np.array([700, 580]), sensor_rad=120)
  79. # s.trackers[2].speed = np.array([-22, 10])
  80. s.add_target('tgt', np.array([200, 300]))
  81. s.add_target('tgt', np.array([100, 300]))
  82. s.add_target('tgt', np.array([500, 300]))
  83. s.add_target('tgt', np.array([300, 500]))
  84. s.add_edges([[0, 1], [0, 2], [1,2], [0,3]])
  85. s.run(logging.INFO, True)