2 min read

PyTorch and YOLOv5 Computer Vision

PyTorch and YOLOv5 Computer Vision
Photo by Chris Ried / Unsplash

Perhaps the best computer vision results I have ever achieved. These two things go perfect together, and I'm happy I was able to finally use PyTorch for once. Anyway, here are the results:

If you want to see the previous results with OpenCV and YOLOv4, check out this blog post.

So here's how I did it. (By the way, this is all under 20 lines of code, which is crazy.)

First I had to import PyTorch and PILLOW for images. Then a variable input to ask which image to use. Then we download the model from PyTorch Hub. Then we open the image stored in the variable. Then we inference it with some cool stuff, I don't know what the hell is happening in this part, but it works so good enough. Print results and show the results as an image or save the image. And then finally there's some other stuff needed which I have no clue what it does. Here's the source code:

import torch
from PIL import Image

# Image chosen by user
image = input("What image do you want to use? ")

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# Images
img1 = Image.open(image)  # PIL image
imgs = [img1]  # batch of images

# Inference
results = model(imgs, size=640)  # includes NMS

# Results
results.print()  
results.save()  # or .show()

results.xyxy[0]  # img1 predictions (tensor)
results.pandas().xyxy[0]  # img1 predictions (pandas)

Good results in this run, and as always, I'm gonna get back to work and see what other cool shit I can make. Maybe some simple application which allows me to zoom in on these images so I can get a better view at results. We'll see though. Good shit happening lately! The power of Freethinking never stops.