diff --git a/agent.py b/agent.py index 6e1269b..18eba90 100644 --- a/agent.py +++ b/agent.py @@ -29,7 +29,8 @@ class UAV: self.send_rc_control = False self.manual_control = ManualControl(self) - self.status = Status.INIT + self.auto_control = AutoControl(self) + self.status = Status.DETECT self.frame_queue = Queue(maxsize=1) # 用于存放视频帧, 1用于避免帧堆积 @@ -88,9 +89,14 @@ class UAV: self.keep_recording = False cv2.destroyAllWindows() - def auto_control(self): - if self.status == Status.INIT: - pass + def auto_track_land(self): + + if self.status == Status.INIT or self.status == Status.TRACKING: + front_img = self.frame_queue.get() + if self.status == Status.DETECT: + self.auto_control.detect(front_img=front_img) + elif self.status == Status.TRACK: + self.auto_control.track(front_img=front_img) diff --git a/control.py b/control.py index cf228c2..c5a8161 100644 --- a/control.py +++ b/control.py @@ -91,8 +91,7 @@ class AutoControl: self.target = [0, 0, 0, 0] self.pid = PIDController(1, 0, 0, time.time()) - def detect(self): - front_img = self.uav.frame_queue.get() + def detect(self, front_img): cv2.imshow("ROI select", front_img[:, :, 0:3]) self.gROI = cv2.selectROI("ROI select", front_img[:, :, 0:3], False) if (not self.gROI): @@ -101,11 +100,16 @@ class AutoControl: self.gTracker = Tracker(tracker_type="KCF") self.gTracker.initWorking(front_img[:, :, 0:3], self.gROI) print("start tracking") - self.uav.status = Status.TRACKING + self.uav.status = Status.TRACK cv2.destroyWindow("ROI select") - def track(self): - front_img = self.uav.frame_queue.get() + def track(self, front_img): + _item = self.gTracker.track(front_img) + if _item.message: + self.target = _item.getMessage()['target'] + + else: + return None class FakeUAV: