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)}