Attribute Prediction

Sample inference script for torchscript exported image-tagger.

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['attribute_name'], item['value']) for item in mappings}

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

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

# Path to your image
image_path = '/path/to/your/image'
image =
# 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)
    y = torch.sigmoid(y).squeeze()
    # All classes with probabilities > 0.5 are considered present in
    # the input. You can tweak this 0.5 threshold if you desire.
    idxs = torch.where(y > 0.5)[0].cpu().numpy()
    present_attributes_and_values = []
    for idx in idxs:
print("Attribute/Value pairs for input:", present_attributes_and_values)
Last updated on Dec 21, 2022

Removing the risk from vision AI.

Only 13% of vision AI projects make it to production, with Hasty we boost that number to 100%.

Start for free Request a demo