Which files that I need to change if I want to test with only an image?

[Cabbagehust2507]

Hi @hysts , I've trained and evaluated successfully with my own dataset. And now I want to test with only an image, but I'm not sure about any files I need to change. Can you help me this problem? I've tried to change some files and some functions but it's not work.
Thank you so much!

[hysts]

Hi, @Cabbagehust2507

There is no code for inference per se in this repository, but you can load your trained model as follows.

import pytorch_image_classification

config = pytorch_image_classification.get_default_config()
config.merge_from_file('/path/to/your/config.yaml')

device = torch.device(config.device)

model = pytorch_image_classification.create_model(config)
checkpoint = torch.load('/path/to/your/checkpoint.pth')

model.load_state_dict(checkpoint['model'])
model.to(device)
model.eval()

Then, you can use the snippet below to show top5 predictions.

transform = pytorch_image_classification.transforms.create_transform(
    config, is_train=False)

image = cv2.imread('/path/to/your/image.jpg')
data = transform(PIL.Image.fromarray(image))
with torch.no_grad():
    pred = model(data.unsqueeze(0).to(device))
prob = F.softmax(pred, dim=1).cpu()

scores, indices = prob.topk(k=5)
scores = scores.numpy().ravel()
indices = indices.numpy().ravel()
names = [dic[index] for index in indices] # Here, `dic` is a dictionary that maps indices to label names. You need to define it before this line.
for name, score in zip(names, scores):
    print(f'{name}: {score:.3f}')

plt.imshow(image[:, :, ::-1])
plt.show()

[Cabbagehust2507]

It's quite clear to understand, thank you for all your assistance <3