콘텐츠로 건너뛰기

PyTorch 허브

📚 This guide explains how to load YOLOv5 🚀 from PyTorch Hub at https://pytorch.org/hub/ultralytics_yolov5.

시작하기 전에

요구 사항.txt를 Python>=3.8.0 환경을 포함하여 PyTorch>=1.8. 모델데이터 세트는 최신 YOLOv5 릴리스에서 자동으로 다운로드됩니다.

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

프로팁: 복제는 https://github.com/ultralytics / yolov5 불필요 😃

PyTorch 허브에 YOLOv5 로드

간단한 예

이 예는 PyTorch 허브에서 사전 학습된 YOLOv5s 모델을 다음과 같이 로드합니다. model 를 호출하고 추론을 위해 이미지를 전달합니다. 'yolov5s' 은 가장 가볍고 빠른 YOLOv5 모델입니다. 사용 가능한 모든 모델에 대한 자세한 내용은 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

자세한 예시

이 예는 다음을 보여줍니다. 일괄 추론 와 함께 PIL 그리고 OpenCV 이미지 소스. results인쇄 를 콘솔로 이동합니다, 저장runs/hub, 표시 를 사용하여 지원되는 환경에서 화면을 표시하고 텐서 또는 판다 데이터프레임.

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 zidane.jpg에 대한 추론 결과 YOLO bus.jpg에 대한 추론 결과

모든 추론 옵션은 다음을 참조하세요. YOLOv5 AutoShape() 앞으로 메서드.

추론 설정

YOLOv5 모델에는 신뢰도 임계값, IoU 임계값 등과 같은 다양한 추론 속성이 포함되어 있으며, 이를 설정할 수 있습니다:

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

장치

모델은 생성 후 모든 디바이스로 전송할 수 있습니다:

model.cpu()  # CPU
model.cuda()  # GPU
model.to(device)  # i.e. device=torch.device(0)

모델을 직접 만들 수도 있습니다. device:

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

프로팁: 입력 이미지는 추론 전에 올바른 모델 장치로 자동 전송됩니다.

무음 출력

다음을 사용하여 모델을 자동으로 로드할 수 있습니다. _verbose=False:

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

입력 채널

기본 3개가 아닌 4개의 입력 채널로 사전 학습된 YOLOv5s 모델을 로드합니다:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", channels=4)

이 경우 모델은 더 이상 사전 학습된 입력 레이어와 동일한 모양이 아닌 첫 번째 입력 레이어를 제외하고 사전 학습된 가중치로 구성됩니다. 입력 레이어는 무작위 가중치로 초기화된 상태로 유지됩니다.

수업 수

기본값 80이 아닌 10개의 출력 클래스로 사전 학습된 YOLOv5s 모델을 로드합니다:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", classes=10)

이 경우 모델은 더 이상 사전 학습된 출력 레이어와 동일한 모양이 아닌 출력 레이어를 제외한 사전 학습된 가중치로 구성됩니다. 출력 레이어는 무작위 가중치로 초기화된 상태로 유지됩니다.

강제 재장전

위의 단계에서 문제가 발생하면 설정 force_reload=True 기존 캐시를 삭제하고 PyTorch 허브에서 최신 YOLOv5 버전을 강제로 새로 다운로드하면 도움이 될 수 있습니다.

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

스크린샷 추론

데스크톱 화면에서 추론을 실행하려면 다음과 같이 하세요:

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)

다중GPU 추론

YOLOv5 모델을 스레드 추론을 통해 여러 GPU에 병렬로 로드할 수 있습니다:

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()

교육

추론이 아닌 트레이닝을 위해 YOLOv5 모델을 로드하려면 autoshape=False. 무작위로 초기화된 가중치로 모델을 로드하려면(처음부터 학습하려면) 다음을 사용합니다. pretrained=False. 이 경우 자체 교육 스크립트를 제공해야 합니다. 또는 YOLOv5 사용자 지정 데이터 자습서 교육 모델 트레이닝을 위해

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

Base64 결과

API 서비스와 함께 사용하세요. 자세한 내용은 https://github.com/ultralytics/yolov5 /pull/2291 및 플라스크 REST API 예제를 참조하세요.

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

잘린 결과

결과를 반환하고 감지 크롭으로 저장할 수 있습니다:

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

판다 결과

결과는 판다스 데이터프레임으로 반환할 수 있습니다:

results = model(im)  # inference
results.pandas().xyxy[0]  # Pandas DataFrame
판다 출력(확장하려면 클릭)
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

정렬된 결과

결과를 열별로 정렬할 수 있습니다(예: 차량 번호판 숫자 감지를 왼쪽에서 오른쪽(x축)으로 정렬):

results = model(im)  # inference
results.pandas().xyxy[0].sort_values("xmin")  # sorted left-right

박스 자르기 결과

결과를 반환하고 감지 크롭으로 저장할 수 있습니다:

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

JSON 결과

결과는 다음과 같이 변환된 후 JSON 형식으로 반환될 수 있습니다. .pandas() 데이터 프레임을 사용하는 .to_json() 메서드를 사용합니다. JSON 형식은 다음을 사용하여 수정할 수 있습니다. orient 인수를 사용합니다. 판다 보기 .to_json() 문서 를 참조하세요.

results = model(ims)  # inference
results.pandas().xyxy[0].to_json(orient="records")  # JSON img1 predictions
JSON 출력(확장하려면 클릭)
[
    {
        "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"
    }
]

사용자 지정 모델

이 예제에서는 사용자 지정 20클래스 VOC-훈련된 YOLOv5s 모델 'best.pt' PyTorch 허브.

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 및 OpenVINO 모델

PyTorch Hub는 사용자 지정 학습된 모델을 포함하여 대부분의 YOLOv5 내보내기 형식에 대한 추론을 지원합니다. 모델 내보내기에 대한 자세한 내용은 TFLite, ONNX, CoreML, TensorRT 내보내기 튜토리얼을 참조하세요.

💡 ProTip: TensorRT may be up to 2-5X faster than PyTorch on GPU benchmarks 💡 ProTip: ONNX and OpenVINO may be up to 2-3X faster than PyTorch on CPU benchmarks

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

지원 환경

Ultralytics 는 바로 사용할 수 있는 다양한 환경을 제공하며, 각 환경에는 다음과 같은 필수 종속성이 사전 설치되어 있습니다. CUDA, CUDNN, Python, 및 PyTorch와 같은 필수 종속 요소를 설치하여 프로젝트를 시작할 수 있습니다.

프로젝트 상태

YOLOv5 CI

이 배지는 모든 YOLOv5 GitHub Actions 지속적 통합(CI) 테스트가 성공적으로 통과되었음을 나타냅니다. 이러한 CI 테스트는 교육, 검증, 추론, 내보내기벤치마크 등 다양한 주요 측면에서 YOLOv5 의 기능과 성능을 엄격하게 확인합니다. 24시간마다 그리고 새로운 커밋이 있을 때마다 테스트를 수행하여 macOS, Windows 및 Ubuntu에서 일관되고 안정적인 작동을 보장합니다.

📅 Created 12 months ago ✏️ Updated 18 days ago

댓글