update

.env

@@ -3,7 +3,7 @@ PATCH_WIDTH=256
 # 图块宽度
 PATCH_HEIGHT=256
 # 测试精度时置信度阈值
-CONFIDENCE_THRESHOLD=0.6
+CONFIDENCE_THRESHOLD=0.90
 # 图片输入大小
 IMG_INPUT_SIZE=256
 # 工作线程数

bmodel_application.py

@@ -30,15 +30,19 @@ class Predictor:
         # 输入的大图高度和宽度
         self.input_image_h = 0
         self.input_image_w = 0
-        self.print_network_info()
         # 报警的置信度阈值
-        self.confidence_threshold = 0.6
+        self.confidence_threshold = 0.85
         # 存储连续的帧的检测结果
         self.continuous_detection_result = []
         # 连续浑浊的patch计数器
         self.continuous_counter_mat = np.array(0)
         # 判定有浑浊水体时的连续帧数量
-        self.max_continuous_frames = 15
+        self.max_continuous_frames = 50
+        # 比例判定的起始帧数
+        self.start_frame_num = self.max_continuous_frames
+        # 比例判定的阈值
+        self.ratio_threshold = 0.85
+        self.print_network_info()
 
     def __call__(self, img) -> bool:
         # 预处理，获取输入图像的patch序列和左上角坐标
@@ -76,7 +80,12 @@ class Predictor:
             'Height': self.net_h,
             'Width': self.net_w,
             'Mean': self.mean,
-            'Std': self.std
+            'Std': self.std,
+            'Input Image Size': [self.input_image_h, self.input_image_w],
+            'Confidence Threshold': self.confidence_threshold,
+            'Max Continuous Frames': self.max_continuous_frames,
+            'Ratio Threshold': self.ratio_threshold,
+            'Start Frame Num': self.start_frame_num
         }
 
         print("=" * 50)
@@ -154,7 +163,7 @@ class Predictor:
         # 第三层 时间滤波
         self.continuous_detection_result.append(predicted_class)
         flag = self.update_continuous_counter(predicted_class) # 连续帧滤波
-        if len(self.continuous_detection_result) >= 20:
+        if len(self.continuous_detection_result) >= self.start_frame_num:
             flag = self.discriminate_ratio() and flag
             print(f'是否存在浑浊水体，综合判别结果：{flag}')
             self.continuous_detection_result.pop(0)
@@ -176,7 +185,7 @@ class Predictor:
         positive_index = np.array(pre_class_arr, dtype=np.int32) > 0
         negative_index = np.array(pre_class_arr, dtype=np.int32) == 0
         # 给负样本区域置零
-        self.continuous_counter_mat[negative_index] -= 1
+        self.continuous_counter_mat[negative_index] -= 3
         # 给正样本区域加1
         self.continuous_counter_mat[positive_index] += 1
         # 保证不出现负数
@@ -192,7 +201,7 @@ class Predictor:
         # 方式一：60%以上的帧存在浑浊水体
         water_pre_arr = np.array(water_pre_list, dtype=np.float32)
         water_pre_arr_sum = np.sum(water_pre_arr, axis=0)
-        bad_water = np.array(water_pre_arr_sum >= 0.6 * len(water_pre_list), dtype=np.int32)
+        bad_water = np.array(water_pre_arr_sum >= self.ratio_threshold * len(water_pre_list), dtype=np.int32)
         bad_flag = bool(np.sum(bad_water, dtype=np.int32) > 2)  # 大于两个patch符合要求才可以
         # print('比例信息：',water_pre_arr_sum )
         print(f'浑浊比例判别：该时间段是否存在浑浊水体：{bad_flag}')
@@ -215,7 +224,7 @@ class Predictor:
         return new_predicted_conf_mat.flatten(), new_predicted_class_mat.flatten()
 def argsparser():
     parser = argparse.ArgumentParser(prog=__file__)
-    parser.add_argument('--input','-i', type=str, default=r'4_video_20251223163145', help='path of input, must be image directory')
+    parser.add_argument('--input','-i', type=str, default=r'./4_video_20251223163145', help='path of input, must be image directory')
     parser.add_argument('--bmodel','-b', type=str, default='./shufflenet_f32.bmodel', help='path of bmodel')
     parser.add_argument('--dev_id','-d', type=int, default=0, help='tpu id')
     args = parser.parse_args()

labelme/crop_patch.py

@@ -15,9 +15,9 @@ patch_h = int(os.getenv('PATCH_HEIGHT', 256))
 
 def main():
     # TODO：需要修改为标注好的图片路径
