Tải YOLOv5 từ PyTorch Hub

📚 Hướng dẫn này giải thích cách tải YOLOv5 🚀 từ PyTorch Hub tại https://pytorch.org/hub/ultralytics_yolov5.

Trước khi Bạn Bắt đầu

Cài đặt requirements.txt trong môi trường Python>=3.8.0, bao gồm PyTorch>=1.8. Các model và tập dữ liệu sẽ tự động tải xuống từ bản release YOLOv5 mới nhất.

pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt

Mẹo

Việc clone ultralytics/yolov5 là không bắt buộc — PyTorch Hub sẽ tự động lấy mã nguồn.

Tải YOLOv5 với PyTorch Hub

Ví dụ đơn giản

Ví dụ này tải một model YOLOv5s đã được huấn luyện trước từ PyTorch Hub dưới dạng model và truyền vào một hình ảnh để thực hiện inference. 'yolov5s' là model YOLOv5 nhẹ nhất và nhanh nhất. Để biết chi tiết về tất cả các model khả dụng, vui lòng xem README.

import torch

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

# Image
im = "https://ultralytics.com/images/zidane.jpg"

# Inference
results = model(im)

results.pandas().xyxy[0]
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 1  433.50  433.50   517.5  714.5    0.687988     27     tie
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

Ví dụ chi tiết

Ví dụ này minh họa batched inference với các nguồn ảnh PIL và OpenCV. Các results có thể được in ra console, lưu vào runs/hub, hiển thị lên màn hình trên các môi trường hỗ trợ, và trả về dưới dạng tensors hoặc pandas dataframes.

import cv2
import torch
from PIL import Image

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

# Images
for f in "zidane.jpg", "bus.jpg":
    torch.hub.download_url_to_file("https://ultralytics.com/images/" + f, f)  # download 2 images
im1 = Image.open("zidane.jpg")  # PIL image
im2 = cv2.imread("bus.jpg")[..., ::-1]  # OpenCV image (BGR to RGB)

# Inference
results = model([im1, im2], size=640)  # batch of images

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

results.xyxy[0]  # im1 predictions (tensor)
results.pandas().xyxy[0]  # im1 predictions (pandas)
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 1  433.50  433.50   517.5  714.5    0.687988     27     tie
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

For all inference options see YOLOv5 AutoShape() forward method.

Cài đặt Inference

Các model YOLOv5 chứa nhiều thuộc tính inference như confidence threshold, IoU threshold, v.v. có thể được thiết lập bởi:

model.conf = 0.25  # NMS confidence threshold
model.iou = 0.45  # NMS IoU threshold
model.agnostic = False  # NMS class-agnostic
model.multi_label = False  # NMS multiple labels per box
model.classes = None  # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
model.max_det = 1000  # maximum number of detections per image
model.amp = False  # Automatic Mixed Precision (AMP) inference

results = model(im, size=320)  # custom inference size

Thiết bị (Device)

Các model có thể được chuyển sang bất kỳ thiết bị nào sau khi khởi tạo:

model.cpu()  # CPU
model.cuda()  # GPU
model.to(device)  # i.e. device=torch.device(0)

Các model cũng có thể được khởi tạo trực tiếp trên một device bất kỳ:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", device="cpu")  # load on CPU

Mẹo

Các hình ảnh đầu vào sẽ tự động được chuyển sang đúng thiết bị của model trước khi thực hiện inference.

Tắt thông báo (Silence Outputs)

Các model có thể được tải ở chế độ im lặng (không thông báo) với _verbose=False:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", _verbose=False)  # load silently

Kênh đầu vào (Input Channels)

Để tải một model YOLOv5s đã huấn luyện trước với 4 kênh đầu vào thay vì 3 kênh mặc định:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", channels=4)

Trong trường hợp này, model sẽ bao gồm các trọng số đã được huấn luyện trước ngoại trừ lớp đầu vào đầu tiên, vì nó không còn cùng hình dạng (shape) với lớp đầu vào đã huấn luyện trước. Lớp đầu vào này sẽ được khởi tạo lại bằng các trọng số ngẫu nhiên.

