51 lines
1.2 KiB
Python
51 lines
1.2 KiB
Python
import cv2
|
|
import glob
|
|
import torch
|
|
from torchvision import transforms
|
|
from PIL import Image
|
|
import numpy as np
|
|
# from resnet import resnet
|
|
from torchvision.models import resnet18
|
|
from const import label_name, input_size
|
|
|
|
|
|
def test():
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(device)
|
|
|
|
net = resnet18(weights=None)
|
|
net.load_state_dict(torch.load("./model/resnet_epoch_31.pth", weights_only=True))
|
|
|
|
im_list = glob.glob("./dataset/test/*/*.jpg")
|
|
np.random.shuffle(im_list)
|
|
|
|
net.to(device)
|
|
|
|
test_transform = transforms.Compose([
|
|
transforms.Resize((input_size, input_size)),
|
|
transforms.ToTensor()
|
|
])
|
|
|
|
for im_path in im_list:
|
|
net.eval()
|
|
im_data = Image.open(im_path)
|
|
|
|
inputs = test_transform(im_data)
|
|
inputs = torch.unsqueeze(inputs, dim=0)
|
|
|
|
inputs = inputs.to(device)
|
|
outputs = net.forward(inputs)
|
|
print("outputs", outputs)
|
|
|
|
_, 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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test()
|