Classification

Sample inference script for exported torchscript image classifier

The following sample code should be run from the export directory:

import torch
import numpy as np
from PIL import Image
import json

with open('class_mapping.json') as data:
    mappings = json.load(data)

class_mapping = {item['model_idx']: item['class_name'] for item in mappings}

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = torch.jit.load('model.pt').to(device)

image_path = '/path/to/your/image'
image = Image.open(image_path)
\# Transform your image if the config.yaml shows
\# you used any image transforms for validation data
image = np.array(image)
\# Convert to torch tensor
x = torch.from_numpy(image).to(device)
with torch.no_grad():
    \# Convert to channels first, add batch dimension, convert to float datatype
    x = x.permute(2, 0, 1).unsqueeze(dim=0).float()
    y = model(x)
    \# if model is classifier:
    y = y.argmax(dim=1).squeeze().item()
    predicted_class = class_mapping[y]
print("Predicted class for image:", predicted_class)

Get AI confident. Start using Hasty today.

Our platform is completely free to try. Sign up today to start your two-month trial.