红外小目标跟踪python完整Pipeline实现

bashpip   install opencv-python numpy scipy

pythonimport   cv2   import numpy as np   from scipy.optimize import linear_sum_assignment   # ======================   # 1. 红外图像预处理   # ======================   def preprocess_infrared(img):"""红外图像预处理：灰度化+背景抑制+噪声过滤"""if len(img.shape) == 3:gray = cv2.cvtColor(img,   cv2.COLOR_BGR2GRAY)else:gray = img.copy()# 高斯滤波去除椒盐噪声gray = cv2.GaussianBlur(gray, (3,   3), 0)return gray   # ======================   # 2. 顶帽变换提取高亮小目标   # ======================   def tophat_detect(gray, kernel_size=9):"""形态学顶帽变换：提取比背景亮的小目标，抑制大尺度背景"""kernel =   cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))tophat = cv2.morphologyEx(gray,   cv2.MORPH_TOPHAT, kernel)# 阈值分割得到二值图_, binary = cv2.threshold(tophat,   0, 255, cv2.THRESH_OTSU)return binary   # ======================   # 3. LCM局部对比度检测（红外小目标经典算法）   # ======================   def lcm_detect(gray, win_size=5, bg_win_size=15, threshold=2.0):"""局部对比度度量LCM：计算每个像素与周围背景的对比度亮目标对比度高，噪声/背景对比度低"""gray_float =   gray.astype(np.float32)# 计算局部背景均值和标准差bg_mean = cv2.blur(gray_float,   (bg_win_size, bg_win_size))bg_sq_mean = cv2.blur(gray_float **   2, (bg_win_size, bg_win_size))bg_std =   np.sqrt(np.maximum(bg_sq_mean - bg_mean ** 2, 1e-6))# 计算局部对比度图lcm_map = (gray_float - bg_mean) /   bg_stdlcm_map[lcm_map < threshold] =   0# 低于阈值的背景置零# 非极大值抑制，提取局部峰值点lcm_map = cv2.dilate(lcm_map,   np.ones((3,3), np.uint8))binary = np.uint8(lcm_map > 0) *   255return binary   # ======================   # 4. 从二值图提取候选目标框   # ======================   def get_detections(binary, min_area=2, max_area=100):"""提取连通域，返回检测框列表 [x,y,w,h]"""contours, _ =   cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)detections = []for cnt in contours:area = cv2.contourArea(cnt)if min_area < area <   max_area:x, y, w, h =   cv2.boundingRect(cnt)detections.append([x, y, w,   h])return np.array(detections) if   len(detections) > 0 else np.empty((0, 4))

python#   ======================   # 单个轨迹的卡尔曼滤波器   # ======================   class KalmanTrack:def __init__(self, bbox):# 状态向量：[x,   y, vx, vy]中心坐标+速度self.kf = cv2.KalmanFilter(4,   2)# 状态转移矩阵（匀速模型）self.kf.transitionMatrix =   np.array([[1, 0, 1, 0],[0, 1, 0, 1],[0, 0, 1, 0],[0, 0, 0, 1]], np.float32)# 测量矩阵self.kf.measurementMatrix =   np.array([[1, 0, 0, 0],[0, 1, 0, 0]], np.float32)# 初始化状态x, y, w, h = bboxself.kf.statePost =   np.array([[x + w/2], [y + h/2], [0], [0]], np.float32)# 过程噪声协方差self.kf.processNoiseCov =   np.eye(4, dtype=np.float32) * 0.03# 测量噪声协方差self.kf.measurementNoiseCov =   np.eye(2, dtype=np.float32) * 1.0self.time_since_update = 0# 未更新帧数self.id = 0self.hits = 0# 匹配成功次数self.is_confirmed = False# 是否为有效轨迹def predict(self):"""预测下一帧位置"""pred = self.kf.predict()self.time_since_update += 1return pred[:2].flatten()# 返回预测的中心坐标 [x, y]def update(self, bbox):"""用检测结果更新"""x, y, w, h = bboxmeasurement = np.array([[x +   w/2], [y + h/2]], np.float32)self.kf.correct(measurement)self.time_since_update = 0self.hits += 1# 连续匹配3帧以上确认为有效轨迹if self.hits >= 3:self.is_confirmed = Truedef get_bbox(self, default_size=6):"""获取当前目标框（极小目标默认固定尺寸）"""x, y =   self.kf.statePost[:2].flatten()return [int(x -   default_size/2), int(y - default_size/2), default_size, default_size]   # ======================   # 计算IOU（用于匹配代价）   # ======================   def bbox_iou(box1, box2):x1, y1, w1, h1 = box1x2, y2, w2, h2 = box2inter_x1 = max(x1, x2)inter_y1 = max(y1, y2)inter_x2 = min(x1 + w1, x2 + w2)inter_y2 = min(y1 + h1, y2 + h2)inter_area = max(0, inter_x2 -   inter_x1) * max(0, inter_y2 - inter_y1)union_area = w1 * h1 + w2 * h2 -   inter_areareturn inter_area / (union_area +   1e-6)   # ======================   # 纯运动跟踪器主类   # ======================   class MotionTracker:def __init__(self, max_age=10,   min_hits=3, iou_threshold=0.1):self.tracks = []self.next_id = 1self.max_age = max_age# 最大丢失帧数self.min_hits = min_hits# 确认轨迹最小匹配次数self.iou_threshold =   iou_thresholddef update(self, detections):# 1. 所有轨迹先做预测for track in self.tracks:track.predict()# 2. 计算代价矩阵（IOU越大代价越小）pred_boxes = [t.get_bbox() for   t in self.tracks]det_boxes = detections.tolist()   if len(detections) > 0 else []if len(pred_boxes) > 0 and   len(det_boxes) > 0:cost_matrix =   np.zeros((len(pred_boxes), len(det_boxes)))for i, pred in   enumerate(pred_boxes):for j, det in   enumerate(det_boxes):cost_matrix[i, j] =   1 - bbox_iou(pred, det)# 匈牙利算法匹配row_ind, col_ind =   linear_sum_assignment(cost_matrix)matched_trks = set()matched_dets = set()for i, j in zip(row_ind,   col_ind):if cost_matrix[i, j]   < (1 - self.iou_threshold):self.tracks[i].update(det_boxes[j])matched_trks.add(i)matched_dets.add(j)unmatched_dets = [j for j   in range(len(det_boxes)) if j not in matched_dets]else:unmatched_dets =   list(range(len(det_boxes)))# 3. 未匹配的检测新建轨迹for j in unmatched_dets:new_track =   KalmanTrack(det_boxes[j])new_track.id = self.next_idself.next_id += 1self.tracks.append(new_track)# 4. 删除长时间丢失的轨迹self.tracks = [t for t in   self.tracks if t.time_since_update < self.max_age]# 5. 返回已确认的轨迹return [t.get_bbox() + [t.id]   for t in self.tracks if t.is_confirmed]

