下载
中文
注册

目标追踪逻辑代码

在实现目标检测的前提下,结合小车的基础控制部分,将小车的速度调整依赖到目标检测的推理结果上,就能实现目标追踪。

  1. “python/src/scenes/tracking.py”为目标追踪的核心代码,示例代码定义追踪的运行逻辑。
    class Tracking(BaseScene):
        def __init__(self, memory_name, camera_info, msg_queue):
            super().__init__(memory_name, camera_info, msg_queue)
            self.model = None
     
        def init_state(self):
            log.info(f'start init {self.__class__.__name__}')
            model_path = os.path.join(os.getcwd(), 'weights', 'tracking.om')
            if not os.path.exists(model_path):
                log.error(f'Cannot find the offline inference model(.om) file needed for {self.__class__.__name__}  scene.')
                return True
            self.model = YoloV5(model_path) #加载模型
            log.info(f'{self.__class__.__name__} model init succ.')
            self.ctrl.execute(SetServo(servo=[90, 65])) #设置舵机角度
            return False
     
        def loop(self):
            ret = self.init_state() #执行初始化
            if ret:
                log.error(f'{self.__class__.__name__} init failed.')
                return
            frame = np.ndarray((self.height, self.width, 3), dtype=np.uint8, buffer=self.broadcaster.buf) #获取共享内存中的图片
            log.info(f'{self.__class__.__name__} loop start')
            last_action = None
            last_not_seen = True
            forward_speed_slow = 30
            forward_speed_fast = 40
            while True:
                action = None
                if self.stop_sign.value:
                    break
                if self.pause_sign.value:
                    continue
     
                img_bgr = frame.copy()
                bboxes = self.model.infer(img_bgr)
                log.info(f'{bboxes}')
                if not bboxes:
                    if last_not_seen:
                        action = Stop()
                    else:
                        last_not_seen = True
                        continue
                else:
                    if len(bboxes) > 1:
                        ori_box = sorted(bboxes, key=lambda x: x[-1], reverse=True)[0][:4]
                    else:
                        ori_box = bboxes[0][:4]
                    x1, y1, x2, y2 = ori_box
                    x, y = (x1 + x2) // 2, (y1 + y2) // 2 #计算目标中心点的x与y坐标
                    h, w = y2 - y1, x2 - x1 #计算目标的宽高
     
                    if h * w < 141 * 128 or y < 110: #进行距离判断,如果过远就加速,否则减速
                        speed = forward_speed_fast
                    else:
                        speed = forward_speed_slow
     
                    if x < 400:
                        action = TurnLeft(degree=1.1, speed=speed) #左转
                    elif x > 1000:
                        action = TurnRight(degree=1.1, speed=speed) #右转
                    else:
                        action = Advance(speed=speed) #直行
     
                    if h * w > 800 * 500 or y > 390: #如果距离过近则停车
                        action = Stop()
     
                if action is None or action == last_action:
                    continue
                self.ctrl.execute(action)
                last_action = action
  2. 初始化并导入Yolov5目标检测模型。
    def __init__(self, memory_name, camera_info, msg_queue):
            super().__init__(memory_name, camera_info, msg_queue)
            self.model = None
     
        def init_state(self):
            log.info(f'start init {self.__class__.__name__}')
            model_path = os.path.join(os.getcwd(), 'weights', 'tracking.om')
            if not os.path.exists(model_path):
                log.error(f'Cannot find the offline inference model(.om) file needed for {self.__class__.__name__}  scene.')
                return True
            self.model = YoloV5(model_path)
            log.info(f'{self.__class__.__name__} model init succ.')
            self.ctrl.execute(SetServo(servo=[90, 65]))
            return False
  3. 在得到正确导入结果后,开启循环,不断获取推理结果,并根据结果估算智能小车和追踪目标之间的距离,再根据计算出的结果下达不同的运动指令,设置慢速和快速跟进的两个速度。
    def loop(self):
            ret = self.init_state()
            if ret:
                log.error(f'{self.__class__.__name__} init failed.')
                return
            frame = np.ndarray((self.height, self.width, 3), dtype=np.uint8, buffer=self.broadcaster.buf)
            log.info(f'{self.__class__.__name__} loop start')
            last_action = None
            last_not_seen = True
            forward_speed_slow = 30
            forward_speed_fast = 40
  4. 获取推理结果的外接框。
    bboxes = self.model.infer(img_bgr)
  5. 计算出目标框的中心点的位置和目标框的宽高大小。
    log.info(f'{bboxes}')
                if not bboxes:
                    if last_not_seen:
                        action = Stop()
                    else:
                        last_not_seen = True
                        continue
                else:
                    if len(bboxes) > 1:
                        ori_box = sorted(bboxes, key=lambda x: x[-1], reverse=True)[0][:4]
                    else:
                        ori_box = bboxes[0][:4]
                    x1, y1, x2, y2 = ori_box
                    x, y = (x1 + x2) // 2, (y1 + y2) // 2
                    h, w = y2 - y1, x2 - x1

    根据计算出的目标框的大小来判断小车和目标之间的距离,再调整小车的行进速度。根据目标近大远小的简单规则,存在两个判断条件,如果目标框的面积小于一定值,就说明小车与目标距离较远,需要快速接近目标,另外如果识别框的中心点的纵坐标大于0,也就是在摄像头视角里的上半部分,也说明小车距离目标较远,也需快速接近目标,反之亦然。

    if h * w < 141 * 128 or y < 110:
                        speed = forward_speed_fast
                    else:
                        speed = forward_speed_slow
  6. 另外如果前方目标在小车的偏左或偏右的位置,也可以采用同样的判断方法,即判断目标框的中心点的横坐标落在小车摄像头视角画面中的左侧还是右侧,进而下发对应的微调转向的命令,实现跟踪目标的方向调整。
    if x < 400:
                        action = TurnLeft(degree=1.1, speed=speed)
                    elif x > 1000:
                        action = TurnRight(degree=1.1, speed=speed)
                    else:
                        action = Advance(speed=speed)
     
                    if h * w > 800 * 500 or y > 390:
                        action = Stop()