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
Đố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.
- Sổ 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
Trạng thái dự án
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.