-    input_path = r'D:\code\water_turbidity_det\frame_data\train\20260107\4_ch26_20260105141754'
+    input_path = r'D:\code\water_turbidity_det\label_data_tem\4_ch26_20260113145353_1'
     # TODO: 需要修改为保存patch的根目录
-    output_path_root = r'D:\code\water_turbidity_det\label_data_tem\train'
+    output_path_root = r'D:\code\water_turbidity_det\20260115\test'
 
     # 读取标注文件
     label_path = os.path.join(input_path, 'label.txt')

labelme/fixed_label.py

@@ -164,7 +164,7 @@ def main():
     global patch_h
     global scale
     # TODO: 需要更改为准备标注的图像路径，使用当前目录下的000000.jpg，结果保存在当前目录下label.txt
-    img_path = r'D:\code\water_turbidity_det\frame_data\4_ch26_20260105141754\000000.jpg'
+    img_path = r'D:\code\water_turbidity_det\label_data_tem\4_ch26_20260113145353_1\000300.jpg'
     play_video(img_path)
     img = cv2.imread(img_path)
     draw_tool.set_path(img_path)

labelme/random_del.py

@@ -4,7 +4,7 @@ import random
 
 def main():
     # TODO：需要修改图像路径
-    path = r'D:\code\water_turbidity_det\frame_data\1_ch25_20260105084338'
+    path = r'D:\code\water_turbidity_det\label_data_tem\2_ch52_20260113011503_0'
     del_rate = 0.3
     img_path = [i for i in os.listdir(path) if i.split('.')[-1] in ['jpg', 'png'] ]
     is_del = input(f'{path}路径内共有{len(img_path)}张图片，是否删除{del_rate}比例的图片？(y/n): ')

labelme/video_depart.py

@@ -5,8 +5,8 @@ import shutil
 def main():
     # 视频路径
     # TODO: 修改视频路径为自己的视频路径,每次指定一个视频
-    path = r'D:\code\water_turbidity_det\video\20260107\4_ch26_20260105141754.mp4'
-    output_rootpath = r'D:\code\water_turbidity_det\frame_data'  # 输出路径的根目录
+    path = r'/video/20260115/4_ch26_20260113145353.mp4'
+    output_rootpath = r'D:\code\water_turbidity_det\tem_test'  # 输出路径的根目录
     # 抽帧间隔
     interval = 20
     # 我们将图像输出到根目录下的子目录中，子目录和视频名称相同

pth2onnx.py

@@ -10,11 +10,11 @@ if __name__ == '__main__':
 
     # 载入模型框架
     # model = SimpleModel()
-    # model = resnet50(pretrained=False)
-    model = shufflenet_v2_x1_0()
+    model = resnet50(pretrained=False)
+    # model = shufflenet_v2_x1_0()
     # model = shufflenet_v2_x2_0()
     # model = squeezenet1_0()
-    model_name = "shufflenet"
+    model_name = "resnet50"
     if model_name == "squeezenet":
         # 获取SqueezeNet的最后一个卷积层的输入通道数
         final_conv_in_channels = model.classifier[1].in_channels
@@ -25,7 +25,7 @@ if __name__ == '__main__':
             nn.ReLU(inplace=True),
             nn.AdaptiveAvgPool2d((1, 1))
         )
-    if model_name == "shufflenet":
+    if model_name == "shufflenet" or model_name == "shufflenet-x2" or model_name == "resnet50":
         model.fc = nn.Linear(int(model.fc.in_features), 2, bias=True)
     model.load_state_dict(torch.load(rf'./{model_name}.pth')) # xxx.pth表示.pth文件, 这一步载入模型权重
     print("加载模型成功")
@@ -42,7 +42,6 @@ if __name__ == '__main__':
                       )  # xxx.onnx表示.onnx文件, 这一步导出为onnx模型, 并不做任何算子融合操作。
 
     # 验证模型
-
     onnx_model = onnx.load(f"{model_name}.onnx")  # 使用不同变量名
     onnx.checker.check_model(onnx_model)  # 验证模型完整性
     # 使用ONNX Runtime进行推理

train.py

@@ -83,7 +83,9 @@ class Trainer:
 
         # 创建数据集对象
         self.train_dataset = ImageFolder(root=train_dir, transform=self.train_transforms)
+        print(f"训练样本数量: {len(self.train_dataset)}")
         self.val_dataset = ImageFolder(root=val_dir, transform=self.val_transforms)
+        print(f"验证样本数量: {len(self.val_dataset)}")
         # 创建数据加载器 (Windows环境下设置num_workers=0避免多进程问题)
         self.train_loader = DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.workers)  # 可迭代对象，返回(inputs_tensor, label)
         self.val_loader = DataLoader(self.val_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.workers)

video_test.py

