6

calculator.py

@@ -12,46 +12,31 @@ def set_chinese_font():
     rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei', 'Microsoft YaHei', 'SimSun']
     rcParams['axes.unicode_minus'] = False
 
-def region_ndti(img, mask):
-    """
-
-    :param img: bgr image
-    :param mask:
-    :return:
-    """
+def region_ndti(img, mask, is_night):
+    """对于可见光图像计算ndti，对于红外图像直接取值"""
     if len(img.shape) != 3:
         raise RuntimeError('请输入三通道彩色图像')
 
     if len(mask.shape) != 2:
         raise RuntimeError('请输入单通道掩膜图像')
+
     img = img.astype(np.float32)
     mask = mask.astype(np.float32)
-    #b = img[:,:,0]
+    # 取通道
     g = img[:,:,1]
     r = img[:,:,2]
-
-    ndti = (r - g) / (r + g + 1e-6)
+    # 判断是否为红外
+    if is_night: # 红外直接取值
+        ndti = r
+    else:  # 可见光计算ndti
+        ndti = (r - g) / (r + g + 1e-6)
+    # 掩膜归一化
     if np.max(mask) > 2:
-        mask /= 255
+        mask /= 255.
+    # 仅保留掩膜区域的ndti
     ndti *= mask
     return ndti
 
-def region_night(img, mask):
-    if len(img.shape) != 3:
-        raise RuntimeError('请输入三通道彩色图像')
-
-    if len(mask.shape) != 2:
-        raise RuntimeError('请输入单通道掩膜图像')
-    img = img.astype(np.float32)
-    mask = mask.astype(np.float32)
-    ir = img[:,:,2]
-
-    if np.max(mask) > 2:
-        mask /= 255
-    ir *= mask
-    return ir
-
-
 def plot_histogram_opencv(matrix, start, end, step, title):
     """使用OpenCV绘制直方图，避免GUI冲突"""
     # 计算直方图
