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