PyTorch Trung tâm

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

Trước khi bắt đầu

Cài đặt requirements.txt trong một PythonMôi trường >=3.8.0 , bao gồm PyTorch>=1,8. Mô hình và bộ dữ liệu tải xuống tự động từ phiên bản mới nhất YOLOv5 phát hành.

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

💡 ProTip: Nhân bản https://github.com/ultralytics/yolov5là không bắt buộc 😃

Tải YOLOv5 với PyTorch Trung tâm

Ví dụ đơn giản

Ví dụ này tải mô hình YOLOv5s được đào tạo sẵn từ PyTorch Hub as model và chuyển một hình ảnh để suy luận. 'yolov5s' là nhẹ nhất và nhanh nhất YOLOv5 mẫu. Để biết chi tiết về tất cả các mô hình có sẵn, vui lòng xem ĐỌC SÁCH.

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 hàng loạt với PIL và Mở CV nguồn hình ảnh. results có thể là In để an ủi, Lưu đến runs/hub, hiển thị để sàng lọc trên các môi trường được hỗ trợ và trả về dưới dạng tensor hoặc Gấu trúc khung dữ liệu.

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

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

Đối với tất cả các tùy chọn suy luận, hãy xem YOLOv5 AutoShape() Về phía trước phương pháp.

Cài đặt suy luận

YOLOv5 các mô hình chứa các thuộc tính suy luận khác nhau như ngưỡng tin cậy, ngưỡng IoU, v.v. có thể được đặt bởi:

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

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

Thiết bị

Các 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

💡 ProTip: Hình ảnh đầu vào được tự động chuyển đến đúng thiết bị kiểu máy trước khi suy luận.

Kết quả đầu ra im lặng

Mô hình có thể được tải âm thầm với _verbose=False:

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

Kênh đầu vào

Để tải mô hình YOLOv5s được đào tạo sẵn 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 đào tạo trước ngoại trừ lớp đầu vào đầu tiên, không còn có hình dạng giống như lớp đầu vào được đào tạo trước. Lớp đầu vào sẽ vẫn được khởi tạo bởi các trọng số ngẫu nhiên.

Số lượng lớp học

Để tải mô hình YOLOv5s được đào tạo sẵn với 10 lớp đầu ra thay vì 80 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 đào tạo trước ngoại trừ các lớp đầu ra, không còn có hình dạng giống như các lớp đầu ra được đào tạo trước. Các lớp đầu ra sẽ vẫn được khởi tạo bởi các trọng số ngẫu nhiên.

Lực lượng tải lại

Nếu bạn gặp sự cố với các bước trên, hãy 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 mới nhất YOLOv5 Phiên bản từ PyTorch Trung tâm.

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 để 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

YOLOv5 Các mô hình có thể được tải vào nhiều GPU song song với suy luận luồng:

import torch
import threading


def run(model, im):
    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()

Đào tạo

Để tải một YOLOv5 Mô hình đào tạo thay vì suy luận, thiết lập autoshape=False. Để tải một mô hình với trọng số được khởi tạo ngẫu nhiên (để đào tạo từ đầu), hãy sử dụng pretrained=False. Bạn phải cung cấp kịch bản đào tạo của riêng bạn trong trường hợp này. Ngoài ra, hãy xem YOLOv5 Hướng dẫn đào tạo dữ liệu tùy chỉnh để đào tạo 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 https://github.com/ultralytics/yolov5/pull/2291 và ví dụ API Flask REST để biết chi tiết.

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 xén

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

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

Kết quả gấu trúc

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 gấu trúc (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ả được sắp xếp

Kết quả có thể được sắp xếp theo cột, tức là để sắp xếp phát hiện chữ số 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 xén hộp

Kết quả có thể được trả về và lưu dưới dạng cây trồng 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 chuyển đổi thành .pandas() Khung dữ liệu sử dụng .to_json() phương pháp. Định dạng JSON có thể được sửa đổi bằng cách sử dụng orient lý lẽ. Xem gấu trúc .to_json() tư 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"
  }
]

Mô hình tùy chỉnh

Ví dụ này tải một lớp 20 tùy chỉnh VOC-được đào tạo mô hình YOLOv5s 'best.pt' với PyTorch Trung tâm.

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

TensorRT, ONNX và OpenVINO Mô hình

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

💡 ProTip: TensorRT có thể nhanh hơn tới 2-5 lần so với PyTorch trên điểm chuẩn GPU 💡 ProTip: ONNX và OpenVINO có thể nhanh hơn tới 2-3 lần so với PyTorch trên đ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ụ thuộc thiết yếu như CUDA, CUDNN,Pythonvà PyTorch, để khởi động các dự án của bạn.

Máy tính xách tay 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 Docker

Tình trạng dự án

Huy hiệu này cho biết rằng tất cả YOLOv5 Các bài kiểm tra Tích hợp liên tục (CI) GitHub Actions đã vượt qua thành công. Các bài kiểm tra CI này kiểm tra nghiêm ngặt chức năng và hiệu suất của YOLOv5 trên các khía cạnh chính khác nhau: đào tạo, xác nhận, suy luận, xuất khẩu 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 tiến hành 24 giờ một lần và theo mỗi cam kết mới.

Đã tạo 2023-11-12, Cập nhật 2023-12-03
Tác giả: glenn-jocher (3)