@@ -312,25 +297,25 @@ def main(video_dir, mask_dir, id):
             # 处理当前帧
             # 缩放
             frame = cv2.resize(frame, (frame.shape[1] // scale, frame.shape[0] // scale))
+            # cv2.imshow('original', frame)
             # 滤波
             # frame = cv2.GaussianBlur(frame, (5, 5), 1.5)
             # 计算掩膜区域的NDTI值
-            if is_night:
-                roi_ndti = region_night(frame, mask)
-            else:
-                roi_ndti = region_ndti(frame, mask)
+            roi_ndti = region_ndti(frame, mask, is_night)
             # 给ndti着色
             color_ndti = colorful_ndti_matrix(roi_ndti)
 
             # 绘制直方图
             roi_index = mask > 0
             roi_ndti_fla = roi_ndti[roi_index]  # 仅保留感兴趣区的计算结果
-            plot_histogram_seaborn_simple(roi_ndti_fla,-1,1,0.01, str(id))
+            # plot_histogram_seaborn_simple(roi_ndti_fla,-1,1,0.01, str(id))
 
             # 打印统计信息
             # 3σ原则
             mean_value  = np.mean(roi_ndti_fla)
+            print('调整前平均值:', mean_value)
             std_value = np.std(roi_ndti_fla)
+            print('调整前平均值:', std_value)
             min_value = np.min(roi_ndti_fla)
             max_value = np.max(roi_ndti_fla)
 
@@ -383,39 +368,41 @@ def main(video_dir, mask_dir, id):
 if __name__ == '__main__':
     set_chinese_font()
 
-    # single_img(r'D:\code\water_turbidity_det\data\day_202511211129\1_img_202511211412.jpg',
+    # single_img(r'D:\code\water_turbidity_det\data\day_202511211129\1_device_capture.jpg',
     #             r'D:\code\water_turbidity_det\draw_mask\mask\1_main_20251119102036_@1000000.png')
     # pass
-    path1 = r'D:\code\water_turbidity_det\data\day_202511211129\1_video_202511211128.dav'
-    path2 = r'D:\code\water_turbidity_det\data\day_202511211129\2_video_202511211128.dav'
-    path3 = r'D:\code\water_turbidity_det\data\day_202511211129\3_video_202511211128.dav'
-    path4 = r'D:\code\water_turbidity_det\data\day_202511211129\4_video_202511211128.dav'
-
-    path1 = r'D:\code\water_turbidity_det\data\night\1_video_2025_1120.dav'
-    path2 = r'D:\code\water_turbidity_det\data\night\2_video_20251120_1801.dav'
-    path3 = r'D:\code\water_turbidity_det\data\night\3_video_20251120_1759.dav'
-    path4 = r'D:\code\water_turbidity_det\data\night\4_video_20251120_1800.dav'
-
-    t1 = threading.Thread(target=main, kwargs={'video_dir': path1,
-                                               'mask_dir': r'D:\code\water_turbidity_det\draw_mask\mask\1_main_20251119102036_@1000000.png',
-                                               'id':1})
-    t2 = threading.Thread(target=main, kwargs={'video_dir': path2,
-                                               'mask_dir': r'D:\code\water_turbidity_det\draw_mask\mask\2_main_20251119102038_@1000000.png',
-                                               'id':2})
-    t3 = threading.Thread(target=main, kwargs={'video_dir': path3,
-                                               'mask_dir': r'D:\code\water_turbidity_det\draw_mask\mask\3_main_20251119102042_@1000000.png',
-                                               'id':3})
-    t4 = threading.Thread(target=main, kwargs={'video_dir': path4,
-                                               'mask_dir': r'D:\code\water_turbidity_det\draw_mask\mask\4_main_20251119102044_@1000000.png',
-                                               'id':4})
-    # threads = [t1, t2, t3, t4]
-    threads = [t4]
-    for t in threads:
-        t.start()
-    for t in threads:
-        if t.is_alive():
-            t.join()
-
+    # path1 = r'D:\code\water_turbidity_det\data\day_202511211129\1_video_202511211128.dav'
+    # path2 = r'D:\code\water_turbidity_det\data\day_202511211129\2_video_202511211128.dav'
+    # path3 = r'D:\code\water_turbidity_det\data\day_202511211129\3_video_202511211128.dav'
+    # path4 = r'D:\code\water_turbidity_det\data\day_202511211129\4_video_202511211128.dav'
+    #
+    # path1 = r'D:\code\water_turbidity_det\data\night\1_video_2025_1120.dav'
+    # path2 = r'D:\code\water_turbidity_det\data\night\2_video_20251120_1801.dav'
+    # path3 = r'D:\code\water_turbidity_det\data\night\3_video_20251120_1759.dav'
+    # path4 = r'D:\code\water_turbidity_det\data\night\4_video_20251120_1800.dav'
+    #
+    # t1 = threading.Thread(target=main, kwargs={'video_dir': path1,
+    #                                            'mask_dir': r'D:\code\water_turbidity_det\draw_mask\mask\1_main_20251119102036_@1000000.png',
+    #                                            'id':1})
+    # t2 = threading.Thread(target=main, kwargs={'video_dir': path2,
+    #                                            'mask_dir': r'D:\code\water_turbidity_det\draw_mask\mask\2_main_20251119102038_@1000000.png',
+    #                                            'id':2})
+    # t3 = threading.Thread(target=main, kwargs={'video_dir': path3,
+    #                                            'mask_dir': r'D:\code\water_turbidity_det\draw_mask\mask\3_main_20251119102042_@1000000.png',
+    #                                            'id':3})
+    # t4 = threading.Thread(target=main, kwargs={'video_dir': path4,
+    #                                            'mask_dir': r'D:\code\water_turbidity_det\draw_mask\mask\4_main_20251119102044_@1000000.png',
+    #                                            'id':4})
+    # # threads = [t1, t2, t3, t4]
+    # threads = [t4]
+    # for t in threads:
+    #     t.start()
+    # for t in threads:
+    #     if t.is_alive():
+    #         t.join()
+    main(video_dir=r'D:\code\water_turbidity_det\data\records_202512031553\video4_20251129053218_20251129060528.mp4',
+         mask_dir=r'D:\code\water_turbidity_det\draw_mask\mask\4_device_capture.png',
+         id=4)
     cv2.destroyAllWindows()
 
 

draw_mask/draw.py

@@ -77,4 +77,4 @@ class DrawRectangle:
 
 if __name__ == '__main__':
     drawer = DrawRectangle(2)
-    drawer.draw(r"D:\code\water_turbidity_det\data\4_main_20251119102044_@1000000.jpg")
+    drawer.draw(r"D:\code\water_turbidity_det\draw_mask\mask\4_device_capture.jpg")

play.py

@@ -0,0 +1,68 @@
+import cv2
+import numpy as np
+from calculator import region_ndti, colorful_ndti_matrix
+
+
+def main(video_path, video_mask_path, video_id):
+
+    records = []
+
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        print("错误：无法打开视频文件。")
+        print("可能的原因：文件路径错误，或系统缺少必要的解码器。")
+        exit()
+    scale = 4
+    mask = cv2.imread(video_mask_path, cv2.IMREAD_GRAYSCALE)
+    if mask is None:
+        raise RuntimeError('mask open failed')
+    mask = cv2.resize(mask, (mask.shape[1] // scale, mask.shape[0] // scale))
+
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    print("fps:", fps)
+    while True:
+
+        ret, frame = cap.read()  # ret指示是否成功读取帧，frame是图像数据
+        if not ret:
+            print("视频播放完毕。")
+            break
+        # 判断是否为夜间模式
+        if np.mean(np.abs(frame[:, :, 0] - frame[:, :, 1])) < 0.1:
+            is_night = True
+        else:
+            is_night = False
+        frame = cv2.resize(frame, (frame.shape[1] // scale, frame.shape[0] // scale))
+
+        # 计算掩膜区域的NDTI值
+        roi_ndti = region_ndti(frame, mask, is_night)
+
+
+        # 给ndti着色
+        color_ndti = colorful_ndti_matrix(roi_ndti)
+
+
+        # 仅保留感兴趣区的计算结果
+        roi_index = mask > 0
+        roi_ndti_fla = roi_ndti[roi_index]
+        # 3σ原则
+        # mean_value = np.mean(roi_ndti_fla)
+        # std_value = np.std(roi_ndti_fla)
+        # up_limit = mean_value + 3 * std_value
+        # down_limit = mean_value - 3 * std_value
+        # roi_ndti_fla = roi_ndti_fla[(up_limit >= roi_ndti_fla) & (roi_ndti_fla >= down_limit)]
+
+        text = f'pixs:{int(np.sum(roi_ndti_fla))} mean:{roi_ndti_fla.mean():.3f} std:{roi_ndti_fla.std():.3f} {'night' if is_night else 'day'}'
+        cv2.putText(frame, text, (10, 25), cv2.FONT_HERSHEY_DUPLEX, 0.6, (0, 255, 0), 2, cv2.LINE_AA)
+        cv2.imshow(f'original_{video_id}', frame)
+        cv2.waitKey(30)
+
+        # 退出逻辑
+        key = cv2.waitKey(10) & 0xFF
+        if key  == ord('q'):
+            break
+
+if __name__ == '__main__':
+    main(video_path=r'D:\code\water_turbidity_det\data\day_202511211129\1_video_202511211128.dav',
+         video_mask_path=r'D:\code\water_turbidity_det\draw_mask\mask\1_device_capture.png',
+         video_id=1)
+    cv2.destroyAllWindows()