resnet parm

core/test/test_resnet.py

@@ -1,3 +1,4 @@
+import os
 import cv2
 import glob
 import torch
@@ -5,21 +6,32 @@ from torchvision import transforms
 from PIL import Image
 import numpy as np
 from core.nets.resnet import resnet
-from core.const import label_name, input_size
+from core.const import mode, label_name, input_size
 
 
 def test():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(device)
 
-    net = resnet()
-    net.load_state_dict(torch.load("./model/resnet_epoch_15.pth", weights_only=True))
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    core_dir = os.path.dirname(script_dir)
+    model_dir = os.path.join(core_dir, "models")
+    dataset_dir = os.path.join(core_dir, "dataset", mode, "test")
 
-    im_list = glob.glob("./dataset/test/*/*.jpg")
+    print("model_dir", model_dir)
+
+    net = resnet()
+    net.load_state_dict(torch.load(os.path.join(model_dir, "resnet_epoch_50.pth"), weights_only=True))
+
+    print("111")
+
+    im_list = glob.glob(os.path.join(dataset_dir, "*", "*.jpg"))
     np.random.shuffle(im_list)
 
     net.to(device)
 
+    print("222")
+
     test_transform = transforms.Compose([
         transforms.Resize(input_size),
         transforms.CenterCrop(input_size),
@@ -27,7 +39,11 @@ def test():
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
     ])
 
+    print("333")
+
+    ary = []
     for im_path in im_list:
+        #  print("im_path", im_path)
         net.eval()
         im_data = Image.open(im_path)
 
@@ -40,11 +56,15 @@ def test():
         _, pred = torch.max(outputs.data, dim=1)
         print(label_name[pred.cpu().numpy()[0]], " ", im_path)
 
-        img = np.asarray(im_data)
-        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-        cv2.imshow("img", img)
-        cv2.waitKey(0)
+        result = label_name[pred.cpu().numpy()[0]]
 
+        if result in im_path:
+            ary.append(True)
+        else:
+            ary.append(False)
+
+    print("ary", ary)
+    print(sum(ary) / len(ary))
 
 if __name__ == "__main__":
     test()

core/test/test_resnet_onnx.py

@@ -0,0 +1,57 @@
+import cv2
+import glob
+import numpy as np
+from PIL import Image
+import onnxruntime as ort  # 【修改1】替换torch导入
+from torchvision import transforms  # 保留transforms，仍用于预处理
+from core.const import label_name, input_size
+
+
+def test():
+    # 【修改2】选择ONNX Runtime的执行提供程序，自动选择CPU或CUDA
+    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if 'CUDAExecutionProvider' in ort.get_available_providers() else ['CPUExecutionProvider']
+    print(f"使用设备: {providers[0]}")
+
+    # 【修改3】加载ONNX模型，替代原来的PyTorch模型加载
+    session = ort.InferenceSession("./model/resnet_final.onnx", providers=providers)
+    
+    # 获取输入名称（用于后续推理时指定输入）
+    input_name = session.get_inputs()[0].name
+
+    im_list = glob.glob("./dataset/test/*/*.jpg")
+    np.random.shuffle(im_list)
+
+    # 预处理完全不变
+    test_transform = transforms.Compose([
+        transforms.Resize(input_size),
+        transforms.CenterCrop(input_size),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+
+    for im_path in im_list:
+        im_data = Image.open(im_path)
+
+        inputs = test_transform(im_data)
+        inputs = torch.unsqueeze(inputs, dim=0)  # 这里还在用torch，下面会改
+
+        # 【修改4】将输入转为numpy，ONNX Runtime需要numpy输入
+        inputs = inputs.numpy()
+        if providers[0] == 'CPUExecutionProvider':
+            inputs = inputs.astype(np.float32)
+
+        # 【修改5】ONNX Runtime推理，输出直接是numpy数组
+        outputs = session.run(None, {input_name: inputs})[0]
+
+        # 【修改6】解析结果，直接用numpy操作
+        pred = np.argmax(outputs, axis=1)
+        print(label_name[pred[0]], " ", im_path)
+
+        img = np.asarray(im_data)
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+        cv2.imshow("img", img)
+        cv2.waitKey(0)
+
+
+if __name__ == "__main__":
+    test()