water_turbidity_det/play.py

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()