1

.env

@@ -13,4 +13,6 @@ CUDA_VISIBLE_DEVICES=0
 # batch size
 BATCH_SIZE=128
 # pretrained
-PRETRAINED=True
+PRETRAINED=True
+# 是否使用bias
+USE_BIAS=True

.gitignore

@@ -52,4 +52,5 @@ __pycache__/
 *.json
 *.profile
 # 日志文件
-runs/*
+*/runs/*
+*_f32*/

accuracy.py

@@ -4,7 +4,7 @@ import numpy as np
 
 def argsparser():
     parser = argparse.ArgumentParser(prog=__file__)
-    parser.add_argument('--input', type=str, default='./prediction_result.json', help='path of prediction')
+    parser.add_argument('--input', '-i',type=str, default='./prediction_result.json', help='path of prediction json')
     args = parser.parse_args()
     return args
 
@@ -75,18 +75,18 @@ def main(args):
     # 分离真值和预测值
     y_true = data[:, 0]  # 第一列是真值
     y_pred = data[:, 1]  # 第二列是预测值
-    # 计算评估指标
-    metrics = calculate_metrics(y_true, y_pred)
-
-    # 打印结果
-    print("=== 分类评估结果 ===")
-    print(f"总体精度: {metrics['accuracy']:.4f}")
-    print(f"混淆矩阵：TP FN")
-    print(f"        FP TN")
-    print(f"{metrics['confusion_matrix']}")
-    print(f"查准率: {metrics['precision']:.4f}")
-    print(f"查全率: {metrics['recall']:.4f}")
-    print(f"F1分数: {metrics['f1']:.4f}")
+    # # 计算评估指标
+    # metrics = calculate_metrics(y_true, y_pred)
+    #
+    # # 打印结果
+    # print("=== 分类评估结果 ===")
+    # print(f"总体精度: {metrics['accuracy']:.4f}")
+    # print(f"混淆矩阵：TP FN")
+    # print(f"        FP TN")
+    # print(f"{metrics['confusion_matrix']}")
+    # print(f"查准率: {metrics['precision']:.4f}")
+    # print(f"查全率: {metrics['recall']:.4f}")
+    # print(f"F1分数: {metrics['f1']:.4f}")
     from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
     # 计算各项指标
     report = classification_report(y_true, y_pred)

bmodel_test.py

@@ -80,9 +80,9 @@ class Predictor:
 
 def argsparser():
     parser = argparse.ArgumentParser(prog=__file__)
-    parser.add_argument('--input', type=str, default='./test', help='path of input, must be image directory')
-    parser.add_argument('--bmodel', type=str, default='./shufflenet_f32.bmodel', help='path of bmodel')
-    parser.add_argument('--dev_id', type=int, default=0, help='tpu id')
+    parser.add_argument('--input','-i', type=str, default='./test', 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()
 
     return args
@@ -124,7 +124,6 @@ def main(args):
 
 
 
-
 if __name__ == '__main__':
     args = argsparser()
     main(args)

pth2onnx.py

@@ -5,16 +5,16 @@ import onnxruntime as ort
 import numpy as np
 from torchvision.models import resnet50, shufflenet_v2_x1_0, shufflenet_v2_x2_0
 from torch import nn
-from simple_model import  SimpleModel
+# from simple_model import  SimpleModel
 if __name__ == '__main__':
 
     # 载入模型框架
     # model = SimpleModel()
-    # model = resnet50(pretrained=True)
-
-    model = shufflenet_v2_x1_0()
-    model.fc = nn.Linear(int(model.fc.in_features), 2, bias=False)
-    model.load_state_dict(torch.load(r'./shufflenet.pth')) # xxx.pth表示.pth文件, 这一步载入模型权重
+    model = resnet50(pretrained=False)
+    # model = shufflenet_v2_x1_0()
+    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("加载模型成功")
     model.eval() # 设置模型为推理模式
 
@@ -22,7 +22,7 @@ if __name__ == '__main__':
    # print(model)
     torch.onnx.export(model,
                       example_input,
-                      "shufflenet.onnx",
+                      f"{model_name}.onnx",
                       opset_version=13,
                       export_params=True,
                       do_constant_folding=True,
@@ -30,10 +30,10 @@ if __name__ == '__main__':
 
     # 验证模型
 
-    onnx_model = onnx.load("shufflenet.onnx")  # 使用不同变量名
+    onnx_model = onnx.load(f"{model_name}.onnx")  # 使用不同变量名
     onnx.checker.check_model(onnx_model)  # 验证模型完整性
     # 使用ONNX Runtime进行推理
-    ort_session = ort.InferenceSession("shufflenet.onnx")
+    ort_session = ort.InferenceSession(f"{model_name}.onnx")
     ort_inputs = {ort_session.get_inputs()[0].name: example_input.detach().numpy()}
     ort_outs = ort_session.run(None, ort_inputs)
 

test.py → video_test.py

@@ -26,6 +26,7 @@ class Predictor:
         self.weights_path = weights_path
         self.num_classes = num_classes
         self.model = None
+        self.use_bias = os.getenv('USE_BIAS', True)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         print(f"当前设备: {self.device}")
         # 加载模型
@@ -51,13 +52,13 @@ class Predictor:
         # 替换最后的分类层以适应新的分类任务
         if hasattr(self.model, 'fc'):
             # ResNet系列模型
-            self.model.fc = nn.Linear(int(self.model.fc.in_features), self.num_classes, bias=False)
+            self.model.fc = nn.Linear(int(self.model.fc.in_features), self.num_classes, bias=self.use_bias)
         elif hasattr(self.model, 'classifier'):
             # Swin Transformer等模型
-            self.model.classifier = nn.Linear(int(self.model.classifier.in_features), self.num_classes, bias=False)
+            self.model.classifier = nn.Linear(int(self.model.classifier.in_features), self.num_classes, bias=self.use_bias)
         elif hasattr(self.model, 'head'):
             # Swin Transformer使用head层
-            self.model.head = nn.Linear(int(self.model.head.in_features), self.num_classes, bias=False)
+            self.model.head = nn.Linear(int(self.model.head.in_features), self.num_classes, bias=self.use_bias)
 
         else:
             raise ValueError(f"Model {name} does not have recognizable classifier layer")
@@ -220,9 +221,9 @@ def main():
     # 初始化模型实例
     # TODO:修改模型网络名称/模型权重路径/视频路径
     predictor = Predictor(model_name='shufflenet',
-                          weights_path=r'/shufflenet.pth',
+                          weights_path=r'./shufflenet.pth',
                           num_classes=2)
-    input_path = r'frame_data/train/20251225/4_video_202511211127'
+    input_path = r'D:\code\water_turbidity_det\frame_data\test\20251225\video4_20251129120320_20251129123514'
     # 预处理图像
     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']]
@@ -268,7 +269,7 @@ def main():
 
         cv2.waitKey(20)
         # 方式1判别
-        if len(water_pre_list) > 100:
+        if len(water_pre_list) > 20:
             flag = discriminate_ratio(water_pre_list) and flag
             water_pre_list = []
             print('综合判别结果：', flag)