File size: 3,687 Bytes
654303c |
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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image
import pandas as pd
import io
import ast
import os
from sklearn.metrics import confusion_matrix, classification_report
import timm
class EfficientNetB0Alpha(nn.Module):
def __init__(self, num_classes=26):
super().__init__()
self.model = timm.create_model('efficientnet_b0', pretrained=True, in_chans=1, num_classes=num_classes)
def forward(self, x):
return self.model(x)
class Dataset(Dataset):
def __init__(self, csv_path, transform=None, image_col='image', label_col='label'):
self.data = pd.read_csv(csv_path)
self.transform = transform
self.image_col = image_col
self.label_col = label_col
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
img_data = self.data.iloc[idx][self.image_col]
label = self.data.iloc[idx][self.label_col]
if isinstance(img_data, str):
img_dict = ast.literal_eval(img_data)
img_bytes = img_dict['bytes']
else:
img_bytes = img_data['bytes']
img = Image.open(io.BytesIO(img_bytes)).convert('L')
if self.transform:
img = self.transform(img)
return img, label
def load_model(model_path, num_classes, device):
if not os.path.exists(model_path):
raise FileNotFoundError(f"Model file not found at {model_path}")
model = EfficientNetB0Alpha(num_classes=num_classes)
checkpoint = torch.load(model_path, map_location=device, weights_only=True)
model.load_state_dict(checkpoint['model_state_dict'])
model.to(device)
model.eval()
return model
def evaluate(model, test_loader, device, class_names):
model.eval()
correct, total = 0, 0
all_preds, all_labels = [], []
with torch.no_grad():
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
all_preds.extend(predicted.cpu().numpy())
all_labels.extend(labels.cpu().numpy())
accuracy = 100 * correct / max(total, 1)
print(f"Test Accuracy: {accuracy:.2f}%")
print("\nClassification Report:")
print(classification_report(all_labels, all_preds, target_names=class_names, digits=2))
cm = confusion_matrix(all_labels, all_preds)
print("\nConfusion Matrix (True Labels: rows, Predicted Labels: columns):")
print(pd.DataFrame(cm, index=class_names, columns=class_names))
return accuracy, cm
def main():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_classes = 26
model_path = "saved_models/best_model.pth"
test_csv = "dataset/test.csv"
batch_size = 32
print("Device being used:", device)
test_transform = transforms.Compose([
transforms.Grayscale(num_output_channels=1),
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5], std=[0.5])
])
test_dataset = Dataset(test_csv, transform=test_transform)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)
class_names = [chr(65 + i) for i in range(26)]
model = load_model(model_path, num_classes, device)
evaluate(model, test_loader, device, class_names)
if __name__ == "__main__":
main() |