How to use .pt file locally for inferencing an image

Hi,
I am a beginner in computer vision. I followed excatly the steps illustrated on YOLOV5 github and I used colab and finally downloaded the weights in ‘best.pt’ file. Now, I want to use this file locally for inferencing images. I saw different pieces of code and I tried a few but I got noting. How can I do that?

Assuming you’ve stored the model.state_dict() in the best.pt file, you would have to create the model instance in a new script, load the state_dict, and could perform the inference.
Something like this should work:

model = YOLO(...) # create new model object
model.load_state_dict(torch.load(path_to_state_dict))
model.eval()

# load single image or dataset
img = PIL.Image.open(path_to_image)
...
x = transform(img) # transform to tensor

with torch.inference_mode():
    output = model(x)
1 Like

Alright. I understand the code, but I have one issue. I realized that in the first line ‘model = YOLO()’ I must create an instance of a class that holds the structures and parameters of yolov5. I don’t understand how to do this. Do I need to download yolov5 from github or what? Please explain in more details.

Yes, you need the class definition from the repository to create the model object in the same way as you’ve done it during training.

1 Like

Okay. I downloaded yolov5 from github into my project folder. What is next? Excuse me but it isn’t clear for me.

You will have to import the yolo model (e.g. import models.yolo.Model as YOLO) such that you can instantiate an untrained model as per model = YOLO(...). From there you can work with the code provided.

1 Like