Số lượng lớp (Number of Classes)

Để tải một model YOLOv5s đã huấn luyện trước với 10 lớp đầu ra thay vì 80 lớp mặc định:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", classes=10)

Trong trường hợp này, model sẽ bao gồm các trọng số đã được huấn luyện trước ngoại trừ các lớp đầu ra, vì chúng không còn cùng hình dạng với các lớp đầu ra đã huấn luyện trước. Các lớp đầu ra này sẽ được khởi tạo lại bằng các trọng số ngẫu nhiên.

Buộc tải lại (Force Reload)

Nếu bạn gặp sự cố với các bước trên, việc đặt force_reload=True có thể hữu ích bằng cách xóa cache hiện có và buộc tải lại bản YOLOv5 mới nhất từ PyTorch Hub. Các bản copy đã lưu trong cache nằm tại ~/.cache/torch/hub; việc xóa thư mục đó cũng mang lại hiệu quả tương tự.

model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True)  # force reload

Inference trên màn hình

Để chạy inference trên màn hình desktop của bạn:

import torch
from PIL import ImageGrab

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

# Image
im = ImageGrab.grab()  # take a screenshot

# Inference
results = model(im)

Inference đa GPU (Multi-GPU)

Các model YOLOv5 có thể được tải vào nhiều GPU cùng lúc để thực hiện inference đa luồng:

import threading

import torch

def run(model, im):
    """Performs inference on an image using a given model and saves the output; model must support `.save()` method."""
    results = model(im)
    results.save()

# Models
model0 = torch.hub.load("ultralytics/yolov5", "yolov5s", device=0)
model1 = torch.hub.load("ultralytics/yolov5", "yolov5s", device=1)

# Inference
threading.Thread(target=run, args=[model0, "https://ultralytics.com/images/zidane.jpg"], daemon=True).start()
threading.Thread(target=run, args=[model1, "https://ultralytics.com/images/bus.jpg"], daemon=True).start()

Training

Để tải một model YOLOv5 cho mục đích huấn luyện thay vì inference, hãy đặt autoshape=False. Để tải một model với các trọng số khởi tạo ngẫu nhiên (để huấn luyện từ đầu), hãy sử dụng pretrained=False. Bạn phải tự cung cấp script huấn luyện trong trường hợp này. Ngoài ra, hãy xem Hướng dẫn huấn luyện dữ liệu tùy chỉnh YOLOv5 của chúng tôi để biết thêm về huấn luyện model.

import torch

model = torch.hub.load("ultralytics/yolov5", "yolov5s", autoshape=False)  # load pretrained
model = torch.hub.load("ultralytics/yolov5", "yolov5s", autoshape=False, pretrained=False)  # load scratch

Kết quả Base64

Dành cho việc sử dụng với các dịch vụ API. Xem ví dụ Flask REST API để biết chi tiết.

import base64
from io import BytesIO

from PIL import Image

results = model(im)  # inference

results.ims  # array of original images (as np array) passed to model for inference
results.render()  # updates results.ims with boxes and labels
for im in results.ims:
    buffered = BytesIO()
    im_base64 = Image.fromarray(im)
    im_base64.save(buffered, format="JPEG")
    print(base64.b64encode(buffered.getvalue()).decode("utf-8"))  # base64 encoded image with results

Kết quả cắt (Cropped)

Kết quả có thể được trả về và lưu dưới dạng các hình cắt (crops) phát hiện được:

results = model(im)  # inference
crops = results.crop(save=True)  # cropped detections dictionary

Kết quả Pandas

Kết quả có thể được trả về dưới dạng Pandas DataFrames:

results = model(im)  # inference
results.pandas().xyxy[0]  # Pandas DataFrame

Pandas Output (click to expand)

print(results.pandas().xyxy[0])
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 1  433.50  433.50   517.5  714.5    0.687988     27     tie
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

Kết quả đã sắp xếp

Kết quả có thể được sắp xếp theo cột, ví dụ: sắp xếp các chữ số biển số xe được phát hiện từ trái sang phải (trục x):

