water_turbidity_det/calculator.py

import cv2
import numpy as np
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt
from  matplotlib import rcParams
import seaborn as sns
import threading


def set_chinese_font():
    rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei', 'Microsoft YaHei', 'SimSun']
    rcParams['axes.unicode_minus'] = False

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)
    # 取通道
    g = img[:,:,1]
    r = img[:,:,2]
    # 判断是否为红外
    if is_night: # 红外直接取值
        ndti = r
    else:  # 可见光计算ndti
        ndti = (r - g) / (r + g + 1e-6)
    # 掩膜归一化
    if np.max(mask) > 2:
        mask /= 255.
    # 仅保留掩膜区域的ndti
    ndti *= mask
    return ndti

def plot_histogram_opencv(matrix, start, end, step, title):
    """使用OpenCV绘制直方图，避免GUI冲突"""
    # 计算直方图
    num_bins = int((end - start) / step)
    hist, bins = np.histogram(matrix.ravel(), bins=num_bins, range=(start, end))

    # 创建直方图图像（增加高度以便显示标签）
    hist_img = np.zeros((450, 600, 3), dtype=np.uint8)
    hist_img.fill(255)  # 白色背景

    # 归一化直方图：确保高度在合理范围内
    if len(hist) > 0 and np.max(hist) > 0:
        # 保留10像素的边距
        hist_normalized = cv2.normalize(hist, None, 0, 350, cv2.NORM_MINMAX)
    else:
        hist_normalized = np.zeros_like(hist)

    # 计算合理的矩形宽度（确保至少1像素宽）
    if len(hist) > 0:
        bin_width = max(1, 600 // len(hist))  # 确保最小宽度为1
    else:
        bin_width = 1

    # 绘制直方图矩形
    for i in range(len(hist_normalized)):
        if i >= 600:  # 防止索引越界
            break

        x1 = i * bin_width
        x2 = min(x1 + bin_width, 599)  # 确保不超出图像边界
        y1 = 400 - int(hist_normalized[i])  # 从底部开始计算高度
        y2 = 400

        # 只绘制有高度的矩形
        if y1 < y2:
            cv2.rectangle(hist_img, (x1, y1), (x2, y2), (0, 0, 255), -1)

    # 添加坐标轴
    cv2.line(hist_img, (50, 400), (550, 400), (0, 0, 0), 2)  # x轴
    cv2.line(hist_img, (50, 400), (50, 50), (0, 0, 0), 2)  # y轴

    # 添加标题和标签（调整位置避免被裁剪）
    cv2.putText(hist_img, title, (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2)
    cv2.putText(hist_img, 'gray', (280, 430),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
    cv2.putText(hist_img, 'fre', (5, 200),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)

    return hist_img

def plot_histogram_seaborn_simple(matrix, start, end, step, title):
    """
    使用Seaborn的histplot函数直接绘制（最简单的方法）
    """
    # 解决中文显示问题
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False

    plt.figure(figsize=(12, 6))

    # 直接使用Seaborn的histplot，设置bins参数[3](@ref)
    num_bins = int((end - start) / step)
    sns.histplot(matrix.ravel(), bins=num_bins, kde=True,
                 color='skyblue', alpha=0.7,
                 edgecolor='white', linewidth=0.5)

    plt.title(title+ '_histplot')
    plt.xlabel('灰度值')
    plt.ylabel('像素频率')
    plt.tight_layout()
    # plt.show()
    plt.savefig('temp_' + title +'.png')  # 保存到文件而不是显示
    plt.close()  # 关闭图形，释放资源

    # 如果需要显示，可以用cv2读取并显示
    hist_img = cv2.imread('temp_' + title +'.png')
    cv2.imshow('Histogram '+title, hist_img)

def plot_histogram_seaborn(matrix, hist, bin_edges, start, end, step):
    """
    使用Seaborn绘制更美观的直方图
    """
    # 设置Seaborn样式
    sns.set_style("whitegrid")
    plt.figure(figsize=(12, 6))

    # 将数据转换为适合Seaborn的格式
    flattened_data = matrix.ravel()

    # 使用Seaborn的histplot（会自动计算直方图，但我们用自定义的）
    # 这里我们手动绘制以确保使用我们的自定义bins
    bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2

    # 创建条形图
    bars = plt.bar(bin_centers, hist, width=step * 0.8,
                   alpha=0.7, color=sns.color_palette("viridis", 1)[0],
                   edgecolor='white', linewidth=0.5)

    # 添加KDE曲线（核密度估计）
    sns.kdeplot(flattened_data, color='red', linewidth=2, label='密度曲线')

    plt.title(f'灰度分布直方图', fontsize=16)
    plt.xlabel('灰度值', fontsize=12)
    plt.ylabel('像素频率', fontsize=12)
    plt.legend()
    plt.tight_layout()
    plt.show()



def colorful_ndti_matrix(ndti_matrix):
    """给NDTI矩阵着色"""
    # 存放着色后的NDTI矩阵
    color_ndti_matrix = np.zeros((ndti_matrix.shape[0], ndti_matrix.shape[1], 3), dtype=np.uint8)
    negative_mask = ndti_matrix < 0
    positive_mask = ndti_matrix > 0
    # 处理负值
    if np.any(negative_mask):
        blue_intensity = ndti_matrix.copy()
        blue_intensity[positive_mask] = 0
        blue_intensity[negative_mask] *= -255.
        blue_intensity = np.clip(blue_intensity, 0, 255).astype(np.uint8)
        color_ndti_matrix[:, :,0] = blue_intensity
    # 处理正值区域（红色渐变）
    if np.any(positive_mask):
        # 将0到+1映射到0-255的红色强度
        red_intensity = ndti_matrix.copy()
        red_intensity[negative_mask] = 0
        red_intensity[positive_mask] *= 255
        red_intensity = np.clip(red_intensity, 0, 255).astype(np.uint8)
        color_ndti_matrix[:, :, 2] = red_intensity

    # 返回着色后的ndti矩阵
    return color_ndti_matrix

def img_add(img1, img2, img1_w=1, img2_w=0, gamma=0):

    if len(img1.shape) != len(img2.shape):
        raise ValueError('img1 and img2 must have the same shape')

    # 设置权重参数（透明度）
    alpha = img1_w  # 第一张图像的权重
    beta = img2_w # 第二张图像的权重
    gamma =gamma  # 亮度调节参数

    # 执行加权叠加
    result = cv2.addWeighted(img1, alpha, img2, beta, gamma)
    return result
def callback(event, x, y, flags, param):
    if event == cv2.EVENT_LBUTTONDOWN:
        is_night_ = param['is_night']
        frame = param['Image']
        w_name = param['window_name']
        b = frame[:,:,0].astype(np.float32)
        g = frame[:,:,1].astype(np.float32)
        r = frame[:,:,2].astype(np.float32)
        if is_night_:
            ndti_value = r[y, x]
        else:
            ndti_value = (r[y,x] - g[y,x]) / (r[y,x] + g[y,x])
        cv2.putText(frame, f'{ndti_value:.2f}', (x, y), cv2.FONT_HERSHEY_DUPLEX, 0.6, (0, 255, 0), 2, cv2.LINE_AA)
        cv2.putText(frame, f'paused', (10, 45), cv2.FONT_HERSHEY_DUPLEX, 0.6, (0, 255, 0), 2, cv2.LINE_AA)
        cv2.imshow(w_name, frame)
def single_img(img_path,mask_path, id=0):

    frame = cv2.imread(img_path)
    if frame is None:
        raise RuntimeError('img open failed')
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    if mask is None:
        raise RuntimeError('mask open failed')
    scale = 4 if id != 3 else 2
    mask = cv2.resize(mask, (mask.shape[1] // scale, mask.shape[0] // scale))
    frame = cv2.resize(frame, (frame.shape[1] // scale, frame.shape[0] // scale))
    roi_ndti = region_ndti(frame, mask)
    # 给ndti着色
    color_ndti = colorful_ndti_matrix(roi_ndti)

    # 绘制直方图
    roi_index = mask > 0
    roi_ndti_fla = roi_ndti[roi_index]  # 仅保留感兴趣区的计算结果
    hist_img = plot_histogram_opencv(roi_ndti[roi_index], -1, 1, 0.01, f'NDTI Histogram {id}')

    # 打印统计信息
    # 3σ原则
    mean_value = np.mean(roi_ndti_fla)
    std_value = np.std(roi_ndti_fla)
    min_value = np.min(roi_ndti_fla)
    max_value = np.max(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))} adj_mean:{roi_ndti_fla.mean():.3f} adj_std:{roi_ndti_fla.std():.3f}'
    cv2.putText(frame, text, (10, 25), cv2.FONT_HERSHEY_DUPLEX, 0.6, (0, 255, 0), 2, cv2.LINE_AA)
    print(f"""统计信息:
    总像素数: {np.sum(roi_ndti_fla):,}
    灰度范围: {min_value:.1f} - {max_value:.1f}
    平均值: {mean_value:.2f}
    标准差: {std_value:.2f}
    调整平均值：{roi_ndti_fla.mean():.2f}
    调整标准差：{roi_ndti_fla.std():.2f}
    """
          )
    # 显示当前帧处理结果
    cv2.imshow('original', frame)  # 原图
    cv2.imshow('mask'+ str(id), mask)   # 掩膜
    roi_ndti = np.abs(roi_ndti*255.).astype(np.uint8)
    cv2.imshow('ndti'+ str(id), roi_ndti)  # ndti黑白强度
    cv2.imshow('color_ndti'+ str(id), color_ndti)  # # ndti彩色强度
    # 图像叠加
    add_img = img_add(frame, color_ndti)
    cv2.imshow('add_ori_ndti' + str(id), add_img)
    param = {'Image': frame}
    cv2.setMouseCallback('original', callback, param=param)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

def main(video_dir, mask_dir, id):
    # 视频分析浊度
    video_path = video_dir
    # 加载掩膜
    mask_path = mask_dir
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    if mask is None:
        raise RuntimeError('mask open failed')
    # 除数倍率
    scale = 4 if id != 3 else 2

    mask = cv2.resize(mask, (mask.shape[1] // scale, mask.shape[0] // scale))
    # 视频读取
    cap = cv2.VideoCapture(video_path)
    pause = False
    # 检查视频是否成功打开
    if not cap.isOpened():
        print("错误：无法打开视频文件。")
        print("可能的原因：文件路径错误，或系统缺少必要的解码器。")
        exit()
    else:
        ret, frame = cap.read()

    # 成功打开后，逐帧读取和显示视频
    global_img = [None]
    is_night = False
    while True:
        if not pause:
            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
            # 处理当前帧
            # 缩放
            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值
            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))

            # 打印统计信息
            # 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)

            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))} adj_mean:{roi_ndti_fla.mean():.3f} adj_std:{roi_ndti_fla.std():.3f}'
            cv2.putText(frame, text, (10, 25), cv2.FONT_HERSHEY_DUPLEX, 0.6, (0, 255, 0), 2, cv2.LINE_AA)
            # print(f"""统计信息:
            # 总像素数: {np.sum(roi_ndti_fla):,}
            # 灰度范围: {min_value:.1f} - {max_value:.1f}
            # 平均值: {mean_value:.2f}
            # 标准差: {std_value:.2f}
            # 调整平均值：{roi_ndti_fla.mean():.2f}
            # 调整标准差：{roi_ndti_fla.std():.2f}
            # """
            #       )
            # 显示当前帧处理结果
            #cv2.imshow('original', frame)  # 原图
            #cv2.imshow('mask'+ str(id), mask)   # 掩膜
            #roi_ndti = np.abs(roi_ndti*255.).astype(np.uint8)
            #cv2.imshow('ndti'+ str(id), roi_ndti)  # ndti黑白强度
            #cv2.imshow('color_ndti'+ str(id), color_ndti)  # # ndti彩色强度
            # 图像叠加
            add_img = img_add(frame, color_ndti)
            global_img[0] = add_img
            cv2.imshow('add_ori_ndti'+ str(id), add_img)
            #cv2.imshow('Histogram', hist_img)

        # 播放帧率为25FPS，如果期间按下'q'键则退出循环
        key = cv2.waitKey(500) & 0xFF
        if key == ord(' '):
            pause = not pause
            status = "已暂停" if pause else "播放中"
            print(f"{id} 状态: {status}")

        if key  == ord('q'):
            break
        if pause:

            if global_img[0] is not None:
                param = {'Image': global_img[0],'window_name': 'add_ori_ndti' + str(id), 'is_night': is_night}
                cv2.setMouseCallback('add_ori_ndti' + str(id), callback, param=param)

    # 释放资源并关闭所有窗口
    cap.release()


if __name__ == '__main__':
    set_chinese_font()

    # 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_20251120.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'/video/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()