添加三种不同模式

image_fusion/Image_Registration_test.py

@@ -1,10 +1,8 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-# @Time    :
-# @Author  :
-# @File    : Image_Registration_test.py
 
 import time
+import argparse
 
 import cv2
 import numpy as np
@@ -128,7 +126,7 @@ def Images_matching(img_base, img_target):
         print("Not enough matches are found - {}/{}".format(len(good), 4))
         return 0, None, 0
     else:
-        # print(len(dst_pts), len(src_pts), "配准坐标点")
+        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
@@ -246,50 +244,60 @@ if __name__ == '__main__':
     time_all = 0
     dots = 0
     i = 0
-    # fourcc = cv2.VideoWriter_fourcc(*'XVID')
-    # capture = cv2.VideoCapture("video/20190926_141816_1_8/20190926_141816_1_8/infrared.mp4")
-    # capture2 = cv2.VideoCapture("video/20190926_141816_1_8/20190926_141816_1_8/visible.mp4")
-    # fps = capture.get(cv2.CAP_PROP_FPS)
-    # out = cv2.VideoWriter('output2.mp4', fourcc, fps, (640, 480))
-    # # 持续读取摄像头数据
-    # while True:
-    #     read_code, frame = capture.read()  # 红外帧
-    #     read_code2, frame2 = capture2.read()  # 可见光帧
-    #     if not read_code:
-    #         break
-    #     i += 1
-    #     # frame = cv2.resize(frame, (1920, 1080))
-    #     # frame2 = cv2.resize(frame2, (640, 512))
-    #
-    #     # 转换为灰度图（红外图像处理）
-    #     frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-    #
-    #     # 调用main函数进行融合和检测
-    #     flag, fusion, dot = main(frame2, frame_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()
-    # cv2.destroyAllWindows()
-    # ave = time_all / i
-    # print(ave, "平均时间")
-    # cv2.destroyAllWindows()
+    args = parse_args()
     
-    # === 新增静态图片测试代码 ===
-    # 输入可见光和红外图像路径
-    visible_path = "../test_images/visible.jpg"  # 可见光图片路径
-    infrared_path = "../test_images/infrared.jpg"  # 红外图片路径
+    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(visible_path)
-    img_infrared = cv2.imread(infrared_path)
+        img_visible = cv2.imread(args.visible)
+        img_infrared = cv2.imread(args.infrared)
         
         if img_visible is None or img_infrared is None:
             print("Error: 图片加载失败，请检查路径！")