|
|
import torch.utils.data.distributed |
|
|
import torchvision.transforms as transforms |
|
|
import torchvision.datasets as datasets |
|
|
from torch.autograd import Variable |
|
|
import pandas as pd |
|
|
from PIL import Image |
|
|
|
|
|
classes = ('cat', 'dog') |
|
|
transform_test = transforms.Compose([ |
|
|
transforms.Resize((224, 224)), |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) |
|
|
]) |
|
|
|
|
|
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") |
|
|
model = torch.load("model_epoch20_lr0.0001_best_epoch2.pth") |
|
|
model.eval() |
|
|
model.to(DEVICE) |
|
|
|
|
|
|
|
|
predictions = [] |
|
|
|
|
|
for index in range(1, 501): |
|
|
img_path = f'datasets/datatest/test/{index}.jpg' |
|
|
|
|
|
|
|
|
img = Image.open(img_path) |
|
|
img = transform_test(img) |
|
|
img = img.unsqueeze(0).to(DEVICE) |
|
|
output = model(img) |
|
|
_, pred = torch.max(output.data, 1) |
|
|
predictions.append(pred.data.item()) |
|
|
|
|
|
|
|
|
|
|
|
df = pd.DataFrame(predictions, columns=['prediction']) |
|
|
|
|
|
|
|
|
|
|
|
df.to_excel('predictions.xlsx', index=False) |
|
|
|
|
|
|
|
|
|
|
|
|