File size: 4,405 Bytes
05e9fbe | 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | import os
import time
import argparse
import torch
import numpy as np
from glob import glob
from PIL import Image, ImageFilter
import matplotlib.pyplot as plt
from sklearn.metrics import precision_score, recall_score, accuracy_score, roc_curve, auc
import torchvision.transforms as transforms
from TrainModel import SimpleAE # ensure TrainModel.py in PYTHONPATH
# -------- Dataset for Testing --------
def load_test_set(data_dir, img_size, noise_std):
files = glob(os.path.join(data_dir, '*.png')) + glob(os.path.join(data_dir, '*.jpg'))
X, y = [], []
transform = transforms.Compose([
transforms.Resize((img_size, img_size)),
transforms.Grayscale(num_output_channels=1),
transforms.ToTensor()
])
for p in files:
img = Image.open(p).convert('L')
x = transform(img)
# clean
X.append(x); y.append(1)
# noisy
arr = np.array(img).astype(np.float32)
arr_n = np.clip(arr + np.random.normal(0, noise_std, arr.shape), 0, 255).astype(np.uint8)
img_n = Image.fromarray(arr_n)
X.append(transform(img_n)); y.append(0)
# blurred
img_b = img.filter(ImageFilter.GaussianBlur(radius=2))
X.append(transform(img_b)); y.append(0)
return torch.stack(X), np.array(y)
# -------- Testing Script --------
def main(args):
os.makedirs(args.output_dir, exist_ok=True)
# Load checkpoint with weights_only=False to allow loading full dict
checkpoint = torch.load(args.model_path, map_location='cpu', weights_only=False)
model = SimpleAE()
model.load_state_dict(checkpoint['model_state'])
threshold = checkpoint['threshold']
model.eval()
# Prepare test set
X, y_true = load_test_set(args.data_dir, args.img_size, args.noise_std)
# Inference metrics
errors, times = [], []
with torch.no_grad():
for x in X:
start = time.time()
out = model(x.unsqueeze(0))
times.append(time.time() - start)
err = ((out - x.unsqueeze(0))**2).mean().item()
errors.append(err)
y_pred = np.array([1 if e <= threshold else 0 for e in errors])
# Calculate metrics
acc = accuracy_score(y_true, y_pred)
prec = precision_score(y_true, y_pred)
rec = recall_score(y_true, y_pred)
fpr, tpr, _ = roc_curve(y_true, [-e for e in errors])
roc_auc = auc(fpr, tpr)
avg_time = np.mean(times)
model_size = os.path.getsize(args.model_path) / (1024*1024)
# Save summary
with open(os.path.join(args.output_dir, 'metrics.txt'), 'w') as f:
f.write(f"Accuracy: {acc:.4f}\nPrecision: {prec:.4f}\nRecall: {rec:.4f}\nROC AUC: {roc_auc:.4f}\n")
f.write(f"Avg inference time (s): {avg_time:.6f}\nModel size (MB): {model_size:.2f}\n")
# Plot and save
plt.figure()
plt.hist(errors, bins=50)
plt.title('Reconstruction Error Distribution')
plt.xlabel('MSE Error')
plt.ylabel('Count')
plt.savefig(os.path.join(args.output_dir, 'error_distribution.png'))
plt.close()
plt.figure()
plt.plot(fpr, tpr, label=f'AUC = {roc_auc:.2f}')
plt.plot([0,1],[0,1],'--')
plt.title('ROC Curve')
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.legend(loc='lower right')
plt.savefig(os.path.join(args.output_dir, 'roc_curve.png'))
plt.close()
plt.figure()
plt.plot(sorted(times), marker='.')
plt.title('Inference Time per Sample')
plt.xlabel('Sample index')
plt.ylabel('Time (s)')
plt.savefig(os.path.join(args.output_dir, 'inference_time.png'))
plt.close()
print(f"Test results saved in {args.output_dir}")
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Evaluate QA model')
parser.add_argument('--data_dir', type=str, default='./data', help='Test images directory')
parser.add_argument('--model_path', type=str, default='./models/autoencoder_qc.pth', help='Path to QC model')
parser.add_argument('--output_dir', type=str, default='./TestModel', help='Directory to save results')
parser.add_argument('--img_size', type=int, default=128, help='Image resize')
parser.add_argument('--noise_std', type=float, default=25.0, help='Synthetic noise std')
args = parser.parse_args()
main(args) |