pytorch 图像分类预测本地图片

def prediect(img_path):    

  img_src = cv2.imread(img_path, 0)  #以灰度图的方式读取要预测的图片

  img=img_src.copy()

  img = cv2.resize(img, (28, 28))      

  height,width=img.shape

  dst=np.zeros((height,width),np.uint8)

  for i in range(height):

      for j in range(width):

          dst[i,j]=255-img[i,j]

  img = dst

  img=np.array(img).astype(np.float32)

  img=np.expand_dims(img,0)

  img=np.expand_dims(img,0)#扩展后，为[1，1，28，28]

  img=torch.from_numpy(img)

  # img = img.to(device)          

  output=model(Variable(img))

  prob = F.softmax(output, dim=1)

  prob = Variable(prob)

  prob = prob.cpu().numpy()  #用GPU的数据训练的模型保存的参数都是gpu形式的，要显示则先要转回cpu，再转回numpy模式

  print(prob)  #prob是10个分类的概率

  pred = np.argmax(prob) #选出概率最大的一个

  print(pred.item())