@@ -30,6 +30,15 @@ class Predictor:
         print(f"当前设备: {self.device}")
         self.model = self.load_model()
 
+        # 判定有浑浊水体时的连续帧数量
+        self.max_continuous_frames = 100
+        # 比例判定的起始帧数
+        self.start_frame_num = self.max_continuous_frames
+        # 比例判定的阈值
+        self.ratio_threshold = 0.90
+        # 报警的置信度阈值
+        self.confidence_threshold = 0.90
+
     def load_model(self):
         return load_model(name=self.model_name, num_classes=self.num_classes, weights_path=self.weights_path, device=self.device)
 
@@ -154,28 +163,28 @@ def fileter_prediction(predicted_class, confidence, pre_rows, pre_cols, filter_d
                 new_predicted_conf_mat[i, j] = 1.0
     return new_predicted_conf_mat.flatten(), new_predicted_class_mat.flatten()
 
-def discriminate_ratio(water_pre_list:list):
+def discriminate_ratio(water_pre_list:list, right_ratio:float):
     # 方式一：60%以上的帧存在浑浊水体
     water_pre_arr = np.array(water_pre_list, dtype=np.float32)
     water_pre_arr_sum = np.sum(water_pre_arr, axis=0)
-    bad_water = np.array(water_pre_arr_sum >= 0.6 * len(water_pre_list), dtype=np.int32)
+    bad_water = np.array(water_pre_arr_sum >= right_ratio * len(water_pre_list), dtype=np.int32)
     bad_flag = bool(np.sum(bad_water, dtype=np.int32) > 2)  # 大于两个patch符合要求才可以
     print(f'浑浊比例判别：该时间段是否存在浑浊水体：{bad_flag}')
     return bad_flag
 
 
-def discriminate_count(pre_class_arr, continuous_count_mat):
+def discriminate_count(pre_class_arr, continuous_count_mat,max_continuous_frames):
     """连续帧判别"""
     positive_index = np.array(pre_class_arr,dtype=np.int32) > 0
     negative_index = np.array(pre_class_arr,dtype=np.int32) == 0
     # 给负样本区域置零
-    continuous_count_mat[negative_index] -= 1
+    continuous_count_mat[negative_index] -= 3
     # 给正样本区域加1
     continuous_count_mat[positive_index] += 1
     # 保证不出现负数
     continuous_count_mat[continuous_count_mat<0] = 0
     # 判断浑浊
-    bad_flag = bool(np.sum(continuous_count_mat > 15) > 2)
+    bad_flag = bool(np.sum(continuous_count_mat > max_continuous_frames) > 2)
     print(f'连续帧方式：该时间段是否存在浑浊水体：{bad_flag}')
     return bad_flag
 
@@ -183,10 +192,10 @@ def main():
 
     # 初始化模型实例
     # TODO:修改模型网络名称/模型权重路径/视频路径
-    predictor = Predictor(model_name='shufflenet',
+    predictor = Predictor(model_name='shufflenet-x2',
                           weights_path=r'./shufflenet.pth',
                           num_classes=2)
-    input_path = r'D:\code\water_turbidity_det\frame_data\1video_20251229124533_hunzhuo'
+    input_path = r'D:\code\water_turbidity_det\tem_test\4_ch26_20260113075612_1'
     # 预处理图像
     all_imgs = os.listdir(input_path)
     all_imgs = [os.path.join(input_path, p) for p in all_imgs if p.split('.')[-1] in ['jpg', 'png']]
@@ -235,10 +244,10 @@ def main():
         cv2.waitKey(25)
         water_pre_list.append(new_predicted_class)
         # 方式2判别
-        flag = discriminate_count(new_predicted_class, continuous_count_mat)
+        flag = discriminate_count(new_predicted_class, continuous_count_mat, predictor.max_continuous_frames)
         # 方式1判别
-        if len(water_pre_list) > 25:
-            flag = discriminate_ratio(water_pre_list) and flag
+        if len(water_pre_list) > predictor.start_frame_num:
+            flag = discriminate_ratio(water_pre_list, predictor.ratio_threshold) and flag
         print('综合判别结果：', flag)
 
 

video_test_rtsp.py

@@ -234,7 +234,7 @@ def main():
     predictor = Predictor(model_name='shufflenet',
                           weights_path=r'./shufflenet.pth',
                           num_classes=2)
-    input_path = r'D:\code\water_turbidity_det\frame_data\1video_20251229124533_hunzhuo'
+    input_path = r'D:\code\water_turbidity_det\frame_data\train\20251225\4_video_20251223163145'
     # 预处理图像
     all_imgs = os.listdir(input_path)
     all_imgs = [os.path.join(input_path, p) for p in all_imgs if p.split('.')[-1] in ['jpg', 'png']]