results = model(im)  # inference
results.pandas().xyxy[0].sort_values("xmin")  # sorted left-right

Kết quả JSON

Results can be returned in JSON format once converted to .pandas() dataframes using the .to_json() method. The JSON format can be modified using the orient argument. See pandas .to_json() documentation for details.

results = model(ims)  # inference
results.pandas().xyxy[0].to_json(orient="records")  # JSON img1 predictions

JSON Output (click to expand)

[
    {
        "xmin": 749.5,
        "ymin": 43.5,
        "xmax": 1148.0,
        "ymax": 704.5,
        "confidence": 0.8740234375,
        "class": 0,
        "name": "person"
    },
    {
        "xmin": 433.5,
        "ymin": 433.5,
        "xmax": 517.5,
        "ymax": 714.5,
        "confidence": 0.6879882812,
        "class": 27,
        "name": "tie"
    },
    {
        "xmin": 115.25,
        "ymin": 195.75,
        "xmax": 1096.0,
        "ymax": 708.0,
        "confidence": 0.6254882812,
        "class": 0,
        "name": "person"
    },
    {
        "xmin": 986.0,
        "ymin": 304.0,
        "xmax": 1028.0,
        "ymax": 420.0,
        "confidence": 0.2873535156,
        "class": 27,
        "name": "tie"
    }
]

Model tùy chỉnh

Ví dụ này tải một model YOLOv5s tùy chỉnh với 20 lớp đã huấn luyện trên VOC mang tên 'best.pt' với PyTorch Hub.

import torch

model = torch.hub.load("ultralytics/yolov5", "custom", path="path/to/best.pt")  # local model
model = torch.hub.load("path/to/yolov5", "custom", path="path/to/best.pt", source="local")  # local repo

Các model TensorRT, ONNX và OpenVINO

PyTorch Hub hỗ trợ inference trên hầu hết các định dạng xuất của YOLOv5, bao gồm cả các model đã được huấn luyện tùy chỉnh. Xem Hướng dẫn xuất model TFLite, ONNX, CoreML, TensorRT để biết chi tiết về việc xuất model.

Mẹo

TensorRT có thể nhanh hơn gấp 2-5 lần so với PyTorch trên GPU benchmarks.
ONNX và OpenVINO có thể nhanh hơn gấp 2-3 lần so với PyTorch trên CPU benchmarks.

import torch

model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.pt")  # PyTorch
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.torchscript")  # TorchScript
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.onnx")  # ONNX
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s_openvino_model/")  # OpenVINO
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.engine")  # TensorRT
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.mlmodel")  # CoreML (macOS-only)
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.tflite")  # TFLite
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s_paddle_model/")  # PaddlePaddle

Môi trường được hỗ trợ

Ultralytics cung cấp nhiều môi trường sẵn sàng sử dụng, mỗi môi trường đều được cài đặt sẵn các phụ thuộc thiết yếu như CUDA, CUDNN, Python, và PyTorch để giúp bạn bắt đầu dự án của mình.

Notebook GPU miễn phí:
Google Cloud: GCP Quickstart Guide
Amazon: AWS Quickstart Guide
Azure: AzureML Quickstart Guide
Docker: Docker Quickstart Guide

Trạng thái dự án

Huy hiệu này cho biết tất cả các bài kiểm thử Tích hợp liên tục (CI) của YOLOv5 GitHub Actions đều đã vượt qua thành công. Các bài kiểm thử CI này kiểm tra nghiêm ngặt chức năng và hiệu suất của YOLOv5 trên nhiều khía cạnh quan trọng: huấn luyện, xác thực, suy luận, xuất mô hình, và điểm chuẩn. Chúng đảm bảo hoạt động nhất quán và tin cậy trên macOS, Windows và Ubuntu, với các bài kiểm thử được tiến hành mỗi 24 giờ và sau mỗi lần commit mới.

Contributors

GLglenn-jocher¹³ PDpderrenger¹ RIRizwanMunawar¹ BUBurhan-Q¹

Created 12 thg 11, 2023Updated 6 ngày trước