Run Yolov5 cls on the image and Get results. This script also takes care of transformation

yolov5-cls-image-classification-example.py

import torch
from PIL import Image
import torch.nn.functional as F
from torchvision import transforms as T

# ! python classify/predict.py --weights yolov5s-cls.pt --source https://thumbs.dreamstime.com/z/motorcycle-helmet-22053909.jpg --save-txt 
model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s-cls.pt') 


IMAGENET_MEAN = 0.485, 0.456, 0.406
IMAGENET_STD = 0.229, 0.224, 0.225
def classify_transforms(size=224):
    return T.Compose([T.ToTensor(), T.Resize(size), T.CenterCrop(size), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])


def classify_transforms(size=224):
    return T.Compose([T.ToTensor(), T.Resize(size), T.CenterCrop(size), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])
! wget https://thumbs.dreamstime.com/z/motorcycle-helmet-22053909.jpg
! pwd
imgs = "/content/yolov5/motorcycle-helmet-22053909.jpg"
image = Image.open(imgs)
transformations = classify_transforms()
convert_tensor = transformations(image)
convert_tensor = convert_tensor.unsqueeze(0)
results = model(convert_tensor)
# print(results)
pred = F.softmax(results, dim=1)

for i, prob in enumerate(pred):
    top5i = prob.argsort(0, descending=True)[:5].tolist()
    text = '\n'.join(f'{prob[j]:.2f} {model.names[j]}' for j in top5i)
    print(text)