Test page

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)

/for

Want to print your doc?
This is not the way.

Try clicking the ⋯ next to your doc name or using a keyboard shortcut (

CtrlP

) instead.