YOLO

How to get results of YOLOv5

I know that we can directly use “results.show()” to get the image with objects being figured out. But what if I want to just show some objects that have bigger confidence than a threshold? Then we need to fetch the results one by one manually:

import torch
import cv2

model = torch.hub.load('.', 'custom', path='best.pt', source='local')
model.eval()

colors = {
    14: (0,255,0),
    80: (0,0,255)
}

names = {
    14: "bird",
    80: "squirrel"
}

for index in ["1.jpeg", "2.jpeg", "4.jpeg", "7.jpeg", "3.webp", "5.webp", "6.webp", "8.png"]:
    img_name = f"squirrel_bird{index}"
    image = cv2.imread(img_name)
    results = model(img, size=960)
    for obj in results.pred[0]:
        x1, y1, x2, y2, conf, cat = obj.numpy()
        x1, y1, x2, y2, cat = int(x1), int(y1), int(x2), int(y2), int(cat)
        print(x1, y1, x2, y2, conf, cat)
        if conf > 0.581 and cat in colors.keys():
            cv2.rectangle(image, (x1, y1), (x2, y2), colors[cat], 2)
            cv2.putText(image, f"{names[cat]},{conf:.2f}", (x1, y1+12), cv2.FONT_HERSHEY_SIMPLEX, 0.5, colors[cat], 1, 2)
    cv2.imshow("yolov5", image)
    cv2.waitKey(0)

The key is to get tensors from “results.pred[0]”, and get coordinates/confidence/category from every tensor.

A trick for using YOLOv5

To detect birds and squirrels, we created a dataset to train the YOLOv5 model. After a week’s training with:

python3 -u train.py --data coco.yaml --cfg yolov5s.yaml --weights '' --batch-size 28 --workers 1

The model could recognize birds and squirrels properly except only for this image:

Why does the model recognize the right-side significant squirrel as a bird? Even though I tried a bigger model, the result was the same…

Only after researching the parameters of the function model() of YOLOv5, I found out we can use a different image size: 960 for detecting.

import inspect
import torch
import cv2

model = torch.hub.load('.', 'custom', path='last.pt', source='local')
#model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
model.eval()

image = cv2.imread(img)
results = model(img, size=960)
results.show()

The result is below for model(img, size=960)

Hmm, seems the single-stage YOLOv5 model is nearsighted, just like me…

Export YOLOv5 models for mobile device

Somebody has finished the work about exporting YOLOv5 models to tflite model. To use it, we only need to:

git clone --single-branch --branch tf-export https://github.com/zldrobit/yolov5.git
cd yolov5
# it will download all pytorch models
sh -x weights/download_weights.sh
# export a tflite model from yolov5l
PYTHONPATH=. python3 models/tf.py --weights yolov5l.pt --cfg models/yolov5l.yaml --img 640
# there will be a tflite model file
ls yolov5l-fp16.tflite

The model file yolov5l-fp16.tflite is 91MB, which is a little too big but still could be put into a mobile phone.

The awesome YOLOv5

I just found a repository YOLOv5 from Github. It’s not just the models are accurate and fast but also easy to get and to use.

Just download the code, install some dependent libraries. And you can just run a simple command:

python3 detect.py --weights yolov5l.pt

Then it will automatically download the YOLO v5 Large model, process all images in inference/images, and put the annotated images into inference/output.

Let’s see some images predicted by YOLO v5 Large:

And every image cost no more than 2 seconds to predict on my laptop. Isn’t it awesome? At least it’s much better than my experiments…

Robin on Linux

YOLO