Bỏ qua nội dung

PyTorch Trung tâm

📚 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ạn 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à tập dữ liệu tự động tải xuống từ phiên bản mới nhất YOLOv5 giải phóng .

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

💡 Mẹo chuyên nghiệp: Không cần sao chép https://github.com/ ultralytics / yolov5 😃

Trọng tải YOLOv5 với PyTorch Trung tâm

Ví dụ đơn giản

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

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, đã cho thấy để sàng lọc trên các môi trường được hỗ trợ và trả về dưới dạng tenxơ 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 về 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() phía trước phương pháp.

Thiết lập suy luận

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

Các 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 sẽ tự động được chuyển đến thiết bị mẫu phù hợp trước khi suy luận.

Đầu ra im lặng

Các mô hình có thể được tải một cách im lặng 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 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 đào tạo trước ngoại trừ lớp đầu vào đầu tiên, không còn cùng hình dạng với 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ằng 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 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 đào tạo trước ngoại trừ các lớp đầu ra, không còn cùng hình dạng với 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ằng các trọng số ngẫu nhiên.

Buộc Tải Lại

Nếu bạn gặp vấn đề với các bước trên, hãy thiết lập force_reload=True có thể giúp bằng cách loại bỏ bộ nhớ đệm hiện có và buộc tải xuống mới nhất YOLOv5 phiên bản từ PyTorch Trục.

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

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

Để 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)

Đ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 cho đà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 tập lệnh đào tạo của riêng bạn trong trường hợp này. Hoặc xem YOLOv5 Hướng dẫn đào tạo dữ liệu tùy chỉnh để đào tạo người mẫu.

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ụ về Flask REST API để 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

Kết quả có thể được trả về và lưu dưới dạng dữ liệu 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ả được sắp xếp

Kết quả có thể được sắp xếp theo cột, nghĩa là sắp xếp theo thứ tự 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 cắt hộp

Kết quả có thể được trả về và lưu dưới dạng dữ liệu 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 tranh luận. Xem gấu trúc .to_json() tài liệu để biết thêm 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-mô hình YOLOv5s được đào tạo 'best.pt' với PyTorch Trục.

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 định dạng xuất, bao gồm các mô hình được đào tạo tùy chỉnh. Xem hướng dẫn TFLite, ONNX , CoreML , TensorRT Export để biết chi tiết về cách xuất mô hình.

💡 Mẹo chuyên nghiệp: TensorRT có thể nhanh hơn tới 2-5 lần so với PyTorch về điểm chuẩn GPU 💡 Mẹo chuyên nghiệp: ONNX và OpenVINO có thể nhanh hơn tới 2-3 lần so với PyTorch về đ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 nhiều môi trường sẵn sàng sử dụng, mỗi môi trường đều được cài đặt sẵn các thành phần phụ thuộc cần thiết như CUDA , CUDNN , Python và PyTorch để khởi động dự án của bạn.

Tình trạng dự án

YOLOv5 CI

Huy hiệu này cho biết 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: đào tạo , xác thực , suy luận , xuất và 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 bài kiểm tra được tiến hành sau mỗi 24 giờ và sau mỗi lần cam kết mới.

📅 Được tạo ra cách đây 1 năm ✏️ Đã cập nhật cách đây 1 tháng
Trợ lý Ultralytics glenn-jocher RizwanMunawar Burhan-Q

Bình luận