# utils.py
import torch
import json
from torchvision import transforms

with open('label_mapping.json', 'r') as json_file:
    label_mapping = json.load(json_file)

def load_model(path):
    model = torch.jit.load(path, map_location=torch.device("cpu"))
    return model

def predict(model, image):
    model.eval()
    
    # Transform the image
    transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])

    image = transform(image)

    with torch.no_grad():
        image = image.unsqueeze(0)
        output = model(image)
        probabilities = torch.nn.functional.softmax(output, dim=1)
        _, predicted_class = torch.max(probabilities, 1)

        # Convert predicted class index to label name using label_mapping
        predicted_label = label_mapping[f"{predicted_class.item()}"]
        probability= probabilities[0][predicted_class].item()

    return predicted_label, round(probability, 2)