Spaces:
Sleeping
Sleeping
File size: 2,499 Bytes
e556bb4 1d971c4 e556bb4 1d971c4 e556bb4 1d971c4 e556bb4 1d971c4 e556bb4 1d971c4 e556bb4 1d971c4 e556bb4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | import argparse
import os
import torch
import torch.nn as nn
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader
from sklearn.metrics import classification_report, confusion_matrix
def build_model(arch):
if arch == "mobilenet":
model = models.mobilenet_v2()
num_ftrs = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_ftrs, 2)
elif arch in ("swin_t", "swin_t_finetune"):
model = models.swin_t()
num_ftrs = model.head.in_features
model.head = nn.Linear(num_ftrs, 2)
elif arch == "swin_s":
model = models.swin_s()
num_ftrs = model.head.in_features
model.head = nn.Linear(num_ftrs, 2)
else:
raise ValueError(f"Unknown arch: {arch}")
return model
def evaluate(data_dir, model_path, arch):
data_transforms = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
test_dataset = datasets.ImageFolder(os.path.join(data_dir, 'test'), data_transforms)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
model = build_model(arch)
model.load_state_dict(torch.load(model_path, map_location=device))
model = model.to(device)
model.eval()
all_preds = []
all_labels = []
with torch.no_grad():
for inputs, labels in test_loader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
all_preds.extend(preds.cpu().numpy())
all_labels.extend(labels.cpu().numpy())
print("\nClassification Report:")
print(classification_report(all_labels, all_preds, target_names=test_dataset.classes))
print("\nConfusion Matrix:")
print(confusion_matrix(all_labels, all_preds))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--arch", choices=["mobilenet", "swin_t", "swin_t_finetune", "swin_s"], required=True)
parser.add_argument("--data-dir", default="./data/split")
parser.add_argument("--weights", required=True, help="Path to model weights file")
args = parser.parse_args()
evaluate(args.data_dir, args.weights, args.arch)
|