Attribute Prediction
Sample inference script for torchscript exported image-tagger.
The following sample code should be run from the export directory:
sample_attributer.py
1
import torch
2
import numpy as np
3
from PIL import Image
4
import json
5
6
with open('class_mapping.json') as data:
7
mappings = json.load(data)
8
9
class_mapping = {item['model_idx']: (item['attribute_name'], item['value']) for item in mappings}
10
11
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
12
13
model = torch.jit.load('model.pt').to(device)
14
15
# Path to your image
16
image_path = '/path/to/your/image'
17
image = Image.open(image_path)
18
# Transform your image according to the transforms.json as in
19
# https://help.hasty.ai/model-playground/image-transformations
20
image = np.array(image)
21
# Convert to torch tensor
22
x = torch.from_numpy(image).to(device)
23
with torch.no_grad():
24
# Convert to channels first, add batch dimension, convert to float datatype
25
x = x.permute(2, 0, 1).unsqueeze(dim=0).float()
26
y = model(x)
27
y = torch.sigmoid(y).squeeze()
28
# All classes with probabilities > 0.5 are considered present in
29
# the input. You can tweak this 0.5 threshold if you desire.
30
idxs = torch.where(y > 0.5)[0].cpu().numpy()
31
present_attributes_and_values = []
32
for idx in idxs:
33
present_attributes_and_values.append(
34
class_mapping[idx]
35
)
36
print("Attribute/Value pairs for input:", present_attributes_and_values)
Copied!
Copy link