pythonclass   SimpleTBD:def __init__(self, frame_num=12,   decay=0.88):""":param frame_num: 累积帧数:param decay: 历史帧衰减系数，越新的帧权重越高"""self.frame_num = frame_numself.decay = decayself.frame_buffer = []def update(self, gray_frame):"""输入单帧灰度图，输出多帧累积增强图"""self.frame_buffer.append(gray_frame.astype(np.float32))if len(self.frame_buffer) >   self.frame_num:self.frame_buffer.pop(0)# 加权累积，越新的帧权重越高acc_map =   np.zeros_like(gray_frame, dtype=np.float32)total_weight = 0.0for idx, frame in   enumerate(reversed(self.frame_buffer)):weight = self.decay ** idxacc_map += frame * weighttotal_weight += weightacc_map = acc_map /   total_weight# 归一化到0-255acc_map =   cv2.normalize(acc_map, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)return acc_map

pythonclass   InfraredSmallTargetPipeline:def __init__(self):self.tbd =   SimpleTBD(frame_num=12, decay=0.88)self.tracker =   MotionTracker(max_age=10, min_hits=3, iou_threshold=0.1)# 分级阈值（单位：像素）self.super_far_thresh = 10# 超远距离：<10像素self.mid_thresh = 30# 中距离：10~30像素# 近距离可接入YOLO等AI检测模型，此处用传统检测代替演示def process_frame(self, img):"""处理单帧图像:param img: 输入红外图像:return: 可视化图像、跟踪结果"""# 1. 预处理gray = preprocess_infrared(img)# 2. 超远距离：TBD多帧累积 + LCM检测acc_map = self.tbd.update(gray)lcm_bin = lcm_detect(acc_map,   win_size=5, bg_win_size=15, threshold=1.8)detections =   get_detections(lcm_bin, min_area=2, max_area=200)# 3. 运动跟踪tracks =   self.tracker.update(detections)# 4. 可视化vis = cv2.cvtColor(gray,   cv2.COLOR_GRAY2BGR)# 画检测框（黄色）for x, y, w, h in detections:cv2.rectangle(vis, (x, y), (x+w, y+h), (0,   255, 255), 1)# 画跟踪轨迹（绿色，带ID）for x, y, w, h, track_id in   tracks:cv2.rectangle(vis, (x, y),   (x+w, y+h), (0, 255, 0), 2)cv2.putText(vis,   f"ID:{track_id}", (x, y-5), cv2.FONT_HERSHEY_SIMPLEX,   0.4, (0, 255, 0), 1)return vis, tracks   # ======================   # 运行示例：处理视频/摄像头   # ======================   if __name__ == "__main__":# 替换为你的红外视频路径，0表示调用摄像头cap = cv2.VideoCapture("infrared_video.mp4")pipeline =   InfraredSmallTargetPipeline()while True:ret, frame = cap.read()if not ret:breakvis_frame, tracks =   pipeline.process_frame(frame)cv2.imshow("Infrared Small   Target Tracking", vis_frame)if cv2.waitKey(1) & 0xFF ==   ord('q'):breakcap.release()cv2.destroyAllWindows()