| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from model import CnnNet | |
| # Index order used in training: 0 = background, 1 = helicopter, 2 = drone | |
| CLASS_NAMES = ["background", "helicopter", "drone"] | |
| model = CnnNet() | |
| model.load_state_dict(torch.load("weightsfilefromspectrogram.pth", map_location="cpu")) | |
| model.eval() | |
| def predict(image_path): | |
| # Spectrograms are single-channel, 480x640, pixel values scaled to 0-1 | |
| img = Image.open(image_path).convert("L").resize((640, 480)) | |
| x = torch.tensor(np.array(img) / 255.0, dtype=torch.float32) | |
| x = x.unsqueeze(0).unsqueeze(0) # shape [1, 1, 480, 640] | |
| with torch.no_grad(): | |
| probs = torch.exp(model(x))[0] # the model outputs log-probabilities | |
| return {name: float(p) for name, p in zip(CLASS_NAMES, probs)} |