目标追踪逻辑代码
在实现目标检测的前提下,结合小车的基础控制部分,将小车的速度调整依赖到目标检测的推理结果上,就能实现目标追踪。
- “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
- 初始化并导入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
- 在得到正确导入结果后,开启循环,不断获取推理结果,并根据结果估算智能小车和追踪目标之间的距离,再根据计算出的结果下达不同的运动指令,设置慢速和快速跟进的两个速度。
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
- 获取推理结果的外接框。
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 h, w = y2 - y1, x2 - x1
根据计算出的目标框的大小来判断小车和目标之间的距离,再调整小车的行进速度。根据目标近大远小的简单规则,存在两个判断条件,如果目标框的面积小于一定值,就说明小车与目标距离较远,需要快速接近目标,另外如果识别框的中心点的纵坐标大于0,也就是在摄像头视角里的上半部分,也说明小车距离目标较远,也需快速接近目标,反之亦然。
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()
父主题: 代码实现