File size: 3,326 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 | import os
import argparse
import torch
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
from TrainModel import SimpleAE # Ensure TrainModel.py is in PYTHONPATH
def load_checkpoint(path, device='cpu'):
"""
Load autoencoder model and threshold from checkpoint file.
"""
# Load full checkpoint
checkpoint = torch.load(path, map_location=device, weights_only=False)
model = SimpleAE().to(device)
model.load_state_dict(checkpoint['model_state'])
threshold = checkpoint['threshold']
model.eval()
return model, threshold
def preprocess_image(image_path, img_size):
"""
Load and preprocess image: resize, grayscale, to tensor.
"""
img = Image.open(image_path).convert('L')
transform = transforms.Compose([
transforms.Resize((img_size, img_size)),
transforms.ToTensor()
])
tensor = transform(img).unsqueeze(0) # shape: 1x1xHxW
return tensor
def compute_reconstruction_error(model, x_tensor, device='cpu'):
"""
Compute mean squared reconstruction error for input tensor.
"""
x = x_tensor.to(device)
with torch.no_grad():
recon = model(x)
mse = torch.mean((recon - x) ** 2).item()
return mse
def map_error_to_quality(mse, threshold):
"""
Convert reconstruction error to a quality percentage in [0,100].
- If mse <= threshold: quality = 100 * (1 - mse/threshold)
- Else: quality = 0
Confidence equals quality.
"""
if mse <= threshold:
quality = (1 - mse / threshold) * 100
else:
quality = 0.0
# Clamp to [0,100]
quality = max(0.0, min(100.0, quality))
confidence = quality
return round(quality, 2), round(confidence, 2)
def main(args):
# Validate inputs
if not os.path.isfile(args.model_path):
raise FileNotFoundError(f"Model checkpoint not found: {args.model_path}")
if not os.path.isfile(args.image_path):
raise FileNotFoundError(f"Image file not found: {args.image_path}")
# Load model and threshold
model, threshold = load_checkpoint(args.model_path, device='cpu')
# Preprocess image
x = preprocess_image(args.image_path, args.img_size)
# Compute error
mse_error = compute_reconstruction_error(model, x, device='cpu')
# Map to quality and confidence
quality_pct, confidence_pct = map_error_to_quality(mse_error, threshold)
# Output detailed results
print(f"Image: {os.path.basename(args.image_path)}")
print(f"Reconstruction MSE Error: {mse_error:.6f}")
print(f"Threshold (95th percentile): {threshold:.6f}")
print(f"Quality Score: {quality_pct:.2f}%")
print(f"Confidence: {confidence_pct:.2f}%")
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Assess quality of a radiology image using trained QA model')
parser.add_argument('--model_path', type=str, default='./models/autoencoder_qc.pth', help='Path to QA model checkpoint')
parser.add_argument('--image_path', type=str, required=True, help='Path to radiology image (PNG/JPG)')
parser.add_argument('--img_size', type=int, default=128, help='Image resize dimension (square)')
args = parser.parse_args()
main(args)
|