PyTorch——模型推断——单张推断OpenCV（一）

#coding= utf-8
import os
import torch
from data_pipe import get_data
from model import SimpleNet
import numpy as np
import cv2
from PIL import Image


class Infer(object):

    def __init__(self):
        self.model = SimpleNet()
        self.model.load_state_dict(torch.load("./models/model_10.pth"))
        self.model.eval()

    def _infer(self, img_tensor):
        with torch.no_grad():
            result = self.model(img_tensor)
        if result > 0.5:
            result = 1
        else:
            result = 0
        return result

    def predict(self, path):
        img_path_list = [os.path.join(path ,x) for x in os.listdir(path)]
        for img_path in img_path_list:
            print(img_path)
            img = cv2.imread(img_path)
            img = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)
            img_tensor = torch.from_numpy(np.asarray(img)).permute(2,0,1).float()/255.0
            img_tensor = img_tensor.reshape((1, 3, 32, 32))
            result = self._infer(img_tensor)
            print(result)


if __name__ == "__main__":
    path = "./test_images"
    Infer().predict(path)