PyTorch Trung tâm
📚 This guide explains how to load YOLOv5 🚀 from PyTorch Hub at 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.
💡 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
Đố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:
Mô hình cũng có thể được tạo trực tiếp trên bất kỳ device
:
💡 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
:
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:
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:
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.
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)
Đa-GPU Suy luận
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 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()
Đà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:
Kết quả gấu trúc
Kết quả có thể được trả về dưới dạng Pandas DataFrames:
Đầu ra gấu trúc (nhấp để mở rộng)
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):
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:
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 GPU Điểm chuẩn 💡 ProTip: ONNX và OpenVINO có thể nhanh hơn tới 2-3 lần so với PyTorch trên CPU Điểm chuẩn
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 với 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.
- Tự do GPU Sổ ghi chép:
- Google Đám mây: 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.