forked from salwaalkhatib/pytorch-intro
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathinference.py
More file actions
23 lines (18 loc) · 797 Bytes
/
inference.py
File metadata and controls
23 lines (18 loc) · 797 Bytes
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from utils import load_checkpoint, validate, MODELS
from dataset.mydataset import MyDataset
import torch
from torchvision import transforms
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
val_transform = transforms.Compose([
transforms.Resize((224, 224)), # Resize images to 224x224
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
test_dataset = MyDataset(data_dir='chest_xray/test', transform=val_transform)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)
# load best checkpoint
model = MODELS["vgg16"].to(DEVICE)
#TODO: LOAD THE BEST WEIGHTS INTO THE MODEL WITH load_checkpoint
# inference
acc = validate(model, test_loader, DEVICE)
print("[INFO] Accuracy: {:.4f}".format(acc))