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)

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)

Related Posts

• The awesome YOLOv5

  I just found a repository YOLOv5 from Github. It's not just the models are accurate…

• A trick for using YOLOv5

  To detect birds and squirrels, we created a dataset to train the YOLOv5 model. After…

• Export YOLOv5 models for mobile device

  Somebody has finished the work about exporting YOLOv5 models to tflite model. To use it,…

Leave a Reply Cancel reply

Your email address will not be published. Required fields are marked *

Comment *

Name *

Email *

Website

Save my name, email, and website in this browser for the next time I comment.