jiyuhang 2 mesi fa
parent
commit
fe910ba811
4 ha cambiato i file con 23 aggiunte e 14 eliminazioni
  1. 1 2
      bmodel_application.py
  2. 0 0
      model/__init__.py
  3. 8 0
      model/model_zoon.py
  4. 14 12
      video_test.py

+ 1 - 2
bmodel_application.py

@@ -183,7 +183,7 @@ class Predictor:
         self.continuous_counter_mat[self.continuous_counter_mat<0] = 0
         # 判断浑浊
         bad_flag = bool(np.sum(self.continuous_counter_mat > self.max_continuous_frames) > 2)  # 两个以上的patch满足条件
-        print('连续帧信息:', self.continuous_counter_mat)
+        # print('连续帧信息:', self.continuous_counter_mat)
         print(f'连续帧判别:该时间段是否存在浑浊水体:{bad_flag}')
         return bad_flag
 
@@ -219,7 +219,6 @@ def argsparser():
     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()
-
     return args
 
 def main(args):

+ 0 - 0
model/__init__.py


+ 8 - 0
model/model_zoon.py

@@ -0,0 +1,8 @@
+
+
+
+
+
+def load_model(name:str):
+
+    pass

+ 14 - 12
video_test.py

@@ -206,24 +206,25 @@ def discriminate_ratio(water_pre_list:list):
     # 方式一: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 >= 0.6 * 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_cont(pre_class_arr, continuous_count_mat):
+def discriminate_count(pre_class_arr, continuous_count_mat):
     """连续帧判别"""
     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] = 0
+    continuous_count_mat[negative_index] -= 1
     # 给正样本区域加1
     continuous_count_mat[positive_index] += 1
+    # 保证不出现负数
+    continuous_count_mat[continuous_count_mat<0] = 0
     # 判断浑浊
-    bad_flag = np.max(continuous_count_mat) > 30
-    if bad_flag:
-        print(f'连续帧方式:该时间段是否存在浑浊水体:{bool(bad_flag)}')
+    bad_flag = bool(np.sum(continuous_count_mat > 15) > 2)
+    print(f'连续帧方式:该时间段是否存在浑浊水体:{bad_flag}')
     return bad_flag
 
 def main():
@@ -233,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\train\20251230\4video_20251229160514'
+    input_path = r'D:\code\water_turbidity_det\frame_data\1video_20251229124533_hunzhuo'
     # 预处理图像
     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']]
@@ -280,14 +281,15 @@ def main():
         cv2.imshow('image_filter', new_img)
 
         cv2.waitKey(25)
+        water_pre_list.append(new_predicted_class)
+        # 方式2判别
+        flag = discriminate_count(new_predicted_class, continuous_count_mat)
         # 方式1判别
         if len(water_pre_list) > 25:
             flag = discriminate_ratio(water_pre_list) and flag
-            water_pre_list = []
-            print('综合判别结果:', flag)
-        water_pre_list.append(new_predicted_class)
-        # 方式2判别
-        flag = discriminate_cont(new_predicted_class, continuous_count_mat)
+        print('综合判别结果:', flag)
+
+
 
 if __name__ == "__main__":
     main()