Image Tagging task
Sample inference script for torchscript exported image-tagger.
The following script should be run from the model export directory:
Hello, thank you for using the code provided by CloudFactory. Please note that some code blocks might not be 100% complete and ready to be run as is. This is done intentionally as we focus on implementing only the most challenging parts that might be tough to pick up from scratch. View our code block as a LEGO block - you can’t use it as a standalone solution, but you can take it and add it to your system to complement it.
python
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['tag_string'] 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)
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_tags = []
for idx in idxs:
present_tags.append(
class_mapping[idx]
)
print("Tags for input:", present_tags)