基于运动补偿的前景检测算法
这段代码实现了基于运动补偿的前景检测算法。
主要功能包括:
- 运动补偿模块:使用基于网格的 KLT 特征跟踪算法计算两帧之间的运动,然后通过单应性变换实现帧间运动补偿。
- 前景检测模块:结合两帧运动补偿结果,通过帧间差分计算前景掩码。
- 异常处理:添加了图像加载检查和异常捕获,提高了代码的健壮性。
- 路径处理:自动创建保存目录,避免因目录不存在导致的错误。
使用时需要提供三帧连续图像:两个参考帧和当前帧。代码会计算出前景掩码并保存为图像文件。
import cv2
import numpy as np
import os
import syskernel_size = 3def motion_compensate(frame1, frame2):"""使用基于网格的KLT特征跟踪实现两帧之间的运动补偿参数:frame1: 前一帧图像(BGR格式)frame2: 当前帧图像(BGR格式)返回:compensated: 运动补偿后的图像mask: 补偿区域的掩码avg_dst: 平均运动距离motion_x: x方向平均运动量motion_y: y方向平均运动量homography_matrix: 单应性变换矩阵"""# 设置LK光流参数lk_params = dict(winSize=(15, 15), maxLevel=3,criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 30, 0.003))# 图像预处理和网格点生成width = frame2.shape[1]height = frame2.shape[0]scale = 2 # 放大图像以获得更精确的跟踪# 调整图像大小以提高特征点检测精度frame1_grid = cv2.resize(frame1, (960 * scale, 540 * scale), dst=None, interpolation=cv2.INTER_CUBIC)frame2_grid = cv2.resize(frame2, (960 * scale, 540 * scale), dst=None, interpolation=cv2.INTER_CUBIC)width_grid = frame2_grid.shape[1]height_grid = frame2_grid.shape[0]gridSizeW = 32 * 2 # 网格宽度gridSizeH = 24 * 2 # 网格高度# 生成网格点作为特征点p1 = []grid_numW = int(width_grid / gridSizeW - 1)grid_numH = int(height_grid / gridSizeH - 1)for i in range(grid_numW):for j in range(grid_numH):# 将点放置在每个网格中心point = (np.float32(i * gridSizeW + gridSizeW / 2.0), np.float32(j * gridSizeH + gridSizeH / 2.0))p1.append(point)p1 = np.array(p1)pts_num = grid_numW * grid_numHpts_prev = p1.reshape(pts_num, 1, 2)# 计算光流pts_cur, st, err = cv2.calcOpticalFlowPyrLK(frame1_grid, frame2_grid, pts_prev, None, **lk_params)# 选择跟踪成功的点good_new = pts_cur[st == 1] # 当前帧中的跟踪点good_old = pts_prev[st == 1] # 前一帧中的跟踪点# 计算运动距离和位移motion_distance = []translate_x = []translate_y = []for i, (new, old) in enumerate(zip(good_new, good_old)):a, b = new.ravel()c, d = old.ravel()motion_distance0 = np.sqrt((a - c) * (a - c) + (b - d) * (b - d))# 过滤异常大的运动值if motion_distance0 > 50:continuetranslate_x0 = a - ctranslate_y0 = b - dmotion_distance.append(motion_distance0)translate_x.append(translate_x0)translate_y.append(translate_y0)motion_dist = np.array(motion_distance)motion_x = np.mean(np.array(translate_x)) if translate_x else 0motion_y = np.mean(np.array(translate_y)) if translate_y else 0avg_dst = np.mean(motion_dist) if motion_dist.size > 0 else 0# 计算单应性变换矩阵if len(good_old) < 15:# 点太少时使用近似恒等变换homography_matrix = np.array([[0.999, 0, 0], [0, 0.999, 0], [0, 0, 1]])else:# 使用RANSAC算法估计单应性矩阵homography_matrix, status = cv2.findHomography(good_new, good_old, cv2.RANSAC, 3.0)# 应用单应性变换进行运动补偿compensated = cv2.warpPerspective(frame1, homography_matrix, (width, height), flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP)# 生成掩码以指示变换区域vertex = np.array([[0, 0], [width, 0], [width, height], [0, height]], dtype=np.float32).reshape(-1, 1, 2)homo_inv = np.linalg.inv(homography_matrix)vertex_trans = cv2.perspectiveTransform(vertex, homo_inv)vertex_transformed = np.array(vertex_trans, dtype=np.int32).reshape(1, 4, 2)im = np.zeros(frame1.shape[:2], dtype='uint8')cv2.polylines(im, vertex_transformed, 1, 255)cv2.fillPoly(im, vertex_transformed, 255)mask = 255 - imreturn compensated, mask, avg_dst, motion_x, motion_y, homography_matrixdef FD_mask(lastFrame1, lastFrame2, currentFrame, save_path='data/mask.jpg'):"""使用两帧运动补偿计算前景掩码参数:lastFrame1: 第一参考帧(BGR格式)lastFrame2: 第二参考帧(BGR格式)currentFrame: 当前帧(BGR格式)save_path: 结果掩码保存路径"""# 图像预处理:高斯模糊和灰度转换lastFrame1 = cv2.GaussianBlur(lastFrame1, (11, 11), 0)lastFrame1 = cv2.cvtColor(lastFrame1, cv2.COLOR_BGR2GRAY)lastFrame2 = cv2.GaussianBlur(lastFrame2, (11, 11), 0)lastFrame2 = cv2.cvtColor(lastFrame2, cv2.COLOR_BGR2GRAY)currentFrame = cv2.GaussianBlur(currentFrame, (11, 11), 0)currentFrame = cv2.cvtColor(currentFrame, cv2.COLOR_BGR2GRAY)# 计算第一参考帧到第二参考帧的运动补偿img_compensate1, mask1, avg_dist1, motion_x1, motion_y1, homo_matrix = motion_compensate(lastFrame1, lastFrame2)frameDiff1 = cv2.absdiff(lastFrame2, img_compensate1)# 计算当前帧到第二参考帧的运动补偿img_compensate2, mask2, avg_dist2, motion_x2, motion_y2, homo_matrix2 = motion_compensate(currentFrame, lastFrame2)frameDiff2 = cv2.absdiff(lastFrame2, img_compensate2)# 融合两个差分结果frameDiff = (frameDiff1 + frameDiff2) / 2# 保存结果os.makedirs(os.path.dirname(save_path), exist_ok=True)cv2.imwrite(save_path, frameDiff)print(f'前景掩码已保存至: {save_path}')return frameDiffif __name__ == "__main__":# 示例:加载三帧图像并计算前景掩码# 请确保这些图像存在,或者修改为您自己的图像路径try:lastFrame1 = cv2.imread('data/Test_images/images/phantom05_0600.jpg')lastFrame3 = cv2.imread('data/Test_images/images/phantom05_0602.jpg')currentFrame = cv2.imread('data/Test_images/images/phantom05_0604.jpg')if lastFrame1 is None or lastFrame3 is None or currentFrame is None:print("错误: 无法加载图像,请检查文件路径!")else:FD_mask(lastFrame1, lastFrame3, currentFrame)except Exception as e:print(f"程序执行出错: {e}")