#!/usr/bin/env python # -*- coding: utf-8 -*- import time import argparse import cv2 import numpy as np from ultralytics import YOLO from skimage.metrics import structural_similarity as ssim # 添加YOLOv8模型初始化 yolo_model = YOLO("best.pt") # 可替换为yolov8s/m/l等 yolo_model.to('cuda') # 启用GPU加速 def calculate_en(img): """计算信息熵（处理灰度图）""" hist = cv2.calcHist([img], [0], None, [256], [0, 256]) hist = hist / hist.sum() return -np.sum(hist * np.log2(hist + 1e-10)) def calculate_sf(img): """计算空间频率（处理灰度图）""" rf = np.sqrt(np.mean(np.square(np.diff(img, axis=0)))) cf = np.sqrt(np.mean(np.square(np.diff(img, axis=1)))) return np.sqrt(rf ** 2 + cf ** 2) def calculate_mi(img1, img2): """计算互信息（处理灰度图）""" hist_2d = np.histogram2d(img1.ravel(), img2.ravel(), 256)[0] pxy = hist_2d / hist_2d.sum() px = np.sum(pxy, axis=1) py = np.sum(pxy, axis=0) return np.sum(pxy * np.log2(pxy / (px[:, None] * py[None, :] + 1e-10) + 1e-10)) def calculate_ssim(img1, img2): """计算SSIM（处理灰度图）""" return ssim(img1, img2, data_range=255) # 裁剪线性RGB对比度拉伸：（去掉2%百分位以下的数，去掉98%百分位以上的数，上下百分位数一般相同，并设置输出上下限） def truncated_linear_stretch(image, truncated_value=2, maxout=255, min_out=0): """ :param image: :param truncated_value: :param maxout: :param min_out: :return: """ def gray_process(gray, maxout=maxout, minout=min_out): truncated_down = np.percentile(gray, truncated_value) truncated_up = np.percentile(gray, 100 - truncated_value) gray_new = ((maxout - minout) / (truncated_up - truncated_down)) * gray gray_new[gray_new < minout] = minout gray_new[gray_new > maxout] = maxout return np.uint8(gray_new) (b, g, r) = cv2.split(image) b = gray_process(b) g = gray_process(g) r = gray_process(r) result = cv2.merge((b, g, r)) # 合并每一个通道 return result # RGB图片配准函数，采用白天的可见光与红外灰度图，计算两者Surf共同特征点，之间的仿射矩阵。 def Images_matching(img_base, img_target): """ :param img_base: :param img_target:匹配图像 :return: 返回仿射矩阵 """ start = time.time() orb = cv2.ORB_create() # 对可见光图像进行对比度拉伸 # img_base = truncated_linear_stretch(img_base) img_base = cv2.cvtColor(img_base, cv2.COLOR_BGR2GRAY) sift = cv2.SIFT_create() # 使用sift算子计算特征点和特征点周围的特征向量 st1 = time.time() kp1, des1 = sift.detectAndCompute(img_base, None) # 1136 1136, 64 kp2, des2 = sift.detectAndCompute(img_target, None) en1 = time.time() # print(en1 - st1, "特征提取") # 进行KNN特征匹配 # FLANN_INDEX_KDTREE = 0 # 建立FLANN匹配器的参数 # indexParams = dict(algorithm=FLANN_INDEX_KDTREE, trees=5) # 配置索引，密度树的数量为5 # searchParams = dict(checks=50) # 指定递归次数 # flann = cv2.FlannBasedMatcher(indexParams, searchParams) # 建立匹配器 # matches = flann.knnMatch(des1, des2, k=2) # 得出匹配的关键点 list: 1136 # FLANN_INDEX_KDTREE = 1 # index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5) # search_params = dict(checks=50) # flann = cv2.FlannBasedMatcher(index_params, search_params) # matches = flann.knnMatch(des1, des2, k=2) st2 = time.time() matcher = cv2.BFMatcher() matches = matcher.knnMatch(des1, des2, k=2) de2 = time.time() # print(de2 - st2, "特征匹配") good = [] # 提取优秀的特征点 for m, n in matches: if m.distance < 0.75 * n.distance: # 如果第一个邻近距离比第二个邻近距离的0.7倍小，则保留 good.append(m) # 134 src_pts = np.array([kp1[m.queryIdx].pt for m in good]) # 查询图像的特征描述子索引 # 134, 2 dst_pts = np.array([kp2[m.trainIdx].pt for m in good]) # 训练(模板)图像的特征描述子索引 if len(src_pts) <= 4: print("Not enough matches are found - {}/{}".format(len(good), 4)) return 0, None, 0 else: print(len(dst_pts), len(src_pts), "配准坐标点") H = cv2.findHomography(dst_pts, src_pts, cv2.RANSAC, 4) # 生成变换矩阵 H[0]: 3, 3 H[1]: 134, 1 end = time.time() times = end - start # print("配准时间", times) return 1, H[0], len(dst_pts) def fusions(img_vl, img_inf): """ :param img_vl: 原图像 :param img_inf: 红外图像 :return: """ img_YUV = cv2.cvtColor(img_vl, cv2.COLOR_BGR2YUV) # 如果输入是BGR，需转换 # img_YUV = cv2.cvtColor(img_vl, cv2.COLOR_RGB2YUV) y, u, v = cv2.split(img_YUV) # 分离通道,获取Y通道 Yf = y * 0.5 + img_inf * 0.5 Yf = Yf.astype(np.uint8) fusion = cv2.cvtColor(cv2.merge((Yf, u, v)), cv2.COLOR_YUV2RGB) return fusion def removeBlackBorder(gray): """ 移除缝合后的图像的多余黑边 输入： image：三维numpy矩阵，待处理图像 输出： 裁剪后的图像 """ threshold = 40 # 阈值 nrow = gray.shape[0] # 获取图片尺寸 ncol = gray.shape[1] rowc = gray[:, int(1 / 2 * nrow)] # 无法区分黑色区域超过一半的情况 colc = gray[int(1 / 2 * ncol), :] rowflag = np.argwhere(rowc > threshold) colflag = np.argwhere(colc > threshold) left, bottom, right, top = rowflag[0, 0], colflag[-1, 0], rowflag[-1, 0], colflag[0, 0] # cv2.imshow('name', gray[left:right, top:bottom]) # 效果展示 cv2.waitKey(1) return gray[left:right, top:bottom], left, right, top, bottom def main(matchimg_vi, matchimg_in): """ :param matchimg_vi: 可见光图像 :param matchimg_in: 红外图像 :return: 融合好的图像（带检测结果） """ try: orimg_vi = matchimg_vi orimg_in = matchimg_in h, w = orimg_vi.shape[:2] # 480 640 # (3, 3)//获取对应的配准坐标点 flag, H, dot = Images_matching(matchimg_vi, matchimg_in) if flag == 0: return 0, None, 0, 0.0, 0.0, 0.0, 0.0 else: # 配准处理 matched_ni = cv2.warpPerspective(orimg_in, H, (w, h)) matched_ni, left, right, top, bottom = removeBlackBorder(matched_ni) # 裁剪可见光图像 # fusion = fusions(orimg_vi[left:right, top:bottom], matched_ni) # 不裁剪可见光图像 fusion = fusions(orimg_vi, matched_ni) # 转换为灰度计算指标 fusion_gray = cv2.cvtColor(fusion, cv2.COLOR_RGB2GRAY) cropped_vi_gray = cv2.cvtColor(orimg_vi, cv2.COLOR_BGR2GRAY) matched_ni_gray = matched_ni # 红外图已经是灰度 # 计算指标 en = calculate_en(fusion_gray) sf = calculate_sf(fusion_gray) mi_visible = calculate_mi(fusion_gray, cropped_vi_gray) mi_infrared = calculate_mi(fusion_gray, matched_ni_gray) mi_total = mi_visible + mi_infrared # 添加SSIM容错处理 try: ssim_visible = calculate_ssim(fusion_gray, cropped_vi_gray) ssim_infrared = calculate_ssim(fusion_gray, matched_ni_gray) ssim_avg = (ssim_visible + ssim_infrared) / 2 except Exception as ssim_error: print(f"SSIM计算错误: {ssim_error}") ssim_avg = -1 # 用-1表示计算失败 # YOLOv8目标检测 results = yolo_model(fusion) # 输入融合后的图像 annotated_image = results[0].plot() # 绘制检测框 # 返回带检测结果的图像 return 1, annotated_image, dot, en, sf, mi_total, ssim_avg except Exception as e: print(f"Error in fusion/detection: {e}") return 0, None, 0, 0.0, 0.0, 0.0, 0.0 def parse_args(): # 输入可见光和红外图像路径 visible_image_path = "./test/visible/visibleI0195.jpg" # 可见光图片路径 infrared_image_path = "./test/infrared/infraredI0195.jpg" # 红外图片路径 # 输入可见光和红外视频路径 visible_video_path = "./test/visible.mp4" # 可见光视频路径 infrared_video_path = "./test/infrared.mp4" # 红外视频路径 """解析命令行参数""" parser = argparse.ArgumentParser(description='图像融合与目标检测') parser.add_argument('--mode', type=str, choices=['video', 'image'], default='image', help='输入模式：video（视频流） 或 image（静态图片）') # 区分摄像头或视频文件 parser.add_argument('--source', type=str, choices=['camera', 'file'], help='视频输入类型：camera（摄像头）或 file（视频文件）') # 视频模式参数 parser.add_argument('--video1', type=str, default=visible_video_path, help='可见光视频路径（仅在source=file时需要）') parser.add_argument('--video2', type=str, default=infrared_video_path, help='红外视频路径（仅在source=file时需要）') # 摄像头模式参数 parser.add_argument('--camera_id1', type=int, default=0, help='可见光摄像头ID（仅在source=camera时需要，默认0）') parser.add_argument('--camera_id2', type=int, default=1, help='红外摄像头ID（仅在source=camera时需要，默认1）') parser.add_argument('--output', type=str, default='output.mp4', help='输出视频路径（仅在video模式需要）') # 图片模式参数 parser.add_argument('--visible', type=str, default=visible_image_path, help='可见光图片路径（仅在image模式需要）') parser.add_argument('--infrared', type=str, default=infrared_image_path, help='红外图片路径（仅在image模式需要）') return parser.parse_args() if __name__ == '__main__': time_all = 0 dots = 0 i = 0 args = parse_args() if args.mode == 'video': if args.source == 'file': # ========== 视频流处理模式 ========== if not args.video1 or not args.video2: raise ValueError("视频模式需要指定 --video1 和 --video2 参数") capture = cv2.VideoCapture(args.video2) capture2 = cv2.VideoCapture(args.video1) elif args.source == 'camera': # ========== 摄像头处理模式 ========== capture = cv2.VideoCapture(args.camera_id1) capture2 = cv2.VideoCapture(args.camera_id2) else: raise ValueError("必须指定 --source 参数（camera 或 file）") # 公共视频处理逻辑 fps = capture.get(cv2.CAP_PROP_FPS) if args.source == 'file' else 30 fourcc = cv2.VideoWriter_fourcc(*'XVID') out = cv2.VideoWriter(args.output, fourcc, fps, (640, 480)) while True: ret1, frame_vi = capture.read() # 可见光帧 ret2, frame_ir = capture2.read() # 红外帧 if not ret1 or not ret2: break # 红外图像转灰度 frame_ir_gray = cv2.cvtColor(frame_ir, cv2.COLOR_BGR2GRAY) # 执行融合与检测 flag, fusion, _ = main(frame_vi, frame_ir_gray) if flag == 1: cv2.imshow("Fusion with YOLOv8 Detection", fusion) out.write(fusion) if cv2.waitKey(1) == ord('q'): break # 释放资源 capture.release() capture2.release() out.release() cv2.destroyAllWindows() elif args.mode == 'image': # ========= 图片处理模式 ========== if not args.infrared or not args.visible: raise ValueError("图片模式需要指定 --visible 和 --infrared 参数") # 读取图像 img_visible = cv2.imread(args.visible) img_infrared = cv2.imread(args.infrared) if img_visible is None or img_infrared is None: print("Error: 图片加载失败，请检查路径！") exit() # 转换为灰度图（红外图像处理） img_inf_gray = cv2.cvtColor(img_infrared, cv2.COLOR_BGR2GRAY) # 执行融合与检测 flag, fusion_result, dot, en, sf, mi, ssim_val = main(img_visible, img_inf_gray) if flag == 1: # 展示评价指标 print("\n======== 融合质量评价 ========") print(f"信息熵（EN）: {en:.2f}") print(f"空间频率（SF）: {sf:.2f}") print(f"互信息（MI）: {mi:.2f}") # 条件显示SSIM if ssim_val >= 0: print(f"结构相似性（SSIM）: {ssim_val:.4f}") else: print("结构相似性（SSIM）: 计算失败（已跳过）") print(f"配准点数: {dot}") # 显示并保存结果 # cv2.imshow("Fusion with Detection", fusion_result) cv2.imwrite("output/fusion_result.jpg", fusion_result) # cv2.waitKey(0) # cv2.destroyAllWindows() else: print("融合失败！")