OpenCV pytorch 预测错误与 onnx model

Question

My codes are here. I made a model for the prediction of cell and converted it to onnx then load with OpenCV to predicted with OpenCV but went somethings wrong

Libraries

import cv2
import torchvision.models as models
import torch.onnx
import torchvision.transforms as transforms
import numpy as np

Onnx and model

original_model = models.resnet50(pretrained=True)

opencv_net = cv2.dnn.readNetFromONNX('resnet50.onnx')

Opencv prediction

opencv_net = cv2.dnn.readNetFromONNX('resnet50.onnx')
input_img=cv2.imread('image.bmp',cv2.COLOR_BGR2GRAY)
input_img=input_img.astype(np.float32)
input_img=cv2.resize(input_img,(256,256))

mean=np.array([0.485, 0.456, 0.406]) * 255.0
scale=1/255.0
std=[0.229, 0.224, 0.225]
input_blob = cv2.dnn.blobFromImage(
    image=input_img,
    scalefactor=scale,
    size=(224, 224),  # img target size
    mean=mean,
    #swapRB=True,  # BGR -> RGB
    crop=True  # center crop
)

input_blob[0] /= np.asarray(std, dtype=np.float32).reshape(3, 1, 1)
print("Input blob shape: {}\n".format(input_blob.shape))
preproc_img=input_blob

opencv_net.setInput(preproc_img)
out = opencv_net.forward()
print("OpenCV DNN prediction: \n")
print("* shape: ", out.shape)

test_class_id = np.argmax(out)

test_labels=opencv_net.getLayerNames()
#print((test_labels))

confidence = out[0][test_class_id]
print("* class ID: {}, label: {}".format(test_class_id, test_labels[test_class_id]))
print("* confidence: {:.4f}".format(confidence))

Opencv predictions output

OpenCV DNN prediction: 

* shape:  (1, 2)
* class ID: 1, label: 323
* confidence: 8.4153

**!!!Problem is here!!! Pytorch inference**

original_model.eval()
preproc_img = torch.FloatTensor(preproc_img)
# inference
out = original_model(preproc_img)

print("\nPyTorch model prediction: \n")
print("* shape: ", out.shape)

test_class_id = torch.argmax(out, axis=1).item()
print("* class ID: {}, label: {}".format(test_class_id, test_labels[test_class_id]))

confidence = out[0][test_class_id]
print("* confidence: {:.4f}".format(confidence.item()))

error?

* shape:  torch.Size([1, 1000])
Traceback (most recent call last):
  File "X.py", line 121, in <module>
    print("* class ID: {}, label: {}".format(test_class_id, test_labels[test_class_id]))
IndexError: tuple index out of range

Process finished with exit code 1

How is this possible while i'm using the same img and same model?

Answer 1

I solved this problem. The problem was the model. I was using here resnet 50 pretrained but I need my model so I used these lines to solve and its worked.

model = models.resnet50(pretrained = True)
model.fc = nn.Linear(in_features=2048, out_features=2, bias=True)
weights = torch.load('model_best.pth',map_location ='cpu')
model.load_state_dict(weights)
model.eval()

gives me same opencv outputs

PyTorch model prediction: 

* shape:  torch.Size([1, 2])
* class ID: 1, label: 323
* confidence: 8.4153