Chuyển đến nội dung

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. Các mô hình và bộ dữ liệu được tải xuống tự động từ phiên 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 chuyên nghiệp: Việc nhân bản https://github.com/ultralytics/yolov5 là không bắt buộc 😃

Tải YOLOv5 bằng PyTorch Hub

Ví dụ đơn giản

Ví dụ này tải một mô hình YOLOv5s đã được huấn luyện trước từ PyTorch Hub dưới dạng model và chuyển một ảnh để suy luận. 'yolov5s' là mô hình YOLOv5 nhẹ nhất và nhanh nhất. Để biết chi tiết về tất cả các mô hình có sẵn, 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 cho thấy suy luận theo lô với PIL và OpenCV các nguồn hình ảnh. results có thể là được in ra bảng điều khiển, đã lưu đến runs/hub, đã hiển thị ra màn hình trên các môi trường được hỗ trợ và được trả về dưới dạng các tensor 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

Kết quả suy luận YOLO trên zidane.jpg Kết quả suy luận YOLO trên bus.jpg

Đối với tất cả các tùy chọn suy luận, hãy xem YOLOv5 AutoShape() chuyển tiếp phương thức.

Cài đặt suy luận

Các mô hình YOLOv5 chứa nhiều thuộc tính suy luận khác nhau, chẳng hạn như ngưỡng tin cậy, ngưỡng IoU, v.v., có thể được đặt 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ị

Mô hình có thể được chuyển 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)

Mô hình cũng có thể được tạo trực tiếp trên bất kỳ device:

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

💡 Mẹo chuyên nghiệp: Hình ảnh đầu vào được tự động chuyển đến đúng thiết bị mô hình trước khi suy luận.

Tắt tiếng đầu ra

Có thể tải các mô hình một cách âm thầm bằng _verbose=False:

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

Kênh đầu vào

Để tải một mô hình YOLOv5s đã được 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, mô hình 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, lớp này không còn cùng hình dạng với lớp đầu vào đã được huấn luyện trước. Lớp đầu vào sẽ vẫn được khởi tạo bằng các trọng số ngẫu nhiên.

Số lượng Lớp

Để tải một mô hình YOLOv5s đã được 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, mô hình 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, các lớp này không còn cùng hình dạng với các lớp đầu ra đã được huấn luyện trước. Các lớp đầu ra sẽ vẫn được khởi tạo bằng các trọng số ngẫu nhiên.

Tải lại Bắt buộc

Nếu bạn gặp sự cố với các bước trên, việc cài đặt force_reload=True có thể giúp bằng cách loại bỏ bộ nhớ cache hiện có và buộc tải xuống phiên bản YOLOv5 mới nhất từ PyTorch Hub.

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

Suy luận Ảnh chụp màn hình

Để chạy suy luận trên màn hình máy tính 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)

Suy luận Đa GPU

Các mô hình YOLOv5 có thể được tải song song lên nhiều GPU với suy luận đ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()

Huấn luyện

Để tải mô hình YOLOv5 để huấn luyện thay vì suy luận, hãy đặt autoshape=False. Để tải một mô hình với các trọng số được 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 cung cấp tập lệnh huấn luyện của riêng bạn trong trường hợp này. Ngoài ra, hãy xem YOLOv5 của chúng tôi Hướng Dẫn Huấn Luyện Dữ Liệu Tùy Chỉnh cho việc huấn luyện mô hình.

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

Để sử dụng với các dịch vụ API. Xem ví dụ Flask REST API để biết thêm 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

Kết quả có thể được trả về và lưu dưới dạng crop phát hiện:

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
Đầu ra Pandas (nhấp để mở rộng) 
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, ví dụ: để sắp xếp phát hiện 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ả Cắt Hộp

Kết quả có thể được trả về và lưu dưới dạng crop phát hiện:

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

Kết quả JSON

Kết quả có thể được trả về ở định dạng JSON sau khi được chuyển đổi thành .pandas() khung dữ liệu bằng cách sử dụng .to_json() method. Định dạng JSON có thể được sửa đổi bằng cách sử dụng orient argument. Xem pandas .to_json() tài liệu để biết chi tiết.

results = model(ims)  # inference
results.pandas().xyxy[0].to_json(orient="records")  # JSON img1 predictions
Đầu ra JSON (nhấp để mở rộng) 
[
    {
        "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"
    }
]

Các Mô hình Tùy chỉnh

Ví dụ này tải một lớp 20 tùy chỉnh VOCmô hình YOLOv5s đã được huấn luyệ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 mô hình TensorRT, ONNX và OpenVINO

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

💡 ProTip: TensorRT có thể nhanh hơn PyTorch từ 2-5 lần trên các điểm chuẩn GPU 💡 ProTip: ONNX và OpenVINO có thể nhanh hơn PyTorch từ 2-3 lần trên các điểm chuẩn CPU

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 một loạt các 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 phần phụ thuộc thiết yếu như CUDA, CUDNN, Python và PyTorch, để khởi động các dự án của bạn.

Trạng thái dự án

YOLOv5 CI

Huy hiệu này cho biết rằng tất cả các thử nghiệm Tích hợp Liên tục (CI) YOLOv5 GitHub Actions đều vượt qua thành công. Các thử nghiệm 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 chính: huấn luyện, xác thực, suy luận, xuất và điểm chuẩn. 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 thử nghiệm được thực hiện sau mỗi 24 giờ và sau mỗi cam kết mới.



📅 Đã tạo 1 năm trước ✏️ Cập nhật 5 tháng trước
glenn-jocher UltralyticsAssistant RizwanMunawar Burhan-Q

Bình luận