Spaces:
Build error
Build error
| 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 |