import json import time from PIL import Image import torch from torchvision.transforms import transforms model = torch.load('model.pth').to("cuda") model.eval() transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[ 0.48145466, 0.4578275, 0.40821073 ], std=[ 0.26862954, 0.26130258, 0.27577711 ]) # Normalize image ]) with open("tags.json", "r") as file: tags = json.load(file) allowed_tags = sorted(tags) allowed_tags.extend(["placeholder0", "placeholder1", "placeholder2"]) tag_count = len(allowed_tags) image_path="path/to/your/image.png" start = time.time() img = Image.open(image_path).convert('RGB') tensor = transform(img).unsqueeze(0).to("cuda") # transform and add batch dimension with torch.no_grad(): out = model(tensor) probabilities = torch.nn.functional.sigmoid(out[0]) top10_prob, top10_catid = torch.topk(probabilities, 100) for i in range(top10_prob.size(0)): print(allowed_tags[top10_catid[i]], top10_prob[i].item()) end = time.time() print(f'Executed in {end - start} seconds') print("\n\n", end="")