import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as T

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transforms = T.Compose([T.ToPILImage(), T.Resize(size=224), T.ToTensor()])

def load_saved_model(path: str, device:torch.device = DEVICE) -> nn.Module:
    """
    Load the PyTorch model
    """
    loaded_trace = torch.jit.load(path, map_location=device)
    model = nn.Sequential(loaded_trace, nn.Softmax(dim=1))
    return model

def process_image(x: np.ndarray, device:torch.device = DEVICE) -> torch.Tensor:
    """
    Takes a numpy array provided by gradio as input and returns a ready to be used tensor with the same processing use for training.
    """
    x = transforms(x)
    x = x.unsqueeze(0)
    x = x.to(device)
    return x