Predict file to run yourself

predict.py

@@ -0,0 +1,20 @@
+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)}