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ắt đầu#

Cài đặt requirements.txt trong môi trường Python>=3.8.0, bao gồm PyTorch>=1.8. Models và datasets sẽ tự động được tải xuống từ bản phát hành YOLOv5 mới nhất.

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

Mẹo

Không cần sao chép ultralytics/yolov5 — PyTorch Hub sẽ tự động tải code.

Tải YOLOv5 bằng PyTorch Hub#

Ví dụ đơn giản#

Ví dụ này tải một model YOLOv5s pretrained từ PyTorch Hub dưới dạng model và truyền một ảnh vào để inference. 'yolov5s' là model YOLOv5 nhẹ nhất và nhanh nhất. Để biết chi tiết về tất cả model hiện có, 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. results có thể được in ra console, lưu vào runs/hub, hiển thị trên màn hình trong các môi trường được hỗ trợ và trả về dưới dạng tensor hoặc dataframe pandas.

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(f"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

Để xem tất cả tùy chọn inference, hãy xem method forward AutoShape() của YOLOv5.

Thiết lập inference#

Các model YOLOv5 có nhiều thuộc tính inference, chẳng hạn như confidence threshold, IoU threshold, v.v., có thể được thiết lập bằng cách:

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ị#

Có thể chuyển model sang bất kỳ thiết bị nào sau khi tạo:

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

Model cũng có thể được tạo trực tiếp trên bất kỳ device nào:

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

Mẹo

Ảnh đầu vào sẽ được tự động chuyển sang thiết bị của model tương ứng trước khi inference.

Tắt output#

Có thể tải model mà không tạo output bằng _verbose=False:

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

Các kênh đầu vào#

Để tải model YOLOv5s pretrained 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 pretrained weights ngoại trừ layer đầu vào đầu tiên, vì layer này không còn có cùng shape với layer đầu vào pretrained. Layer đầu vào sẽ tiếp tục được khởi tạo bằng random weights.

Số lượng class#

Để tải model YOLOv5s pretrained với 10 class output thay vì 80 class 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 pretrained weights ngoại trừ các layer output, vì chúng không còn có cùng shape với các layer output pretrained. Các layer output sẽ tiếp tục được khởi tạo bằng random weights.

Buộc tải lại#

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

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

Inference từ ảnh chụp 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#

Có thể tải các model YOLOv5 lên nhiều GPU song song bằng threaded inference:

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 model YOLOv5 cho việc training thay vì inference, hãy thiết lập autoshape=False. Để tải model với weights được khởi tạo ngẫu nhiên (nhằm training từ đầu), hãy sử dụng pretrained=False. Trong trường hợp này, bạn phải cung cấp training script của riêng mình. Ngoài ra, hãy xem Tutorial Train Custom Data YOLOv5 của chúng tôi để biết cách training 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ù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ả đã crop#

Có thể trả về và lưu kết quả dưới dạng các crop detection:

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

Kết quả Pandas#

Có thể trả về kết quả 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#

Có thể sắp xếp kết quả theo cột, chẳng hạn để sắp xếp việc detection chữ số trên biển số xe 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#

Có thể trả về kết quả ở định dạng JSON sau khi chuyển đổi chúng thành dataframe .pandas() bằng method .to_json(). Có thể sửa đổi định dạng JSON bằng argument orient. Xem documentation về .to_json() của pandas để biết chi tiết.

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 'best.pt' gồm 20 class, được training trên VOC, bằng 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

Model TensorRT, ONNX và OpenVINO#

PyTorch Hub hỗ trợ inference trên hầu hết format export của YOLOv5, bao gồm cả các model được training tùy chỉnh. Xem tutorial TFLite, ONNX, CoreML, TensorRT Export để biết chi tiết về việc export model.

Mẹo

TensorRT có thể nhanh hơn PyTorch 2-5× trên GPU benchmarks.
ONNX và OpenVINO có thể nhanh hơn PyTorch 2-3× 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 được cài đặt sẵn các dependency thiết yếu như CUDA, CUDNN, Python và PyTorch, giúp bạn nhanh chóng bắt đầu các project.

Notebook GPU miễn phí:
Google Cloud: Hướng dẫn bắt đầu nhanh GCP
Amazon: Hướng dẫn bắt đầu nhanh AWS
Azure: Hướng dẫn bắt đầu nhanh AzureML
Docker: Hướng dẫn bắt đầu nhanh với Docker

Trạng thái dự án#

Huy hiệu này cho biết tất cả các bài kiểm thử GitHub Actions của YOLOv5 về Tích hợp liên tục (CI) đề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 năng của YOLOv5 trên nhiều khía cạnh quan trọng: huấn luyện, validation, suy luận, export và benchmark. Chúng đảm bảo hoạt động nhất quán và đáng tin cậy trên macOS, Windows và Ubuntu, với các bài kiểm thử được thực hiện mỗi 24 giờ và sau mỗi commit mới.

Cộng tác viên

RAraimbekovm¹ PDpderrenger¹ GLglenn-jocher¹

Được tạo 2 tháng trướcĐã cập nhật tuần trước